# While Loops and Variable Vertical Offsets in GTC

###### tags: `cycle 3`
shaped by: EddieD

## Problem

The Lagrangian to Eulerian vertical remapping motif in the FV3 dynamical core has been implemented using a single `gt4py` stencil in `fv3core` using a `while` loop frontend feature. This has already been implemented in the legacy toolchain using variable vertical offsets. These are offsets in the *k*-direction that can be integer scalars or fields rather than only integer literals. In this case, the value of the varaible offset is added to the current *k*-level. These features need to be implemented in the GTC toolchain as well.

The typical use case in a stencil is like the following,

def while_loop_stencil(
    pe1: FloatField,
    pe2: FloatField,
    lev: IntFieldIJ,
    qout: FloatField,
):
    with computation(FORWARD), interval(...):                
        while pe1[0, 0, lev + 1] < pe2[0, 0, lev + 1]:
            qout += pe1[0, 0, lev] * pe2[0, 0, lev]
            lev += 1

In this example, the fields `pe1` and `pe2` correspond to two pressure varaibles that are being compared, `lev` is a 2D integer field for keeping track of the current `k`-level, and `qout` is an output field containing the summation.

## Appetite

* Probably half-cycle, as the work is in already in progress.

## Solution

The implementation is separated into two parts:

1. The variable vertical offsets
2. The `while` loop construct

The legacy toolchain versions of each feature are implemented in `gt4py` pull requests [396](https://github.com/GridTools/gt4py/pull/396) and [410](https://github.com/GridTools/gt4py/pull/410), respectively.

### 1. Variable Vertical Offsets in GTC

The proposed solution is to add a new node type called `VariableOffset` that extends the existing `CartesianOffset` type by allowing the `k` offset to be any expression. A validator verifies that the expression has an integer data type. For the purposes of extent analysis, the `VariableOffset` returns an arbitrary large integer value for `k`. The generic `VaraibleOffset` in `common` is mirrored in the `GTIR`, `OIR`, `GTCpp`, and `CUIR` intermediate representations. Code generation is not modified since the `CartesianOffset` templates can be reused for `VariableOffset` nodes.

### 2. While Loops in GTC

The `while` statement in GTC is implemented similarly to an `if` statement with a conditional expression and a body of statements. The `While` node in the `GTIR` is lowered to a `MaskStmt` in OIR much like the `ScalarIfStmt`. An `is_loop` attribute is added to the `MaskStmt` to indicate whether it should be used to generate an `if` statement or a `while`, but this implementation is certainly open to discussion. In the `GtCpp` backend, the `MaskStmt` with loop is lowered back to a `While` statement, much like the `IfStmt`, but remains a `MaskStmt` in CUIR. The code generation macros are updated accordingly.

## Potential Rabbit Holes

1. There has been some active discussion about whether such a feature belongs as part of a DSL as it tends toward a general purpose language. Another issue is the difficulty of applying optimizations to such a generic feature. The proposed `while` and `for` loop abstractions could potentially be lifted to a more general conditional concept in the OIR.

2. Keeping track of the current `k`-level requires passing an IJ integer field into the stencil to be used as a counter. A full 3D field would not work so a temporary cannot be used. Either typed temporaries, e.g., `tmp: Field[IJK, int] = 0` would be required to replace this, or a compiler pass that can infer that the field used in the variable vertical offset should be 2D.

3. As in a general purpose language, there are currently no safety checks on the `while` loop conditions meaning that infinite loops are possible. The backend could instead generate code that guarantees the variable offset when added to the current *k*-level does not exceed the domain size in the *k*-direction.

## No Gos

The aforementioned discussion about whether such general purpose loop constructs belong in a DSL for weather and climate applications results in the conclusion, that no, indeed they do not.