Using Julia for the next generation of HPC software for geodynamic modelling

Following the long-standing paradigm in HPC, computational geodynamic codes have been typically written in high-level statically typed and compiled languages, namely C/C++ and Fortran. The low productivity rates of these languages led to the so-called two-language problem, where dynamic languages such as Python or MATLAB are used for prototyping purposes, before porting the algorithms to high-performance languages. The Julia programming language aims at bridging the productivity and advantages of such dynamic languages without sacrificing the performance provided by static languages. The high performance of Julia, combined with high-productivity rates and other powerful tools, such as advanced meta-programming (i.e. code generation), make Julia a suitable candidate for the next generation of HPC-ready scientific software.

We introduce the open-source and Julia-written package JustRelax.jl (https://github.com/PTsolvers/JustRelax.jl) as a way forward for the next generation of geodynamic codes. JustRelax.jl is a production-ready API for a collection of highly-efficient numerical solvers (Stokes equations, diffusion, etc.) based on the embarrassingly parallel pseudo-transient method. We rely on ParallelStencil.jl (https://github.com/omlins/ParallelStencil.jl), which leverages the advanced meta-programming capabilities of Julia to generate efficient computational kernels agnostic to the back-end system (i.e. Central Processing Unit (CPU) or Graphics Processing Unit (GPU)). Using ImplicitGlobalGrid.jl (https://github.com/eth-cscs/ImplicitGlobalGrid.jl) to handle the MPI and CUDA-aware MPI communication, these computational kernels run seamlessly in local shared-memory workstations and distributed memory and multi-GPU HPC systems with little effort for the front-end user.

Efficient computation of the (local) physical properties of different materials is another key feature required in geodynamic codes, for which we employ GeoParams.jl (https://github.com/JuliaGeodynamics/GeoParams.jl). This package provides lightweight, optimised, and reproducible computation of different material properties (e.g. advanced rheological laws, density, seismic velocity, etc.), amongst other available features. GeoParams.jl is also carefully designed to support CPU and GPU devices, and be fully compatible with other external packages, such as ParallelStencil.jl and existing auto-differentiation packages.

We finally show high-resolution examples of geodynamic models run on multi-GPU HPC systems employing the presented open-source Julia tools.