Developing a Domain-Specific Language in Python to Achieve Productivity, Portability and Performance for a Weather and Climate Model
Domain-specific languages (DSLs) help separate high performance computing aspects from model code used in development as well as production. The increase of diversity and complexity of modern hardware architectures significantly challenges how large, complex models are written to accommodate for different targets since the complexity of the code increased significantly in order to accommodate for execution models on these targets. Traditional approaches such as Fortran+X or C++ struggle with this due to limitations with code-generation and lack of introspection.
The GT4Py compiler framework provides a simple, sequential frontend language in Python for expressing common numerical patterns in weather and climate applications while leveraging code generation and domain specific optimization to achieve high performance. It can rely on established performance portability frameworks (e.g. GridTools C++) in the backend or directly generate low-level code (e.g. C++/CUDA). GT4Py leverages the Python parser to generate a Python AST which is then lowered through internal representations processed by analysis, parallelization and optimization compiler passes before being passed to the hardware specific backend for code generation. The generated code is compiled and wrapped into a Python module which can then be called by the user.
This presentation will give an overview of the GT4Py domain-specific language and compiler. It will also summarize experience from applying it to the dynamical core and physical parameterizations of the FV3GFS model. We will conclude by giving a status update on performance and discuss our roadmap and opportunities for collaboration.