GPU4GEO: Frontier GPU multi-physics solvers 

PI: Paul J. Tackley 

Co-PIs: Taras Gerya, Boris J. P. Kaus, Alan Edelman, Ludovic Räss, Patrick Sanan 

July 1, 2021 – June 30, 2024

Project Summary

Computational Earth science often relies on modelling to understand complex physical systems which cannot be directly observed. Two domains where this is true are geodynamics, concerned with the long-term dynamics of the solid and liquid rock that compose planets like the Earth, and ice sheet dynamics, studying how ice flows in places like Greenland and Antarctica. Despite their apparent differences, the physical models and ensuing computational challenges in these two domains are extremely similar. Both domains require extremely large simulations which are not possible with the state of the art in HPC, thus requiring next-generation machines. Here, we focus on a high-resolution geodynamic model of region surrounding the Alps, and large-scale models of Earth’s larges ice sheets, Greenland and Antarctica.

The GPU4GEO project proposes software tools which provide a way forward in these two domains by exploiting two powerful emerging paradigms in HPC: massively parallel relaxation-based solvers, and HPC with Julia.

Solvers based on pseudo-transient relaxation on staggered grids show great promise for solving a wide range of mechanical multi-physics problems in geoscience, at scale and on GPU-accelerated clusters. These offer reasonable complexity with growing problem size yet map to compact local operations, ideal for achieving high performance on the current and nextgeneration of supercomputers, featuring nodes with multiple GPUs rewarding fine-grained parallelization with homogeneous tasks.

Our proposed tools encapsulate these proven techniques and make them available in a reusable, extensible, and highperformance framework, so that they may be applied within existing application codes and used to develop new ones. Supported physics will include: diffusion processes; variable-viscosity Stokes flow with nonlinear rheology, elasticity, and plasticity; 2-phase (Darcy) flow, coupled to reactions or heat diffusion; and coupled Stokes, heat diffusion, and Darcy flow models.

We propose a Julia package, JustRelax.jl, which provides solvers for many physical systems of interest, and building blocks to address more. It useMPI for distributed-memory parallelism and supports operations on multi-core CPUs and GPUs.

We use Julia as our main language because (a) we can draw on existing, high-performance Julia implementations, (b) CSCS has a pole position in Interactive Supercomputing with Julia, being one of the first centers having shown to scale Julia on 1000’s of GPUs on their supercomputer (c) Julia’s high-level abstraction and multiple dispatch capabilities make it amenable to portability between backends (e.g. multicore CPUs and AMD, NVIDIA, and (upcoming) Intel GPUs) (d) Julia solves the two-language problem and promises increased productivity, through its central design feature of being both high-level and performance-oriented, (e) it enables powerful meta-programming to render high-performance code generation straightforward, and (f) it can be compiled to shared libraries callable from C or other languages, allowing to integrate with many legacy HPC codes.

We champion this last cutting-edge paradigm in which a C library, librelax, is derived from our Julia package, JustRelax.jl, thus making our tools available to a wide range of existing HPC applications and frameworks. Through this interface, we are easily able to integrate with the PETSc library and the many applications which rely on it, by defining a new nonlinear solver implementation, SNESRelax.

To ensure that this cutting-edge technique is properly optimized, we have direct support from core Julia developers and Julia Computing. We also propose close collaboration with Dr. Samuel Omlin (CSCS), who is not only a leading researcher on the methods on which we base our solvers, but also an expert on the advanced I/O (ADIOS2) and in-situ visualization tools we will integrate with our software.

We will integrate our tools with several geodynamics, ice sheet dynamics, and other geophysical flow applications, in Julia, C, C++, and Fortran. With this integration, our applications will be prepared to make a pair of Tier-0 allocation requests, one in geodynamics and one in ice sheet modelling, for machines like the successor to Piz Daint or LUMI.

To allow for meaningful performance comparisons and to engage the community, we define a set of “bake-off problems”: well-defined benchmark cases by which different applications and frameworks may compare performance.

This highly collaborative project would represent an exceptional confluence of expertise and would unlock exciting results for all parties. It provides a broad set of applications, including two with plans for very large simulations, with computational tools extremely well-suited to answer pressing scientific questions on current and next-generation GPUaccelerated architectures. It will forcefully demonstrate Julia’s place in the HPC ecosystem, both as a primary language and as a backend language. It will leverage CSCS’s investments in next-generation GPU-based systems, interactive HPC,and high-productivity development. It will provide tools of use to a wide community, and simulations which will demonstrate their effectiveness at scale, with simulations of interest to both experts and the broader scientific community and general public.