OMEGA: The Ocean Model for E3SM Global Applications: A New High Performance Computing Code for Exascale Architectures

The US Department of Energy (DOE) conducts climate simulations on some of the world’s largest supercomputers. These exascale machines use heterogeneous architectures with both CPUs and GPUs, and scientific codes must adapt to make full use of this computing power. Programmers have several choices: (1) writing vendor-specific GPU code like CUDA, HIP, and SYCL; (2) adding OpenACC pragmas for the GPUs; or (3) using performance-portable programming models such as Kokkos, YAKL, Raja that provide programmer-friendly abstractions that are translated to vendor-specific backends.

Here we introduce Omega: The Ocean Model for E3SM Global Applications, which is specifically designed for modern exascale computers. It uses the third strategy of calling external libraries that have been optimized for a variety of architectures.

Omega is an unstructured-mesh ocean model based on TRiSK numerical methods. It will be the new ocean component of the DOE’s Energy Exascale Earth System Model (E3SM). The algorithms in Omega follow those of the current ocean component, MPAS-Ocean, but it will be written in C++ rather than Fortran to take advantage of the Kokkos performance portability library. Omega spatial operators are written as Kokkos kernels to run efficiently on both CPUs and GPUs. Work on Omega began in 2023 with a new C++ framework for unstructured mesh partitioning, halo exchanges, parallel IO, and Kokkos interfaces. The current version, Omega-0, solves the shallow water equations and the tracer advection-diffusion equation. In this presentation we share our experiences in transitioning a large ocean code from supercomputers with traditional CPU architectures to exascale computers with heterogeneous CPU/GPU architectures.