Update on OMEGA
Key Points:
- MPAS-Ocean is being rewritten using C++/Kokkos to improve exascale performance.
- The Ocean Model for E3SM Global Applications (OMEGA) will be included in E3SMv4 and enable long, high-resolution simulations.
- Initial performance results for a preliminary OMEGA version shows it is 1.5x faster than MPAS-Ocean on CPU and 40x faster than MPAS-Ocean on GPU
To enable unprecedented ocean model resolutions in E3SM for predicting changes to energy availability and resilience, the Ocean Model for E3SM Global Applications (OMEGA) is under development for E3SMv4. The primary objective of OMEGA is to achieve performance portability on heterogeneous architectures, such as those found on DOE’s exascale computing machines Frontier and Aurora. Leveraging the project’s experience with creating the Gordon Bell prize winning EAMxx model (the C++ next generation E3SM Atmosphere Model), OMEGA is being developed in C++ using the performance portability library, Kokkos, to enable a single code implementation to achieve high performance across architectures and hardware vendors. The Kokkos approach eliminates the need to develop and maintain different code branches that utilize different GPU hardware-specific back-end languages, while still supporting CPU compatibility. Although OMEGA is primarily a C++/Kokkos re-write of the MPAS-Ocean physics and numerics, it also presents an opportunity to improve the performance of the underlying framework that limited the performance of the MPAS-Ocean model, especially with respect to GPU utilization.
Figure 1. Performance of OMEGA-0, a preliminary model configuration, on CPU (a) and GPU (c) and MPAS-Ocean on CPU (b) and GPU (d). Note that the GPU performance of OMEGA is so improved over MPAS-Ocean that the vertical scales in (c) and (d) are different. The OMEGA total performance line would be below the green line in (d)
The OMEGA group has recently reached the first milestone in Phase 3 of E3SM, which is a model that solves the stacked shallow water equations with passive tracers (OMEGA-0). All the horizontal dynamical terms are included in the model, but for now, the layers in the vertical are completely independent. OMEGA-0 has allowed the team to perform verification on the correctness of the code and make early assessments of the model’s performance at scale. Early results show that OMEGA-0 is approximately 1.5 times faster than MPAS-Ocean on CPU (compare panels (a) and (b) in Fig. 1), using a similar stacked shallow water configuration. The GPU performance of OMEGA-0 has been particularly promising, with the model showing a speed-up of 30x on Frontier when comparing full-node CPU to full-node GPU performance (compare panels (a) and (c) in Fig. 1). MPAS-Ocean is 30-50% faster on GPU vs. CPU (compare panels (b) and (d) in Fig. 1). These results are indications that the C++/Kokkos re-write strategy will result in a substantially more performant ocean model in future E3SM versions. This is a great improvement from previous attempts to utilize GPUs with the MPAS-Ocean model and openACC that produced only modest performance gains. The re-write strategy has yielded a 40x improvement over MPAS-Ocean on GPU.
Currently, the OMEGA team is working toward developing OMEGA-1, which will be a 3D ocean model, with some simplified physics and idealized forcing. This model requires the addition of vertical advection, an equation of state, baroclinic pressure gradient, and vertical mixing. In addition, OMEGA will differ from MPAS-Ocean in that it is being developed as a mass conserving, non-Boussinesq model. This will allow the model to naturally represent steric sea level change and improve the representation of mass exchanges in the coupled system. OMEGA will also utilize the latest oceanic equation of state (TEOS-10; Thermodynamic Equation of SeaWater 2010 (TEOS-10) ).
Throughout the OMEGA development cycle, a heavy emphasis has been placed on testing, verification and documentation. Each pull request that introduces a new capability to the code has a design document which is reviewed prior to development, unit tests, verification tests, and documentation for a users and developers guide. Low-level unit tests are performed using CTest, while more complex verification testing, such as convergence testing, is being implemented in a python package called Polaris.
Early indications of the performance of OMEGA-0 show that the final OMEGA model will provide E3SM with the ability to efficiently deliver ultra-high resolution simulations on DOE’s world-leading computing systems. These unprecedented levels of resolution will allow E3SM to be a critical tool for predicting the impacts of subseasonal to decadal variability on energy production, demand, distribution, and infrastructure.
This article is a part of the E3SM “Floating Points” Newsletter, to read the full Newsletter check:
