Open AccessFast, Cheap, and Turbulent—Global Ocean Modeling With GPU Acceleration in Python
Author(s) -
Häfner Dion,
Nuterman Roman,
Jochum Markus
Publication year - 2021
Publication title -
journal of advances in modeling earth systems
Language(s) - English
Resource type - Journals
SCImago Journal Rank - 3.03
H-Index - 58
ISSN - 1942-2466
DOI - 10.1029/2021ms002717
Subject(s) - computer science , python (programming language) , fortran , central processing unit , computational science , parallel computing , cuda , compiler , interfacing , supercomputer , graphics processing unit , subroutine , operating system , computer hardware
Even to this date, most earth system models are coded in Fortran, especially those used at the largest compute scales. Our ocean model Veros takes a different approach: it is implemented using the high‐level programming language Python. Besides numerous usability advantages, this allows us to leverage modern high‐performance frameworks that emerged in tandem with the machine learning boom. By interfacing with the JAX library, Veros is able to run high‐performance simulations on both central processing units (CPU) and graphical processing unit (GPU) through the same code base, with full support for distributed architectures. On CPU, Veros is able to match the performance of a Fortran reference, both on a single process and on hundreds of CPU cores. On GPU, we find that each device can replace dozens to hundreds of CPU cores, at a fraction of the energy consumption. We demonstrate the viability of using GPUs for earth system modeling by integrating a global 0.1° eddy‐resolving setup in single precision, where we achieve 1.3 model years per day on a single compute instance with 16 GPUs, comparable to over 2,000 Fortran processes.