ai2cm / pace

Implementation of the FV3GFS / SHiELD atmospheric model in Python
https://ai2cm.github.io/pace/
Apache License 2.0
37 stars 12 forks source link

DaCe Orchestration: fix parsing, de-inline grid data scalars & memory leak #319

Closed FlorianDeconinck closed 2 years ago

FlorianDeconinck commented 2 years ago

Purpose

This PR address two issues:

Code changes:

'alas, our distributed compilation system works by compiling a 3,3 layout

top tile, then using those 9 caches on every layout upward.

This setup leads to the values of da_min/da_min_c from the 3,3 layout

to be inlined in the generated code. Those variables are used in runtime

calculation (kinetic energy, etc.) which obviously leads to misbehaving numerics

and errors when the 3,3 layout values are used on larger layouts

The solution we implement here is making use of the fact that callbacks

are never inlined in dace optimization. the current workaround uses the

following functions.

An alternative would be to pass the Grid or the DampingCoefficients to DaCe,

clearly flagging it has a dynamic piece of memory (which would

cancel any inlining) but the feature to do that (dace.struct)

is currently in disarray.

N.B.: another solution is to pass da_min and da_min_c as input, put it seems

odd and adds a lot of boilerplate throughout the model code.

- `dace_constants_args` has been renamed to `dace_compiletime_args` to match DaCe API naming change
- change in field referenced in dyn_core, see comment left in code above `self.ptc` and `self.delpc`:

ToDo: Due to DaCe VRAM pooling creating a memory

leak with the usage pattern of those two fields (to be fixed soon)

We use the C_SW internal to workaround it e.g.:

- self.cgrid_shallow_water_lagrangian_dynamics.delpc

- self.cgrid_shallow_water_lagrangian_dynamics.ptc



## Requirements changes:

- N/A

## Infrastructure changes:

- N/A

## Checklist
Before submitting this PR, please make sure:

- [ ] You have followed the coding standards guidelines established at [Code Review Checklist](https://drive.google.com/file/d/1R0nqOxfYnzaSdoYdt8yjx5J482ETI2Ft/view?usp=sharing).
- [ ] Docstrings and type hints are added to new and updated routines, as appropriate
- [ ] All relevant documentation has been updated or added (e.g. README, CONTRIBUTING docs)