GPU

FourierFlows allows you to easily construct and run problems on GPUs.

Upon calling

using FourierFlows

FourierFlows.jl will check whether any CUDA enabled device is present. If such a device is found then FourierFlows.jl makes sure that CUDA related packages are loaded and also it will overload all methods to work with GPU() device as their argument (instead of the standard CPU() device).

It's easy to construct a grid that lives on the GPU. Calling:

nx, Lx = 16, 2.0
grid = OneDGrid(GPU(); nx, Lx)

OneDimensionalGrid
  ├─────────── Device: GPU
  ├──────── FloatType: Float64
  ├────────── size Lx: 2.0
  ├──── resolution nx: 16
  ├── grid spacing dx: 0.125
  ├─────────── domain: x ∈ [-1.0, 0.875]
  └─ aliased fraction: 0.3333333333333333

returns a grid whose arrays are CuArrays. (Calling OneDGrid(; nx, Lx) defaults to using CPU(), i.e., OneDGrid(CPU(); nx, Lx).)

When we construct the Params, Vars, and Equation for our problem we need to make sure that we create arrays on the appropriate device, i.e., Arrays for CPU or CuArrays for the GPU. Function device_array is useful in constructing appropriately chosen arrays.

FourierFlows.device_array — Function

device_array(device::Device)
device_array(device::Device, T, dim)

Return the proper array type according to the device, i.e., Array for CPU and CuArray for GPU.

source

device_array(grid::AbstractGrid)

Return the proper array type according to the grid's device, i.e., Array for CPU and CuArray for GPU.

source

The FourierFlows.Problem constructor infers the device from the grid that is provided.

problem = Problem(equation, stepper, dt, grid, vars, params)

The FourierFlows.Diffusion module is written in a way such that switching from CPU to GPU is only a matter of calling FourierFlows.Diffusion.Problem() with dev=GPU(). All physics modules in GeophysicalFlows.jl can also seamlessly run on a GPU with dev=GPU() argument.

Selecting GPU device

FourierFlows.jl can only utilize a single GPU. If your machine has more than one GPU available, then functionality within CUDA.jl package enables the user to choose the GPU device that FourierFlows.jl should use. The user is referred to the CUDA.jl Documentation; in particular, CUDA.devices and CUDA.CuDevice.