Separating physics and performance: Using Python to implement a fast and maintainable FV3 dynamical core

Abstract:

The increasing power and efficiency of emerging supercomputers is a boon to climate science; long-term global cloud-resolving simulations are coming into reach and are a promising avenue to reduce the uncertainty of global climate projections. However, most global climate models have been optimized for performance on a specific hardware architecture, and typically cannot run efficiently on GPUs without extensive effort that increases the complexity of the code. Decoupling the expression of the numerical algorithm from hardware specific optimization improves the maintainability of a climate model and makes it easier to achieve high performance on a variety of hardware architectures. Such a separation of concerns is possible using a domain-specific language (DSL) embedded in a high-level programming language such as Python and a compiler toolchain ensuring the generation of high-performance code. The DSL approach also reduces the need for optimization in the user-facing code and allows for the use of abstractions such as function calls without a performance penalty, increasing code reusability and concision. To demonstrate the potential of this strategy we will present our work on a Python implementation of the Finite-Volume Cubed-Sphere (FV3) dynamical core, backed by the GT4Py DSL, which provides access to the Dawn compiler toolchain. We will show our ability to achieve good performance on multiple architectures and discuss our approach to porting the Fortran code to this framework, including the design decisions, testing strategy and infrastructure that enabled the effort, as well as lessons learned.