asQ: parallel-in-time finite element simulations using ParaDiag for geoscientific models and beyond

TL;DR

This work introduces asQ, a Firedrake/PETSc-based library for prototyping ParaDiag-II parallel-in-time methods on finite element models, with emphasis on geoscientific applications. It explains the ParaDiag framework, including linear and nonlinear all-at-once systems, alpha-circulant preconditioning, and a simple performance model to predict speedups. Through a suite of linear and nonlinear test cases (advection, shallow water, and compressible Euler), the authors demonstrate strong time-scaling and substantial speedups in linear settings, while highlighting the limitations posed by averaging errors and complex-block conditioning in nonlinear problems. asQ provides a productive testbed for exploring ParaDiag strategies, enabling rapid testing of problem sets and solver configurations on large-scale HPC systems, and outlining avenues for extending time integrators and preconditioners to improve robustness and efficiency in realistic geophysical simulations.

Abstract

Modern high performance computers are massively parallel; for many PDE applications spatial parallelism saturates long before the computer's capability is reached. Parallel-in-time methods enable further speedup beyond spatial saturation by solving multiple timesteps simultaneously to expose additional parallelism. ParaDiag is a particular approach to parallel-in-time based on preconditioning the simultaneous timestep system with a perturbation that allows block diagonalisation via a Fourier transform in time. In this article, we introduce asQ, a new library for implementing ParaDiag parallel-in-time methods, with a focus on applications in the geosciences, especially weather and climate. asQ is built on Firedrake, a library for the automated solution of finite element models, and the PETSc library of scalable linear and nonlinear solvers. This enables asQ to build ParaDiag solvers for general finite element models and provide a range of solution strategies, making testing a wide array of problems straightforward. We use a quasi-Newton formulation that encompasses a range of ParaDiag methods, and expose building blocks for constructing more complex methods. The performance and flexibility of asQ is demonstrated on a hierarchy of linear and nonlinear atmospheric flow models. We show that ParaDiag can offer promising speedups and that asQ is a productive testbed for further developing these methods.

Figures (15)

• Figure 1: $\psi_{j}=(\lambda_{1,j}/\lambda_{2,j})/(1/(\Delta t\theta))$ in the complex plane for varying $N_{t}$ with $\theta=0.5$ and $\alpha=10^{-4}$. As $N_{t}$ increases the $\psi_{j}$ for low frequencies cluster towards the imaginary axis.
• Figure 2: Space-time parallelism using Firedrake's Ensemble class. ensemble.comm.size is $P_{x} = 5$ and ensemble.ensemble_comm.size is $P_{t} = 5$. The dark blue horizontal lines each represent a spatial communicator Ensemble.comm over which a mesh is partitioned. Every spatial communicator has the same number of ranks and the same mesh partitioning. The grey vertical lines represent communicators in time, Ensemble.ensemble_comm. These communicators are responsible for connecting the ranks on each spatial communicator with the same section of the mesh partition.
• Figure 3: Strong scaling in time of the linear scalar advection equation test case with $N_{x}=128^{2}$, $\sigma=0.8$ is demonstrated by measured wallclock times per timestep ($T_{p}/N_{t}$) for varying $N_{t}$ and $\alpha$. The wallclock time for the serial-in-time method is plotted as the horizontal line $T_{s}$. The prediction of the performance model $T_{s}/S$ (\ref{['eq:speedup']}) is shown for $\alpha=10^{-6}$ using the measured $T_{s}$ and $S=N_{t}/\left(2(m_{p}+1)\right)$ (assuming $T_{c}=0$). Maximum speedup of 528 achieved with $\alpha=10^{-6}$ and $N_{t}=16384$.
• Figure 4: Weak scaling in time profile of the linear scalar advection equations test case with $N_{x}=128^{2}$, $\sigma=0.8$, $\alpha=10^{-4}$, and varying window length $N_{t}$. Measured wallclock times per window for each component of the algorithm: total, block solves, space-time transpose, (I)FFTs, all-at-once function evaluation, and all-at-once Jacobian action.
• Figure 5: Colour plot of the free surface elevation $h$ in the linear rotating shallow water equation testcase.