Whole Program Generation of Massively Parallel Shallow Water Equation Solvers

The study of ocean currents has been an active area of research for decades. As a model close to the water surface, the shallow water equations (SWE) can be used. For realistic simulations, efficient numerical solvers are necessary that exhibit a good node-level performance while still maintaining scalability. When comparing the discretized model and the actual implementation, one often finds that they differ vastly. This gap makes it hard for domain experts to implement their models and high performance computing (HPC) experts are required to ensure an optimal implementation. Using domain-specific languages (DSLs) and code generation techniques can be a useful tool to bridge this gap. In recent years, ExaStencils and its DSL ExaSlang have proven to provide a suitable platform for this. We present an extension from up to now elliptic to hyperbolic partial differential equations (PDEs) in this work, namely the SWE. After setting up a suitable discretization, we demonstrate how it can be mapped to ExaSlang code. This code is still quite similar to the original, mathematically motivated specification and can be easily written by domain experts. Still, solvers generated from this abstract representation can be run on large-scale clusters. We demonstrate this by giving performance and scalability results on the state-of-the-art GPU cluster Piz Daint where we solve for close to a trillion unknowns on 2048 GPUs. From there, we discuss the performance impact of different optimizations such as overlapping computation and communication, or switching to a hybrid CPU-GPU parallelization scheme.