Accessing GPUs and using Oceananigans on GPUs

GP-what?

One of Oceananigans' greatest powers is its ability to run on graphical processing units (GPUs). It's not uncommon to obtain speed-ups between 100x and 400x on a research-grade NVidia GPU such as a Titan V or V100 over a single CPU. These kind of speed-ups are most likely for "large-ish" problems with more than 1,000,000 grid points, or resolutions of roughly 128x128x64 and greater. We also often see greater speed-ups as the "complexity" of the model tendency calculations increases --- which occurs when using higher-order advection schemes, nonlinear equations of state, immersed boundaries, and many turbulence closures at the same time. For many pricing models on clusters or Google Cloud, all this means that GPU simulations end up being cheaper and consume less energy than CPU simulations, and that makes everyone happy.

See the performance benchmarks for more information about Oceananigans performance on CPUs and GPUs.

Using GPUs

The biggest obstacle to using a GPU is finding one and installing CUDA. Once that's done, all that's left is to build our grid on the GPU; for example:

using Oceananigans
grid = RectilinearGrid(GPU(), size=(128, 128, 64), x=(0, 1), y=(0, 1), z=(0, 1))

For more details about using Oceananigans on GPUs, see the documentation.

Fast GPUs and slow GPUs

At the moment only Nvidia GPUs are supported (but we hope that AMD GPUs will be supported in the future). Generally speaking, the GPUs shipped with most laptops (for example) do not outperform their associated CPUs for Oceananigans calculations. Speedups are usually obtained only with "research-grade" Nvidia GPUs like the K80, P100, Titan V, V100, or the A100 (for the fortunate souls with access to an A100). In addition, some users have reported decent performance for high-end gaming-oriented GPUs from the GeForce or Quadro line of Nvidia gaming GPUs, especially when running single-precision models.

Note: if your experience differs from the guidelines above, or if you have more detailed information about the performance of Oceananigans on certain Nvidia GPUs, don't hesitate to update this wiki or let us know in the Discussions!

Accessing GPUs

Users in the academic community have a few options for obtaining access to Nvidia GPUs:

• Drop $15k and buy an A100 for yourself. But actually, the Titan V is a great deal too.
• Check whether high-performance computing resources available in your region include GPU capability. For example, if you're based in the United States, you can apply for time on XSEDE computers to obtain access to GPUs.
• Purchase GPU time through Google Cloud or Amazon EC2.
• See the Julia on Google Colab: Free GPU-Accelerated Shareable Notebooks post on the Julia Discourse.
• Code Ocean's GPU support provides access to Tesla K80 GPUs.