- Shaped by: Daniel, Enrique
- Appetite (FTEs, weeks): full cycle
- Developers: Nina<!-- Filled in at the betting table unless someone is specifically required here -->
## Problem
<!-- The raw idea, a use case, or something we’ve seen that motivates us to work on this -->
The MeteoSwiss operational setup gains roughly 10% speedup by switching from double precision to mixed precision with the OpenACC dycore. Mixed-precision is a compile option in ICON, where parts of the variables and computations are done in single precision and some are done in double precission.
To compete with this performance improvement, it would be good to be able to use mixed-precision in the gt4py dycore as well, which is currently in all double-precision.
## Appetite
<!-- Explain how much time we want to spend and how that constrains the solution -->
This feature has the potential to put the GT4PY dycore much closer to be put into production, therefore the appetite is quite high and could be estimated by 3-4 weeks.
## Solution
<!-- The core elements we came up with, presented in a form that’s easy for people to immediately understand -->
The solution for competing with the 10% performance increase, is simply to transfer the feature over to the gt4py dycore.
For that multiple steps are needed:
* First, one has to add an additional feature in GT4PY, see below.
* Second, in ICON4py and ICON-EXCLAIM, one needs to add type aliases for single and double precision, which can be set by using `--enable-mixed-precision`.
* Then the different fields need to be typed accordingly as in ICON (`wp` for double-precision, `vp` in single-precision). This needs to be done in all the stencils.
* Verify mixed precision in icon-exclaim
* Measure performance vs full double precision
The rest is expected to be already working.
### New feature in gt4py
Using type aliases (python variables containing an actual numerical type) as dtypes in Field annotations is transparently supported by gt4py. However, for explicit construction of literal values and for casting field values (`astype(foo, float32)`), only the actual type constructor functions are supported (e.g. `float32()`, `float64()`, ...). We need to add support for using type aliases in these two cases.
The designed solution would be to add a new pass to replace type aliases uses by the actual type they represent at the FOAST level. In this way, the rest of the toolchain will work without further changes.
```python!
if enable_mixed_precision:
vpfloat = float32
wpfloat = float64
else:
vpfloat = float64
wpfloat = float64
@field_operator
def stencil(a: Field[[V, K], vfloat], b: Field[[V, K], wpfloat]) -> Field[[V, K], vpfloat]:
return vpfloat("3.14159") * (a + as_type(b, vpfloat))
# After the new preprocess pass:
# return float32("3.14159") * (a + as_type(b, float32))
```
The pass should be added as first step of the [postprocessing block](https://github.com/GridTools/gt4py/blob/fef9655a9994d9927fcad6af9a065c322a228fe3/src/gt4py/next/ffront/func_to_foast.py#L94) of the FieldOperator parser, and it should receive both the FOAST node and the closure vars (since they contain the type aliases).
## Rabbit holes
<!-- Details about the solution worth calling out to avoid problems -->
## No-gos
- A different strategy to implement this feature would be to extend the type system with a new `Type` element which represents variables containing types. Although it might make more sense in the long term, it involves changes in several components, so for now we stick with the simple solution (maybe add a TODO in the new pass to mention this possibility.)
## Caveats
* For using probtest the peturb amplitude needs to be adjusted.