Skip to content

Data Layout Optimization improvements

Raphaël Colin requested to merge rewrite-DLO into master

The Data Layout Optimization (DLO) introduced with the MLIR code generator is meant to improve performance by improving the data locality of ionic model state variables. This is done by transforming the layout of the state variables from an Array of Structs (AoS) to an Array of Structs of Array (AoSoA).

Previously, the implementation of this transformation was modifying the layout after the initialization of the state variables, by shifting data around which could then be interpreted by the MLIR-generated code.

This was however, poorly maintainable and wasn't usable by other code generators (like the original CPU code and the GPU MLIR code generator).

This merge request changes the way the state structure of ionic models is generated when DLO is enabled. This allows:

  • Using DLO in the baseline C/C++ ionic models
  • Using DLO in the GPU ionic models
    • This allows switching between CPU / GPU targets without having to change the layout
    • This can actually slightly improve performance (could be around 10% speedup) on GPU for large vector sizes (32 and above)
  • Clearer understanding of the layout (directly visible in the structure)
  • Potential compile-time errors instead of hard to detect runtime issues

Some fixes were also made in this merge request, like fixing the checkpoint feature which wasn't working correctly with DLO enabled

Merge request reports