Authors:
(1) Simone Silvestri, Massachusetts Institute of Technology, Cambridge, MA, USA;
(2) Gregory Wagner, Massachusetts Institute of Technology, Cambridge, MA, USA;
(3) Christopher Hill, Massachusetts Institute of Technology, Cambridge, MA, USA;
(4) Matin Raayai Ardakani, Northeastern University, Boston, MA, USA;
(5) Johannes Blaschke, Lawrence Berkeley National Laboratory, Berkeley, CA, USA;
(6) Valentin Churavy, Massachusetts Institute of Technology, Cambridge, MA, USA;
(7) Jean-Michel Campin, Massachusetts Institute of Technology, Cambridge, MA, USA;
(8) Navid Constantinou, Australian National University, Canberra, ACT, Australia;
(9) Alan Edelman, Massachusetts Institute of Technology, Cambridge, MA, USA;
(10) John Marshall, Massachusetts Institute of Technology, Cambridge, MA, USA;
(11) Ali Ramadhan, Massachusetts Institute of Technology, Cambridge, MA, USA;
(12) Andre Souza, Massachusetts Institute of Technology, Cambridge, MA, USA;
(13) Raffaele Ferrari, Massachusetts Institute of Technology, Cambridge, MA, USA. Authors: Authors: (1) Simone Silvestri, Massachusetts Institute of Technology, Cambridge, MA, USA; (2) Gregory Wagner, Massachusetts Institute of Technology, Cambridge, MA, USA; (3) Christopher Hill, Massachusetts Institute of Technology, Cambridge, MA, USA; (4) Matin Raayai Ardakani, Northeastern University, Boston, MA, USA; (5) Johannes Blaschke, Lawrence Berkeley National Laboratory, Berkeley, CA, USA; (6) Valentin Churavy, Massachusetts Institute of Technology, Cambridge, MA, USA; (7) Jean-Michel Campin, Massachusetts Institute of Technology, Cambridge, MA, USA; (8) Navid Constantinou, Australian National University, Canberra, ACT, Australia; (9) Alan Edelman, Massachusetts Institute of Technology, Cambridge, MA, USA; (10) John Marshall, Massachusetts Institute of Technology, Cambridge, MA, USA; (11) Ali Ramadhan, Massachusetts Institute of Technology, Cambridge, MA, USA; (12) Andre Souza, Massachusetts Institute of Technology, Cambridge, MA, USA; (13) Raffaele Ferrari, Massachusetts Institute of Technology, Cambridge, MA, USA. Table of Links Abstract and 1 Justification Abstract and 1 Justification 2 Performance Attributes 2 Performance Attributes 3 Overview of the Problem 3 Overview of the Problem 4 Current State of the Art 4 Current State of the Art 5 Innovations 5 Innovations 5.1 Starting from scratch with Julia 5.1 Starting from scratch with Julia 5.2 New numerical methods for finite volume fluid dynamics on the sphere 5.2 New numerical methods for finite volume fluid dynamics on the sphere 5.3 Optimization of ocean free surface dynamics for unprecedented GPU scalability 5.3 Optimization of ocean free surface dynamics for unprecedented GPU scalability 6 How performance was measured 6 How performance was measured 7 Performance Results and 7.1 Scaling Results 7 Performance Results and 7.1 Scaling Results 7.2 Energy efficiency 7.2 Energy efficiency 8 Implications 8 Implications 9 Acknowledgments and References 9 Acknowledgments and References 5.1 Starting from scratch with Julia Oceananigans.jl is an open-source library for ocean-flavored fluid dynamics written from scratch in Julia [7]. Julia is a dynamic high-level programming language that leverages Just-In-Time (JIT) compilation and LLVM [24] to achieve performance competitive with traditional HPC languages like C or Fortran. Julia has gathered interest as potential language for HPC [17, 10, 16, 21, 27] and provides easy integration with MPI [47, 38]. Most of Oceananigans.jl software is hardware-agnostic through the Julia package KernelAbstractions.jl [10], which enables performance portability targeting CPUs and different GPU vendors using the JuliaGPU [6, 5] software stack, similar to the capabilities provided by Kokkos [9], OCCA [30], and HIP [1]. To our knowledge, Oceananigans is the first ocean model written from scratch for GPUs, rather than ported from existing CPU code. Starting from scratch and using the Julia programming language allowed us to rethink the typical patterns used in ocean and atmosphere dynamical cores. In particular, we developed a system of composable atomic operators that leverages Julia’s functional programming paradigm and effective inlining capabilities to recursively construct large expression trees for calculus on staggered finite volume grids. Using this composable operator system, we fuse the entire tendency computation for each prognostic variable into a single compute-heavy kernel, each of which depends on only two intermediate diagnostic variables representing hydrostatic pressure and vertical diffusivity (which is treated implicitly using a predictor-corrector method). Such a high degree of abstraction yields a number of innovations: first, kernel fusion maximizes efficiency on GPUs. Second, almost all intermediate quantities are computed on-the-fly, so that Oceananigans is extremely memory efficient and can perform global ocean simulations at resolutions up to 1/4th degree on a single Nvidia V100. Finally, because all compute-heavy kernels rely on a single “tendency kernel function” applied at each grid index i, j, k, we can easily optimize performance by rapidly prototyping techniques to overlap computation and communication. The sparsity of kernels per time-step and small number of temporary variables mean that Oceananigans’ algorithmic structure is markedly different from current ocean models, which typically allocate 10 to 100 times the minimum necessary memory [4] and distribute computations across many small kernels [51]. We argue these algorithmic differences are a major factor in Oceananigans’ energy-efficiency and time-to-solution on GPU systems. 5.2 New numerical methods for finite volume fluid dynamics on the sphere Our results use Oceananigans.HydrostaticFreeSurfaceModel, which solves the hydrostatic Boussinesq equations in a finite volume framework on staggered C-grids [3]. Oceananigans’ hydrostatic model employs an implicit-explicit second-order Adams-Bashforth time stepping scheme. Vertically implicit diffusion is implemented with a backward Euler time-discretization and tridiagonal solver. A major innovation is a new adaptive-order scheme based on weighted essentially non-oscillatory (WENO) reconstructions [42] for advecting momentum and tracers on curvilinear finite-volume grids [43]. This new scheme automatically adapts to changing spatial resolution and permits stable, high-fidelity simulations of ocean turbulence without explicit dissipation or hyper-dissipation. This innovation reduces setup time when changing or increasing resolution while guaranteeing high-fidelity solutions that exhibit the minimum necessary dissipation of sharp, near-grid scale features. 5.3 Optimization of ocean free surface dynamics for unprecedented GPU scalability In hydrostatic ocean models with a free surface, the vertically-averaged, two-dimensional “barotropic mode” has dynamics orders of magnitude faster than the three-dimensional “baroclinic” component, and must be treated by a special “barotropic solver”. Due to communication overhead, barotropic solvers in current ocean models — whether implicit or explicit — are a major bottleneck that accounts for between 40% [22] to 60% [48, 36] of the cost of a typical IPCC-class ocean simulations. Oceananigans’ excellent scalability is enabled by an innovative optimization of the parallel barotropic solver. An increase in computation is traded in for decreased communication latency by leveraging the two-dimensionality of the barotropic problem. Our new barotropic solver is based on explicit subcycling of the barotropic mode. Increasing the width of the barotropic halo to equal the number of explicit subcycles (typically between 10–30) greatly decreases the frequency of communication. As a result, communication is required once per time-step rather than every subcycle, reducing the frequency of communication by a factor of 10 to 30. The cost of the barotropic solver is therefore less than 10% of the total cost of a time step. Due to the sparsity of communication enabled by our novel barotropic solver, all communication operations can be overlapped with computational workloads as sketched in figure 2. This paper is available on arxiv under CC BY 4.0 DEED license. This paper is available on arxiv under CC BY 4.0 DEED license. available on arxiv

Part of HackerNoon's growing list of open-source research papers, promoting free access to academic material.

From Code to Currents: Julia and GPU Magic Transform Ocean Modeling Efficiency

About Author

Comments

TOPICS

THIS ARTICLE WAS FEATURED IN

Related Stories

Untitled Story

A Deep Discussion On Our Study Looking at Near-Inertial Waves and Mesoscale Eddies: Diving In

How Oceananigans.jl Makes High-Resolution Climate Simulations Affordable

From 488m to 10km: Oceananigans Sets New Standards in Climate Science

Scaling the Depths: Oceananigans Achieves Record-Breaking Climate Simulation Performance

From Submesoscales to Global Impact: Oceananigans Powers Climate Predictions

A Deep Discussion On Our Study Looking at Near-Inertial Waves and Mesoscale Eddies: Diving In

How Oceananigans.jl Makes High-Resolution Climate Simulations Affordable

From 488m to 10km: Oceananigans Sets New Standards in Climate Science

Scaling the Depths: Oceananigans Achieves Record-Breaking Climate Simulation Performance

From Submesoscales to Global Impact: Oceananigans Powers Climate Predictions

Light-Mode

Classic

Newspaper

Dark-Mode

Neon Noir

Minty

HN StartUps