Exploring a high-level programming model for the NWP domain using ECMWF microphysics schemes

We explore the domain-specific Python library GT4Py (GridTools for Python) for implementing a representative physical parametrization scheme and the related tangent-linear and adjoint algorithms from the Integrated Forecasting System (IFS) of ECMWF. GT4Py encodes stencil operators in an abstract and hardware-agnostic fashion, thus enabling more concise, readable, and maintainable scientific applications. The library achieves high performance by translating the application into targeted low-level coding implementations. Here, the main goal is to study the correctness and performance portability of the Python rewrites with GT4Py against the reference Fortran code and a number of automatically and manually ported variants created by ECMWF. The present work is part of a larger cross-institutional effort to port weather and climate models to Python with GT4Py. The focus of the current work is the IFS prognostic cloud microphysics scheme, a core physical parametrization represented by a comprehensive code that takes a significant share of the total forecast model execution time. In order to verify GT4Py for numerical weather prediction (NWP) systems, we put additional emphasis on the implementation and validation of the tangent-linear and adjoint model versions which are employed in data assimilation. We benchmark all prototype codes on three European supercomputers characterized by diverse graphics processing unit (GPU) and central processing unit (CPU) hardware, node designs, software stacks, and compiler suites. Once the application is ported to Python with GT4Py, we find excellent portability, competitive GPU performance, and robust execution in all tested scenarios including with single precision.

1 Introduction

Soon after its first public release in 1957, Fortran became the language of choice for weather and climate models (Méndez et al., 2014). On the one hand, its procedural programming style and its built-in support for multi-dimensional arrays have granted Fortran large popularity in the whole scientific computing community. On the other, its low-level nature guarantees fast execution of intensive mathematical operations on vector machines and conventional central processing units (CPUs). In the last decades, these characteristics have permitted weather forecasts to be run several times per day under tight operational schedules on high-performance computing (HPC) systems (Neumann et al., 2019).

In recent years, in response to the simultaneous end of Moore's law and Dennard scaling, and due to the societal challenge to reduce energy consumption, the computer hardware landscape has been undergoing a rapid specialization to prevent unsustainable growth of the power envelope (Müller et al., 2019). As a result, most supercomputers nowadays have a heterogeneous node design, where energy-efficient accelerators such as graphics processing units (GPUs) co-exist with traditional CPUs. Because Fortran has been conceived with CPU-centric machines in mind, efficient programming of hybrid HPC platforms using the core Fortran language can be challenging (Méndez et al., 2014; Lawrence et al., 2018). Indeed, the sustained performance of legacy weather and climate model codes written in Fortran has decreased over the decades (Schulthess et al., 2018), revealing the urgency for algorithmic and software adaptations to remain competitive in the medium and long term (Bauer et al., 2021).

Compiler directives (or pragmas) are an attractive solution for parallelization, both to spread a workload across multiple CPU threads and to offload data and computations to GPU. The most famous incarnations of this programming paradigm are OpenMP (Dagum and Menon, 1998) and OpenACC (Chandrasekaran and Juckeland, 2017). Because compiler directives accommodate incremental porting and enable non-disruptive software development workflows, they are adopted by many weather and climate modeling groups, who are facing the grand challenge of accelerating large code bases with thousands of source files and millions of lines of code, which stem from decades of scientific discoveries and software developments (Lapillonne et al., 2017, 2020; Randall et al., 2022). In order not to threaten the overall readability of the code by exposing low-level instructions, the annotation of Fortran codes with compiler directives can be automated in the pre-processor step of the compilation process using tools such as the CLAW compiler (Clement et al., 2019) or the ECMWF source-to-source translation tool Loki (https://github.com/ecmwf-ifs/loki, last access: 23 September 2024). Although pragma-based programming models can support intrusive hardware-specific code transformations, additional specialized optimizations may still be required, which could finally lead to code duplication and worsen maintainability (Dahm et al., 2023). Moreover, performance and portability are much dependent on the level of support and optimization offered by the compiler stack.

On the contrary, domain-specific languages (DSLs) separate the code describing the science from the code actually executing on the target hardware, thus enabling performance portability, namely application codes that achieve near-optimal performance on a variety of computer architectures (Deakin et al., 2019). Large portions of many modeling systems are being rewritten using multiple and diverse DSLs, not necessarily embedded in Fortran. For instance, the dynamical core of the weather prediction model from the COnsortium for Small-scale MOdeling (COSMO; Baldauf et al., 2011) has been rewritten in C++ using the GridTools library (Afanasyev et al., 2021) to port stencil-based operators to GPUs (Fuhrer et al., 2014, 2018). Similarly, HOMMEXX-NH (Bertagna et al., 2020) is an architecture-portable C++ implementation of the non-hydrostatic dynamical core of the Energy Exascale Earth System model (E3SM; Taylor et al., 2020) harnessing the Kokkos library to express on-node parallelism (Edwards et al., 2014). The GungHo project for a new dynamical core at the UK Met Office (Melvin et al., 2019, 2024) blends the LFRic infrastructure with the PSyclone code generator (Adams et al., 2019). Pace (Ben-Nun et al., 2022; Dahm et al., 2023) is a Python rewrite of the Finite-Volume Cubed-Sphere Dynamical Core (FV3; Harris and Lin, 2013) using GT4Py (GridTools for Python) to accomplish performance portability and productivity. Similarly, various Swiss partners including the Swiss Federal Office of Meteorology and Climatology (MeteoSwiss), ETH Zurich, and the Swiss National Supercomputing Centre (CSCS) are porting the ICOsahedral Non-hydrostatic modeling framework (ICON; Zängl et al., 2015) to GT4Py (Luz et al., 2024). In another related project (Kühnlein et al., 2023), a next-generation model for the IFS at ECMWF is developed in Python with GT4Py building on the Finite Volume Module (FVM; Smolarkiewicz et al., 2016; Kühnlein et al., 2019).

The focus of the portability efforts mentioned above is the model dynamical core – the part of the model numerically solving the fundamental non-linear fluid-dynamics equations. In the present work, we turn the attention to physical parametrizations – which account for the representation of subgrid-scale processes – and additionally address the associated tangent-linear and adjoint algorithms. Parametrizations are being commonly ported to accelerators using OpenACC (e.g., Fuhrer et al., 2014; Yang et al., 2019; Kim et al., 2021). Wrappers around low-level legacy physics codes might then be designed to facilitate adoption within higher-level workflows (Monteiro et al., 2018; McGibbon et al., 2021). Lately, first attempts at refactoring physical parametrizations with respect to portability have been documented in the literature. For instance, Watkins et al. (2023) presented a rewrite of the MPAS-Albany Land Ice (MALI) ice-sheet model using Kokkos. Here, we present a Python implementation of the cloud microphysics schemes CLOUDSC (Ubbiali et al., 2024c) and CLOUDSC2 (Ubbiali et al., 2024d), which are part of the physics suite of the IFS at ECMWF.¹ Details on the formulation and validation of the schemes are discussed in Sect. 2. The proposed Python implementations build upon the GT4Py toolchain, and in the remainder of the paper we use the term CLOUDSC-GT4Py to refer to the GT4Py rewrite of CLOUDSC, while the GT4Py ports of the non-linear, tangent-linear, and adjoint formulations of CLOUDSC2 are collectively referred to as CLOUDSC2-GT4Py. The working principles of the GT4Py framework are illustrated in Sect. 3, where we also advocate the advantages offered by domain-specific software approaches. Section 4 sheds some light on the infrastructure code (Ubbiali et al., 2024b) and how it can enable composable and reusable model components. In Sect. 5, we compare the performance of CLOUDSC-GT4Py and CLOUDSC2-GT4Py, as measured on three leadership-class GPU-equipped supercomputers, to established implementations in Fortran and C/C++. We conclude the paper with final remarks and future development paths.

2 Defining the targeted scientific applications

Several physical and chemical mechanisms occurring in the atmosphere are active on spatial scales that are significantly smaller than the highest affordable model resolution. It follows that these mechanisms cannot be properly captured by the resolved model dynamics but need to be parametrized. Parametrizations express the bulk effect of subgrid-scale phenomena on the resolved flow in terms of the grid-scale variables. The equations underneath physical parametrizations are based on theoretical and semi-empirical arguments, and their numerical treatment commonly adheres to the single-column abstraction; therefore adjustments can only happen within individual columns, with no data dependencies between columns. The atmospheric module of the IFS includes parametrizations dealing with the radiative heat transfer, deep and shallow convection, clouds and stratiform precipitation, surface exchange, turbulent mixing in the planetary boundary layer, subgrid-scale orographic drag, non-orographic gravity wave drag, and methane oxidation (ECMWF, 2023).

The focus of this paper is on the cloud microphysics modules of the ECMWF: CLOUDSC – used in operational forecasting – and CLOUDSC2 – employed in the data assimilation. The motivation is threefold:

• i

  Both schemes are among the most computationally expensive parametrizations, with CLOUDSC accounting for up to 10 % of the total execution time of the high-resolution operational forecasts at ECMWF.

• ii

  They are representative of the computational patterns ubiquitous in physical parametrizations.

• iii

  They already exist in the form of dwarfs. The weather and climate “computational dwarfs”, or simply “dwarfs”, are model components shaped into stand-alone software packages to serve as archetypes of relevant computational motifs (Müller et al., 2019) and provide a convenient platform for performance optimizations and portability studies (Bauer et al., 2020).

In recent years, the Performance and Portability team of ECMWF has created the CLOUDSC and CLOUDSC2 dwarfs. The original Fortran codes for both packages, corresponding respectively to the IFS Cycle 41r2 and 46r1, have been pulled out of the IFS code base, slightly polished² and finally made available in public code repositories (https://github.com/ecmwf-ifs/dwarf-p-cloudsc and https://github.com/ecmwf-ifs/dwarf-p-cloudsc2-tl-ad, last access: 23 September 2024). Later, the repositories were enriched with alternative coding implementations, using different languages and programming paradigms; the most relevant implementations will be discussed in Sect. 5.

2.1 CLOUDSC: cloud microphysics of the forecast model

CLOUDSC is a single-moment cloud microphysics scheme that parametrizes stratiform clouds and their contribution to surface precipitation (ECMWF, 2023). It was implemented in the IFS Cycle 36r4 and has been operational at ECMWF since November 2010. Compared to the pre-existing scheme, it accounts for five prognostic variables (cloud fraction, cloud liquid water, cloud ice, rain, and snow) and brings substantial enhancements in different aspects, including treatment of mixed-phase clouds, advection of precipitating hydrometeors (rain and snow), physical realism, and numerical stability (Nogherotto et al., 2016). For a comprehensive description of the scheme, we refer the reader to Forbes et al. (2011) and the references therein. For all the coding versions considered in this paper, including the novel Python rewrite, the calculations are validated by direct comparison of the output against serialized language-agnostic reference data provided by ECMWF.

2.2 CLOUDSC2: cloud microphysics in the context of data assimilation

The CLOUDSC2 scheme represents a streamlined version of CLOUDSC, devised for use in the four-dimensional variational assimilation (4D-Var) at ECMWF (Courtier et al., 1994). The 4D-Var merges short-term model integrations with observations over a 12 h assimilation window to determine the best possible representation of the current state of the atmosphere. This then provides the initial conditions for longer-term forecasts (Janisková and Lopez, 2023). The optimal synthesis between model and observational data is found by minimizing a cost function, which is evaluated using the tangent-linear of the non-linear forecasting model, while the adjoint model is employed to compute the gradient of the cost function (Errico, 1997; Janisková et al., 1999). For the sake of computational economy, the tangent-linear and adjoint operators are derived from a simplified and regularized version of the full non-linear model. CLOUDSC2 is one of the physical parametrizations included in the ECMWF's simplified model, together with radiation, vertical diffusion, orographic wave drag, moist convection, and non-orographic gravity wave activity (Janisková and Lopez, 2023). In the following, we provide a mathematical and algorithmic representation of the tangent-linear and adjoint versions of CLOUDSC2. For the sake of brevity, in the rest of the paper we will refer to the non-linear, tangent-linear, and adjoint formulations of CLOUDSC2 using CLOUDSC2NL, CLOUDSC2TL, and CLOUDSC2AD, respectively.

Let $F:\mathit{x}\mapsto \mathit{y}$ be the functional description of CLOUDSC2, connecting the input fields x with the output variables y. The tangent-linear operator F^′ of F is derived from the Taylor series expansion

$$\begin{array}{}\text{(1)}& F\left(\mathit{x}+\mathit{\delta }\mathit{x}\right)=\mathit{y}+\mathit{\delta }\mathit{y}=F\left(\mathit{x}\right)+F^{\prime }\left[\mathit{x}\right]\left(\mathit{\delta }\mathit{x}\right)+\mathcal{O}\left(\left|\right|\mathit{\delta }\mathit{x}|{|}^{\mathrm{2}}\right),\end{array}$$

where δx and δy are variations on x and y, and $\left|\right|\cdot \left|\right|$ is a suitable norm. The formal correctness of the coding implementation of F^′ can be assessed through the Taylor test (also called the “V-shape” test), which ensures that the following condition is satisfied up to machine precision:

$$\begin{array}{}\text{(2)}& \underset{\mathit{\lambda }\to \mathrm{0}}{lim}{\displaystyle \frac{F\left(\mathit{x}+\mathit{\lambda }\mathit{\delta }\mathit{x}\right)-F\left(\mathit{x}\right)}{F^{\prime }\left[\mathit{x}\right]\left(\mathit{\lambda }\mathit{\delta }\mathit{x}\right)}}=\mathrm{1}\forall \mathit{x},\mathit{\delta }\mathit{x}.\end{array}$$

The logical steps carried out in the actual implementation of the Taylor test are sketched in Algorithm A1 (Appendix A).

The adjoint operator F^* of F^′ is defined such that for the inner product $<\cdot ,\cdot >$:

$$\begin{array}{}\text{(3)}& <\mathit{\delta }\mathit{x},F^{*}\left[\mathit{y}\right]\left(\mathit{\delta }\mathit{y}\right)>=<\mathit{\delta }\mathit{y},F^{\prime }\left[\mathit{x}\right]\left(\mathit{\delta }\mathit{x}\right)>\forall \mathit{x},\mathit{\delta }\mathit{x},\mathit{y},\mathit{\delta }\mathit{y}.\end{array}$$

In particular, Eq. (3) must hold for y=F(x) and $\mathit{\delta }\mathit{y}=F^{\prime }\left[\mathit{x}\right]\left(\mathit{\delta }\mathit{x}\right)$:

$$\begin{array}{}\text{(4)}& \begin{array}{rl}<\mathit{\delta }\mathit{x},& F^{*}\left[F\left(\mathit{x}\right)\right]\left(F^{\prime }\left[\mathit{x}\right]\left(\mathit{\delta }\mathit{x}\right)\right)>=<F^{\prime }\left[\mathit{x}\right]\left(\mathit{\delta }\mathit{x}\right),\\ & F^{\prime }\left[\mathit{x}\right]\left(\mathit{\delta }\mathit{x}\right)>\forall \mathit{x},\mathit{\delta }\mathit{x}.\end{array}\end{array}$$

The latter condition is at the heart of the so-called symmetry test for F^* (see Algorithm A2 in Appendix A).

3 A domain-specific approach to scientific software development

In scientific software development, it is common practice to conceive a first proof-of-concept implementation of a numerical algorithm in a high-level programming environment like MATLAB/Octave (Lindfield and Penny, 2018), or Python. Because these languages do not require compilation and support dynamic typing, they provide a breeding ground for fast prototyping. However, the direct adoption of interpreted languages in HPC has historically been hindered by their intrinsic slowness. To squeeze more performance out of the underlying silicon, the initial proof of concept is translated into either Fortran, C or C++. This leads to the so-called “two-language problem”, where the programming language used for the germinal prototyping is abandoned in favor of a faster language that might be more complicated to use. The lower-level code can be parallelized for shared memory platforms using OpenMP directives, while distributed memory machines can be targeted using Message Passing Interface (MPI) libraries. The resulting code can later be migrated to GPUs, offering outstanding compute throughput and memory bandwidth, especially for single instruction, multiple data (SIMD) applications. GPU porting is accomplished using either OpenACC or OpenMP directives or via a CUDA (https://docs.nvidia.com/cuda/, last access: 23 September 2024) or HIP (https://rocm.docs.amd.com/projects/HIP/en/latest/, last access: 23 September 2024) rewriting, amongst others. To efficiently run the model at scale on multiple GPUs, a GPU-aware MPI build should be chosen, so as to possibly avoid costly memory transfers between host and device and better overlap computations and communications.

The schematic visualization in Fig. 1a highlights how the above workflow leads to multiple implementations of the same scientific application utilizing different programming models and coding styles. This unavoidably complicates software maintainability: ideally, any modification in the numerical model should be encoded in all implementations, so as to preserve the coherency across the hierarchy. The maintainability problem is exacerbated as the number of lines of code, the pool of platforms to support, and the user base increase. This situation has been known as the “software productivity gap” (Lawrence et al., 2018), and we argue that it cannot be alleviated by relying on general-purpose programming paradigms and monolithic code designs. Instead, it calls for a more synergistic collaboration between domain scientists (which here include model developers, weather forecasters, and weather and climate scientists) and computer experts. A path forward is provided by DSLs through separation of concerns (Fig. 1b) so that domain scientists can express the science using syntactic constructs that are aligned with the semantics of the application domain and hide any architecture-specific detail. The resulting source code is thus hardware-agnostic, more concise, easier to read, and easier to manipulate. A toolchain developed by software engineers then employs automatic code generation techniques to synthesize optimized parallel code for the target computer architecture in a transparent fashion.

Figure 1 Diagrams comparing (a) a well-established workflow in scientific software development and (b) a DSL-based approach resembling the software engineering strategy advocated in this paper. The dashed red-and-blue line in (b) marks the separation-of-concerns between the domain scientists and the computer experts.

3.1 The GT4Py framework

GT4Py (https://github.com/GridTools/gt4py, last access: 23 September 2024) is a Python library to generate high-performance implementations of stencil³ kernels as found in weather and climate applications. The library is developed and maintained by CSCS, ETH Zurich, and MeteoSwiss and benefits from important contributions by international partners such as the Paul Allen Institute for Artificial Intelligence (AI²). The choice of embedding the GT4Py framework in Python has been mainly dictated by the following factors:

• i.

  Python is taught in many academic courses due its clean, intuitive, and expressive syntax; therefore a significant fraction of early-career domain scientists are exposed to the language.

• ii.

  It admits a powerful ecosystem of open-source packages for building end-to-end applications.

• iii.

  It is possible to seamlessly interface Python with lower-level languages with minimal overhead and virtually no memory copies.

• iv.

  Under the thrust of the artificial intelligence and machine learning community (AI/ML), the popularity and adoption of Python across the whole scientific community is constantly growing, as opposed to Fortran (Shipman and Randles, 2023).

The proposed Python implementations of CLOUDSC and CLOUDSC2 are based on the first public release of GT4Py, which only supports Cartesian grids. The latest advancements to support unstructured meshes (contained in the sub-package gt4py.next) are not discussed in this study.

Figure 2 Simplified view of the internal stages carried out by the GT4Py toolchain to generate a high-performance CPU or GPU implementation of the three-dimensional Laplacian stencil starting from its GTScript definition. For the sake of visualization, only two intermediate representations (IRs) are included: the GridTools IR (GTIR) and the implementation IR (IIR).

Figure 2 showcases the main steps undertaken by the GT4Py toolchain to translate the high-level definition of the three-dimensional Laplacian operator into optimized code, which can be directly called from within Python. The stencil definition is given as a regular Python function using the GTScript DSL. GTScript abstracts spatial for-loops away: computations are described for a single point of a three-dimensional Cartesian grid and can be controlled with respect to the vertical index bounds using the interval context manager. Vertical loops are replaced by computation contexts, which define the iteration order along the vertical axis: either PARALLEL (meaning no vertical data dependencies between horizontal planes), FORWARD, or BACKWARD. Each assignment statement within a computation block can be thought of as a loop over a horizontal plane; no horizontal data dependencies are allowed. Neighboring points are accessed through relative offsets, with the first two offsets being the horizontal offsets and the last offset being the vertical offset.

Any function marked with the gt4py.cartesian.gtscript.stencil decorator is translated by the GT4Py frontend into a hierarchy of tree-like intermediate representations (IRs), featuring different levels of abstractions to accommodate diverse optimizations and transformations (Gysi et al., 2021). The lowest-level IR (denoted as implementation IR, or IIR) is consumed by the backends to generate code that is either optimized for a given architecture or suited to a specific purpose. The following backends are currently available:

• NumPy (Harris et al., 2020) is the de facto standard for array computing in Python and can be used for debugging and fast-prototyping.

• GridTools (Afanasyev et al., 2021) is a set of libraries and utilities to write performance-portable applications in the area of weather and climate.

• DaCe (Ben-Nun et al., 2019) is a parallel programming framework, which internally uses the Stateful DataFlow multiGraph (SDFG) data-centric intermediate representation to decouple domain science and performance engineering.

The generated code is compiled under the hood, and Python bindings for the resulting executable are automatically produced; therefore the stencil can finally be executed by passing the input and output fields and by specifying the origin and size of the computation domain. GT4Py provides convenient utilities to allocate arrays with an optimal memory layout for any given backend, relying on NumPy for CPU storages and CuPy (Okuta et al., 2017) for GPU storages. Concerning GPU computing, we highlight that GT4Py supports both NVIDIA and AMD GPUs.

A more realistic and pertinent code sample is provided in Listing 1. It is an abridged GT4Py implementation of the procedure computing the saturation water vapor pressure as a function of air pressure and temperature. The code is extracted from the CLOUDSC2-GT4Py dwarf and highlights two additional features of GTScript: functions and external symbols. Functions can be thought of as macros and can be used to improve composability, reusability, and readability. External symbols are used to encode those scalar parameters (e.g., physical constants) that are kept constant throughout a simulation and might only change between different model setups. External values must be provided at stencil compilation time. The functionalities provided by the package ifs_physics_common will be discussed in the following section.

Listing 1 GTScript (the Python-embedded DSL exposed by GT4Py) functions and stencil computing the saturation water vapor pressure given the air pressure and temperature. Abridged excerpt from the CLOUDSC2-GT4Py dwarf.

4 Infrastructure code

All stencils of CLOUDSC-GT4Py and CLOUDSC2-GT4Py are defined, compiled, and invoked within classes that leverage the functionalities provided by the Sympl package (Monteiro et al., 2018). Sympl is a toolset of Python utilities to write self-contained and self-documented model components. Because the components share a common application public interface (API), they favor modularity, composability, and inter-operability (Schär et al., 2019). These aspects are of utter importance, for instance, when it comes to assessing the impact of process coupling on weather forecasts and climate projections (Ubbiali et al., 2021).

Listing 2 A Python class to compute the saturation water vapor pressure given the air pressure and temperature. Abridged excerpt from the CLOUDSC2-GT4Py dwarf.

Sympl components interact through dictionaries whose keys are the names of the model variables (fields) and whose values are xarray's DataArrays (Hoyer and Hamman, 2017), collecting the grid point values, the labeled dimensions, the axis coordinates, and the units for those variables. The most relevant component exposed by Sympl is TendencyComponent, producing tendencies for prognostic variables and retrieving diagnostics. The class defines a minimal interface to declare the list of input and output fields and initialize and run an instance of the class. This imposes minor constraints on model developers when writing a new physics package.

The bespoke infrastructure code for CLOUDSC-GT4Py and CLOUDSC2-GT4Py is bundled as an installable Python package called ifs_physics_common. Not only does it build upon Sympl, but the package also provides grid-aware and stencil-oriented functionalities. Both the CLOUDSC cloud microphysics and the non-linear, tangent-linear, and adjoint formulations of CLOUDSC2 are encoded as stand-alone TendencyComponent classes settled over a ComputationalGrid. The latter is a collection of index spaces for different grid locations. For instance, (I, J, K) corresponds to cell centers, while (I, J, K-1/2) denotes vertically staggered grid points. For any input and output field, its name, units, and grid location are specified as class properties. When running the component via the dunder method __call__, Sympl transparently extracts the raw data from the input DataArrays according to the information provided in the class definition. This step may involve unit conversion and axis transposition. The resulting storages are forwarded to the method array_call, which carries out the actual computations, possibly by executing GT4Py stencil kernels.

Listing 2 brings a concrete example from CLOUDSC2-GT4Py: a model component leveraging the stencil defined in Listing 1 to compute the saturation water vapor pressure. The class inherits DiagnosticComponent, a stripped-down version of TendencyComponent, which only retrieves diagnostic quantities. Within the instance initializer __init__, the stencil from Listing 1, registered using the decorator ifs_physics_common.framework.stencil.stencil_collection, is compiled using the utility method compile_stencil. The options configuring the stencil compilation (e.g. the GT4Py backend) are fetched from the data class GT4PyConfig.

5 Performance analysis

In this section, we highlight the results from comprehensive performance testing. We compare the developed CLOUDSC-GT4Py and CLOUDSC2-GT4Py codes against reference Fortran versions and various other programming prototypes. The simulations were performed on three different supercomputers:

• i.

  Piz Daint (https://www.cscs.ch/computers/piz-daint, last access: 23 September 2024), an HPE Cray XC40/XC50 system installed at CSCS in Lugano, Switzerland;

• ii.

  MeluXina (https://docs.lxp.lu/, last access: 23 September 2024), an ATOS BullSequana XH2000 machine hosted by LuxConnect in Bissen, Luxembourg, and procured by the EuroHPC Joint Undertaking (JU) initiative;

• iii.

  the Cray HPE EX235a supercomputer LUMI (https://docs.lumi-supercomputer.eu/, last access: 23 September 2024), an EuroHPC pre-exascale machine at the Science Information Technology Center (CSC) in Kajaani, Finland.

On each machine, the CLOUDSC and CLOUDSC2 applications are executed on a single hybrid node that sports one or multiple GPU accelerators alongside the host CPU. An overview of the node architectures for the three considered supercomputers can be found in Table 1.

Table 1 Overview of the node architecture for the hybrid partition of Piz Daint, MeluXina, and LUMI. Only the technical specifications which are most relevant for the purposes of this paper are reported.

Besides the GT4Py codes, we involve up to four alternative lower-level programming implementations, which will be documented in an upcoming publication:

• a.

  The baseline is Fortran enriched with OpenMP directives for multi-threading execution on CPU.

• b.

  An optimized GPU-enabled version based on OpenACC using the single-column coalesced (SCC) loop layout in combination with loop fusion and temporary local array demotion (so-called “k-caching”) is considered. While the SCC loop layout yields more efficient access to device memory and increased parallelism, the k-caching technique significantly reduces register pressure and memory traffic. This is achieved via loop fusion to eliminate most loop-carried dependencies and consequently allows temporaries to be demoted to scalars.

• c.

  The currently best-performing Loki-generated and GPU-enabled variant is also taken into account.

• d.

  Finally, an optimized GPU-enabled version of CLOUDSC including k-caching is used. The code is written in either CUDA or HIP, to target both NVIDIA GPUs (shipped with Piz Daint and MeluXina) and AMD GPUs (available on LUMI).

Table 2 For each coding version of the CLOUDSC and CLOUDSC2 dwarfs considered in the performance analysis, the table reports the compiler suite used to compile the codes on Piz Daint, MeluXina, and LUMI. The codes are compiled with all major optimization options enabled. Those implementations which are either not available or not working are marked with a dash; more details, as well as a high-level description of each coding implementation, are provided in the text.

Table 2 documents the compiler specifications employed for each of the programming implementations, on Piz Daint, MeluXina, and LUMI. We consistently apply the most aggressive optimization, ensuring that the underlying code manipulations do not harm validation. For the different algorithms at consideration, validation is carried out as follows:

• For CLOUDSC and CLOUDSC2NL, the results from each coding version are directly compared with serialized reference data produced on the CPU. For each output field, we perform an element-wise comparison using the NumPy function allclose (https://numpy.org/devdocs/reference/generated/numpy.allclose.html, last access: 23 September 2024). Specifically, the GT4Py rewrites validate on both CPU and GPU, with an absolute and relative tolerance of 10⁻¹² and 10⁻¹⁸ when employing double precision. When reducing the precision to 32 bits, the absolute and relative tolerance levels need to be increased to 10⁻⁴ and 10⁻⁷ on CPU and 10⁻² and 10⁻⁷ on GPU. In the latter case, we observe that the field representing the enthalpy flux of ice still does not pass validation. We attribute the larger deviation from the baseline data on the device to the different instruction sets underneath CPUs and GPUs.

• All implementations of CLOUDSC2TL and CLOUDSC2AD are validated using the Taylor test (see Algorithm A1) and the symmetry test (see Algorithm A2), respectively. In this respect, we emphasize that the GT4Py implementations satisfy the conditions of both tests on all considered computing architectures, regardless of whether double or single precision is employed.

The source repositories for CLOUDSC and CLOUDSC2 dwarfs may include multiple variants of each reference implementation, varying for the optimization strategies. In our analysis, we always take into account the fastest variant of each alternative implementation; for the sake of reproducibility, Table 3 contains the strings identifying the coding versions at consideration and the corresponding NPROMA⁴ employed in the runs. Similarly, for all Python implementations we consider only the most performant backends of GT4Py: the GridTools C++ CPU backend with k-first memory layout, and the DaCe GPU backend.

Table 3 For each reference implementation of the CLOUDSC and CLOUDSC2 dwarfs, the table reports the string identifying the specific variant considered in the performance analysis on Piz Daint, MeluXina, and LUMI. The corresponding NPROMA is provided within parentheses. Those implementations which are either not available or not working are marked with a dash.

Figure 3 Execution time on a single NUMA domain of a hybrid node of the Piz Daint supercomputer for CLOUDSC (a, d), CLOUDSC2NL (b, e), and the symmetry test for CLOUDSC2TL and CLOUDSC2AD (c, f) using either double-precision (a, b, c) or single-precision (d, e, f) floating point arithmetic. The computational domain consists of 65 536 columns and 137 vertical levels. Displayed are the multi-threaded Fortran baseline using OpenMP (grey); two GPU-accelerated Fortran implementations, using either OpenACC directives (lime) or the source-to-source translation tool Loki (yellow); an optimized CUDA C version (green); and the GT4Py rewrite, using either the GridTools C++ CPU backend with k-first data ordering (blue) or the DaCe GPU backend (orange). All numbers should be interpreted as an average over 50 realizations. The panels only show the code versions available and validating at the time of writing.

For the interpretation of the CPU-versus-GPU performance numbers, we note that host codes are executed on all the cores available on a single non-uniform memory access (NUMA) domain of a compute node, while device codes are launched on the GPU attached to that NUMA domain. In a distributed-memory context, this choice allows the same number of MPI ranks to be fit per node, on either CPU or GPU. Table 1 reports the number of NUMA partitions per node for Piz Daint, MeluXina, and LUMI, with the compute and memory resources being evenly distributed across the NUMA domains. Note that the compute nodes of the GPU partition of LUMI have the low-noise mode activated, which reserves one core per NUMA domain to the operating system; therefore only seven out of eight cores are available to the jobs. Moreover, we highlight that each MI250X GPU consists of two graphics compute dies (GCDs) connected via four AMD Infinity Fabric links but not sharing physical memory. From a software perspective, each compute node of LUMI is equipped with eight virtual GPUs (vGPUs), with each vGPU corresponding to a single GCD and assigned to a different NUMA domain.

Figure 4 As Fig. 3 but for the MeluXina supercomputer.

Figures 3–5 visualize the execution times for CLOUDSC (left column), CLOUDSC2NL (center column), and the symmetry test for CLOUDSC2TL and CLOUDSC2AD (right column) for Piz Daint, MeluXina, and LUMI, respectively.⁵ All performance numbers refer to a grid size of 65 536 columns, with each column featuring 137 vertical levels. In each figure, execution times are provided for simulations running either entirely in double precision (corresponding to the 64-bit IEEE format and denoted as FP64; top row) or in single precision (corresponding to the 32-bit IEEE format and denoted as FP32; bottom row). Within each panel, the plotted bars reflect the execution time of the various codes, with a missing bar indicating the corresponding code (non-GT4Py) is either not available or not working properly, specifically as follows:

• The Fortran version of CLOUDSC2AD can only run on a single OpenMP thread on MeluXina (the issue is still under investigation).

• A native GPU-enabled version of CLOUDSC using 32-bit floating point arithmetic does not exist at the time of writing, and no CUDA/HIP implementations are available for CLOUDSC2.

• All Fortran-based implementations of the three formulations of CLOUDSC2 can only use double-precision computations.

• A Loki version of CLOUDSC2TL and CLOUDSC2AD is not available at the time of writing.

Notably, we find the GT4Py rewrite of both CLOUDSC and CLOUDSC2 to be very robust, as the codes execute on every CPU and GPU architecture included in the study and can always employ either double- or single-precision floating point arithmetic. With GT4Py, changing the backend with the respective target architecture, or changing the precision of computations, is as easy as setting a namelist parameter. Moreover, at the time of writing, the GT4Py implementations of the more complex tangent-linear and adjoint formulations of CLOUDSC2 were the first codes enabling GPU execution, again both in double or single precision.

Figure 5 As Fig. 3 but for the LUMI supercomputer.

The performance of the high-level Python with GT4Py compares well against Fortran with OpenACC. The runtimes for GT4Py with its DaCe backend versus OpenACC are similar on Piz Daint, MeluXina, and LUMI. One outlier is the double-precision result on LUMI, for which the OpenACC code appears relatively slow. We suppose this behavior is associated with the insufficient OpenACC support for the HPE Cray compiler. Only the HPE Cray compiler implements GPU offloading capabilities for OpenACC directives on AMD GPUs, meaning that Fortran OpenACC codes require an HPE Cray platform to run on AMD GPUs. In contrast, GT4Py relies on the HIPCC compiler driver developed by AMD to compile device code for AMD accelerators, and this guarantees a proper functioning irrespective of the machine vendor. We further note that the DaCe backend of GT4Py executes roughly 2 times faster on MeluXina's NVIDIA A100 GPUs than on LUMI's AMD Instinct MI250X GPUs. As mentioned above, from a software perspective, each physical GPU module on LUMI is considered two virtual GPUs; therefore the code is actually executed on half of a physical GPU card. Hence we can speculate that using both dies of an AMD Instinct MI250X GPU, performance would be on a par with the NVIDIA A100 GPU.

Another interesting result is that both CLOUDSC-GT4Py and CLOUDSC2-GT4Py are consistently faster than the implementations generated with Loki. Loki allows bespoke transformation recipes to be built to apply changes to programming models and coding styles in an automated fashion. Therefore, GPU-enabled code can be produced starting from the original Fortran by, e.g., automatically adding OpenACC directives. However, because not all optimizations are encoded yet in the transformations, the Loki-generated device code cannot achieve optimal performance. Notwithstanding, source-to-source translators such as Loki are of high relevance for enabling GPU execution with large legacy Fortran code bases.

As used in this paper, GT4Py cannot yet attain the performance achieved by manually optimized native implementations with either Fortran on CPU or CUDA/HIP on GPU. Multi-threaded Fortran can be up to 3 times faster than the GridTools CPU backend of GT4Py using the k-first (C-like) memory layout, while the DaCe GPU backend of GT4Py can be up to a factor of 2 slower than CUDA/HIP. On the one hand, so far the development of GT4Py has been focused on GPU execution (see, e.g., Dahm et al., 2023) because this will be the dominant hardware for time-critical applications in the years to come. On the other hand, we stress that the k-caching CUDA and HIP variants of CLOUDSC were semi-automatically generated by performance engineering experts, starting from an automatic Fortran-to-C transpilation of the SCC variants and manually applying additional optimizations that require knowledge about the specific compute patterns in the application. This process is not scalable to the full weather model and not a sustainable code adaptation method. In contrast, no significant performance engineering (by, e.g., loop fusing and reducing the number of temporary fields) has been applied yet with CLOUDSC-GT4Py and CLOUDSC2-GT4Py.

Figure 6 For the GT4Py rewrites of CLOUDSC (a, d), CLOUDSC2NL (b, e), and the symmetry test for CLOUDSC2TL and CLOUDSC2AD (c, f), fraction of the total execution time spent within the stencil computations (full bars) and the Python side of the application (hatched bars) on Piz Daint, MeluXina, and LUMI. Results are shown for the GridTools C++ CPU backend with k-first data ordering (blue) and the DaCe GPU backend (orange), using either double-precision (a, b, c) or single-precision (d, e, f) floating point arithmetic.

To rule out the possibility that the performance gap between the Python DSL and lower-level codes is associated with overhead originating from Python, Fig. 6 displays the fraction of runtime spent within the stencil code generated by GT4Py and the high-level Python code of the application (infrastructure and framework code; see Sect. 4). Across the three supercomputers, the Python overhead decreases as (i) the complexity and length of computations increase, (ii) the peak throughput and bandwidth delivered by the hardware underneath decrease, and (iii) the floating point precision increases. On average, the Python overhead accounts for 5.4 % of the total runtime on GPU and 0.4 % on CPU; the latter corresponds to 0.7 % relative to the Fortran execution time.

Finally, we observe a significant sensitivity of the GPU performance with respect to the thread block size⁶: for values smaller than 128, performance is degraded across all implementations, with the gap between CUDA/HIP and GT4Py+DaCe being smaller. This shows that some tuning and toolchain optimizations can be performed to improve performance with the DSL approach.

6 Conclusions

The CLOUDSC and CLOUDSC2 cloud microphysics schemes of the IFS at ECMWF have served as demonstrators to show the benefits of a high-level domain-specific implementation approach to physical parametrizations. We presented two Python implementations based on the GT4Py framework, in which the scientific code is insulated from hardware-specific aspects. The resulting application provides a productive user interface with enhanced readability and maintainability and can run efficiently and in a very robust manner across a wide spectrum of compute architectures/systems. The approach can be powerful in the light of the increasingly complex HPC technology landscape, where general-purpose CPUs are increasingly complemented by domain-specific architectures (DSAs) such as GPUs, tensor processor units (TPUs), field-programmable gate arrays (FPGAs), and application-specific integrated circuits (ASICs). In addition to the CLOUDSC scheme used in the IFS forecast model, we have presented results with the GT4Py rewrites of the non-linear, tangent-linear, and adjoint formulations of CLOUDSC2 used in data assimilation. In both CLOUDSC-GT4Py and CLOUDSC2-GT4Py, the stencil kernels are encapsulated within model components sharing a minimal and clear interface. By avoiding any assumption on the host model, the interface aims to provide interoperable and plug-and-play physical packages, which can be transferred more easily between different modeling systems with virtually no performance penalty.

We carried out a comprehensive study to assess the portability of the Python codes across three major supercomputers, differing in terms of the vendor, the node architecture, the software stack, and the compiler suites. We showed that the GPU performance of GT4Py codes is competitive against optimized Fortran OpenACC implementations and performs particularly well when compared to the available codes generated with the Loki source-to-source translation tool. Low-level implementations written in either Fortran for CPUs or CUDA/HIP for GPUs, with additional optimizations that possibly require knowledge about the specific compute patterns, can provide better performance but are extremely challenging to create and maintain for entire models. The CPU performance of GT4Py is currently suboptimal, but this is an expected result given the focus on GPUs in the DSL development so far, and we clearly expect this to improve significantly with upcoming and future GT4Py versions. Indeed, since the DSL can accommodate any specific low-level optimization, attaining the same performance as native, expert-written Fortran and CUDA/HIP models is feasible and will be the target of our future efforts.

The presented results, based on a representative physical parametrization and considering tangent-linear and adjoint versions, add to the notion that weather and climate model codes can execute significantly faster on GPUs (Fuhrer et al., 2018), and the number of HPC systems with accelerators is steadily increasing⁷. Therefore, we envision that CPUs will be increasingly relegated to tasks that are not time-critical.

The current study supports our ongoing efforts and plans to port other physical parametrizations to Python with GT4Py. However, we note that GT4Py was originally devised to express the computational motifs of dynamical cores based on grid point methods, so not all patterns found in the parametrizations are natively supported by the DSL. In this respect, current limitations of the DSL exist for higher-dimensional fields (e.g., arrays storing a tensor at each grid point), but again we expect these to be fully supported for HPC with future extensions of GT4Py. In addition, Python offers a variety of alternative libraries that can be employed in conjunction with GT4Py to generate fast machine code.

Appendix A: Algorithmic description of the Taylor test and symmetry test for CLOUDSC2

In Sect. 2.2, we briefly described the aim and functioning of the Taylor test and the symmetry test for CLOUDSC2. Here, we detail the logical steps performed by the two tests with the help of pseudo-codes encapsulated in Algorithms A1 and A2.

Algorithm A1 The Taylor test assessing the formal correctness of the coding implementation of the tangent-linear formulation of CLOUDSC2, denoted as CLOUDSC2TL. The three-dimensional arrays x and y collect the grid point values for all nin input fields and nout output fields of CLOUDSC2, respectively. The corresponding variations are δx and δy. The grid consists of ncol columns, each containing nlev vertical levels. Note that compared to its functional counterpart $F^{\prime }\left[\mathit{x}\right]:\mathit{\delta }\mathit{x}\mapsto \mathit{\delta }\mathit{y}$, CLOUDSC2TL(x, δx) returns both y and δy. The coding implementation of the non-linear CLOUDSC2 is indicated as CLOUDSC2NL.

Algorithm A2 The symmetry test assessing the formal correctness of the coding implementation of the adjoint formulation of CLOUDSC2, denoted as CLOUDSC2AD. The machine epsilon is indicated as ε; all other symbols have the same meaning as in Algorithm A1. Note that compared to its functional counterpart $F^{*}\left[F\left(\mathit{x}\right)\right]:\mathit{\delta }\mathit{y}\mapsto \mathit{\delta }{\mathit{x}}^{*}$, CLOUDSC2AD(x, δy) returns both y and δx^*.