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
— Functiondevice_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.
device_array(grid::AbstractGrid)
Return the proper array type according to the grid
's device
, i.e., Array
for CPU and CuArray
for GPU.
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
.