feat[cartesian]: 2D temporaries [Experimental]

[FlorianDeconinck]

Description

Canonically, temporaries have been 3D arrays of float. We introduced a few months back the capacity to type hint the type to allow for integer and mixed precision temporaries.

Another missing feature is the capacity to ask for a 2D temporary when the semantic of the code asks for it (aggregate, mask, etc.). We introduce with two flavors:

• The default behavior knows how to resolve the FieldDescriptor type e.g.

tmp: Field[IJ, np.float64] = ...

• A user defined mapping from symbols (as string) to type allows for applications on top of GT4Py to control the verbosity of the syntax

@stencil(backend=..., dtype={'MyType': Field[IJ, np.float64]})
def stencil(...):
   ...
   tmp: MyType = ...

The feature is limited to 2D temporaries for now - and guarded for it.

The feature is experimental because support for gt:X backends is missing.

Requirements

• All fixes and/or new features come with corresponding tests.
• Important design decisions have been documented in the appropriate ADR inside the docs/development/ADRs/ folder.

[FlorianDeconinck]

@egparedes / @havogt

How do we declare 2D temporaries in GridTools? The API of GT_DECLARE_TMP seems to be tied to a full grid domain. What am I missing?

(Feature is experimental because I couldn't figure that one out)

[FlorianDeconinck]

cscs-ci run

[romanc]

The code in the PR is looking great. Some nitpicks inline.

I'd like to see an ADR for this, especially since it's an experimental feature. The list of experimental features here is currently the only place where we track all of them in one place. This list will come in handy at the point when we regularly re-evaluate the state of experimental features (as we promise to do in the ADR of experimental features (to address the concern of stale experimental features)).

[egparedes]

@egparedes / @havogt

How do we declare 2D temporaries in GridTools? The API of GT_DECLARE_TMP seems to be tied to a full grid domain. What am I missing?

Hannes is offline until next week and I'm afraid I cannot help here. @fthaler should also know about this but I forgot to ping him, I'm sorry 🤦‍♂️ ,

[fthaler]

@egparedes / @havogt
How do we declare 2D temporaries in GridTools? The API of GT_DECLARE_TMP seems to be tied to a full grid domain. What am I missing?

Hannes is offline until next week and I'm afraid I cannot help here. @fthaler should also know about this but I forgot to ping him, I'm sorry 🤦‍♂️ ,

There is no way to declare 2D temporaries, because there is no way to fill them with data. All computations within a single spec runs on the same 3D domain, so you can’t fill 2D temporaries in there without possibly getting race conditions. Or what’s the use case?

[FlorianDeconinck]

There is no way to declare 2D temporaries, because there is no way to fill them with data. All computations within a single spec runs on the same 3D domain, so you can’t fill 2D temporaries in there without possibly getting race conditions. Or what’s the use case?

Use cases so far are all in physics parametrization where there's computation in or around particular levels in the atmosphere (so surface, or inversion level, etc.) which demands accretion to be temporarily stored within a stencil computation. We have been going around by defining those temporaries as input arguments - but for some parametrization we end up with dozens of fields and the scientists are not particularly happy about it.

We could go around the GridTools limitation directly in GT4Py with a repository of 2D buffers we pass down - it's a bit of infrastructure but we can make it happen that way.

[FlorianDeconinck]

I have read your comment a bit fast. In gt4py.cartesian we have the guardrails in place to make sure 2D temporaries are not written with race conditions.

[havogt]

Can you give a concrete example? I think we concluded at some point that 2D temporaries are not useful, but don't remember why...

[FlorianDeconinck]

Here are some:

In our shallow convection parametrization:

    with computation(FORWARD), interval(...):
        if MELT_GLAC == 1 and cumulus == 1:
            if K <= ktf-1:
                if ierr == 0:
                    dp = 100.*(po_cup-po_cup[0,0,1])\
                    pwo_eff = 0.5*(pwo+pwo[0,0,1])
                    pwo_solid_phase = (1.-p_liq_ice)*pwo_eff
                    total_pwo_solid_phase = total_pwo_solid_phase+pwo_solid_phase*dp/constants.MAPL_GRAV

That solid phase temporary, will feed later on the stencil on each side of the cloud.

In our microphysics code we have an interpolation code that looks at particular levels for calculation, then expand on the full column. Those use 2D fields that are not useful outside. Because of how we write column code in cartesian they are necessarily 2D temporaries. In the example here pb, pt and the boolean_2d_mask could entirely be inlined

def vertical_interpolation(
    field: FloatField,
    interpolated_field: FloatFieldIJ,
    p_interface_mb: FloatField,
    target_pressure: Float,
    pb: FloatFieldIJ,
    pt: FloatFieldIJ,
    boolean_2d_mask: BoolFieldIJ,
):
    """
    Interpolate to a specific vertical level.

    Only works for non-interface fields. Must be constructed using Z_DIM.

    Arguments:
        field (in): three dimensional field to be interpolated to a specific pressure
        interpolated_field (out): output two dimension field of interpolated values
        p_interface_mb (in): interface pressure in mb
        target_pressure (in): target pressure for interpolation in Pascals
        pb (in): placeholder 2d quantity, can be removed onces 2d temporaries are available
        pt (in): placeholder 2d quantity, can be removed onces 2d temporaries are available
        boolean_2d_mask (in): boolean mask to track when each cell is modified
    """
    # mask tracks which points have been touched. check later on ensures that every point has been touched
    with computation(FORWARD), interval(0, 1):
        boolean_2d_mask = False

    with computation(FORWARD), interval(-1, None):
        pb = 0.5 * (log(p_interface_mb * 100) + log(p_interface_mb[0, 0, 1] * 100))

    with computation(BACKWARD), interval(1, None):
        pt = 0.5 * (log(p_interface_mb[0, 0, -1] * 100) + log(p_interface_mb * 100))
        if log(target_pressure) > pt and log(target_pressure) <= pb and not boolean_2d_mask:
            al = (pb - log(target_pressure)) / (pb - pt)
            interpolated_field = field[0, 0, -1] * al + field * (1.0 - al)
            boolean_2d_mask = True
        pb = pt

    with computation(FORWARD), interval(-1, None):
        pt = 0.5 * (log(p_interface_mb * 100) + log(p_interface_mb[0, 0, -1] * 100))
        pb = 0.5 * (log(p_interface_mb * 100) + log(p_interface_mb[0, 0, 1] * 100))
        if (
            log(target_pressure) > pb
            and log(target_pressure) <= log(p_interface_mb[0, 0, 1] * 100)
            and not boolean_2d_mask
        ):
            interpolated_field = field
            boolean_2d_mask = True

        # ensure every point was actually touched
        if boolean_2d_mask == False:  # noqa
            interpolated_field = field

    # reset masks and temporaries for later use
    with computation(FORWARD), interval(0, 1):
        boolean_2d_mask = False
        pb = 0
        pt = 0

[FlorianDeconinck]

I might add that it's a fine answer to us that GridTools, the C++ lib, will not provide those. We can create the necessary infrastructure in GT4Py, which by it's frontend knows that those Field would be accessed badly and therefore can commit to their dimensions.

[fthaler]

Can you give a concrete example? I think we concluded at some point that 2D temporaries are not useful, but don't remember why...

AFAIR, we just concluded that for all use cases they can just easily be replaced by ‘normal’ 2D fields, i.e., we can’t do useful optimizations with them (in contrast to the 3D ones which can be cached) and they don’t require extent analysis (as we can’t fuse them with following 3D computations due to the synchronization issues). So we decided the advantage of having them is not worth the implementation effort as we won’t gain any performance.

I might add that it's a fine answer to us that GridTools, the C++ lib, will not provide those. We can create the necessary infrastructure in GT4Py, which by it's frontend knows that those Field would be accessed badly and therefore can commit to their dimensions.

From the GridTools point of view that’s probably the easiest solution.

[FlorianDeconinck]

This makes sense to me. I'll give time for Hannes to weigh in, but I am good with this separation of responsibility between gt4py and gridtools.

[romanc]

Looks good for me. One important question about dependencies (in cartesian/frontend/nodes.py) and a bunch of not so important typos in the ADR (thanks for adding one). Otherwise this looks ready from my point of view.

[twicki]

why do we remove the "default" argument (True, True, True) here?

[FlorianDeconinck]

Because we need the Field to respect the dimensions and I have run into a number of "defaults" that are difficult to track when we move to other dimensions. I'd rather have dimensions passed through proper.

[FlorianDeconinck]

cscs-ci run

[FlorianDeconinck]

cscs-ci run

[romanc]

🚀

[twicki]

Looks good

[FlorianDeconinck]

cscs-ci run

[romanc]

cscs-ci run

[romanc]

cscs-ci run

[FlorianDeconinck]

cscs-ci run