Horizontal Regions in GTC backends

# Horizontal Regions in GTC backends ###### tags: `cycle 3` shaped by: Johann ## Problem Horizontal Regions are used in FV3core, and so if FV3core is to use the GTC backends, then support for this needs to be in the GTC IRs and backends. ## Appetite Part of the FV3core on GTC package of projects. ## Solution Initial work has already focused on the support and optimizations in gtir and oir in the [regions-gtc](https://github.com/jdahm/gt4py/tree/regions-gtc) branch. Therefore, the scope of this project is: 0. Merge frontend changes. 1. Review existing gtir and oir implementations. 2. Add to gtcpp. 3. Add to cuir. 4. Add a test suite test with validation. 5. Coordinate with the numpy backend development and add to npir. 6. Coordinate with the dace backend development and add to daceir. Details involved in each step are described below: ### 0. Merge frontend GridTools/gt4py [#424](https://github.com/GridTools/gt4py/pull/424) has the frontend changes. #### Prerequisites for project - [x] Merge current master into branch - [x] Make PR to GridTools/gt4py:master - [x] Ensure frontend tests are on the branch ### 1. Review gtir and oir changes #### Explanation of IR concepts The `HorizontalRegion` nodes in gtir, become `MaskStmt` with a `HorizontalMask` boolean expression in oir. If a number of these mask statements have non-overlapping masks, then that set is converted by an optimization pass to a single statement, where the `value` part of the assignment is a `HorizontalSwitch`. #### Prerequisites for project - [x] Review diff with the master branch - [ ] Finish optimization pass converting statements to single switch ### 2. Add to gtcpp GridTools does not support different extents depending on the position in the IJ plane. This makes a "hack" or workaround in the implementation. Fortunately, we check on the the gtir that no temporaries are assigned to in regions. This limits their usefulness, but it makes the GridTools implementation possible in a predictable way. **Implementation Note:** Technically we could make this less strict and move this check later to the gtcpp backend and check that no non-stage-local temporaries get assigned to, since local scalars in stages are not assigned extents in GridTools. The two subsections below describe the difficult aspects of completing the work. It should be noted that these were similar difficulties in the legacy backed implementation, and remains in a somewhat hacked state there. #### Additional stage inputs and lowering A number of things need to be added to the gtcpp and code-generation. These might even be worth splitting into separate PRs: 1. Positional computations: `i_pos` and/or `j_pos` need to be added as stage inputs on any stage that has a `HorizontalMask`. See [here](https://github.com/GridTools/gridtools/blob/master/tests/regression/positional_stencil.cpp) for an example in GTv2. 2. Domain size: `domain_size_i`, `domain_size_j` need to be added as stage inputs on any stage that requires it based on the `HorizontalMask`. `domain` in the gtcpp generated code could be a lambda capture and then a global parameter made. 3. The `MaskStmt` with `HorizontalMask` need to be lowered to an interpretation that can be easily code generated using the new stage accessors mentioned above. #### Complications of extent analysis Currently oir does not perform any extent analysis, so the `MaskStmt` nodes with `HorizontalMask` masks might actually not be needed to execute the code (depending on the location of the mask in the IJ plane). An additional critical and complicated step in gtcpp is that these nodes are eliminated during extent analysis, since they are needed to determine the extents of other calculations, and at the same time the extent is necessary to determine how far or if at all that statement is executed. ### 3. Add to cuir Similar extent issue as in gtcpp ir. Will need to figure out how to lower the `HorizontalMask` and the `HorizontalSwitch` to something easily code-generated by the cuir backend. ### 4. Test suite addition The test suite should have all the capability we need to test horizontal regions end-to-end. This should be relatively straightforward. ### 5. Numpy backend addition Similar extent issue as in gtcpp ir. Need to coordinate with the numpy backend effort to add this. ### 6. Dace backend addition Coordination needed with other shaped project. ## Potential Rabbit Holes Extent analysis `MaskStmt` elimination needs to happen for each backend, but currently the analysis happens for each backend independently. This should be considered to be moved into oir. However, doing so might end up being a time sink. ## No Gos Non-stage-local temporaries in conditionals.