Heterogeneous computing

Support for heterogeneous computing is currently being worked on.

User-facing interface

GPU support in Trixi.jl is controlled via two keyword arguments to semidiscretize:

• storage_type: the array type used for all internal data structures. Set to CuArray (from CUDA.jl) for NVIDIA GPUs or ROCArray (from AMDGPU.jl) for AMD GPUs. Defaults to nothing, which keeps the standard CPU Array.
• real_type: the floating-point element type. GPU workflows typically benefit from setting this to Float32. Defaults to nothing, which retains the type used when building the semidiscretization (usually Float64).

Both arguments can be used independently. A typical GPU setup looks like:

using CUDA   # or using AMDGPU
ode = semidiscretize(semi, tspan; real_type = Float32, storage_type = CuArray)

The rest of the elixir (callbacks, ODE solver call) remains unchanged. See, e.g., examples/p4est_2d_dgsem/elixir_euler_source_terms.jl for a concrete example.

Internal use of Adapt.jl

Adapt.jl is a package in the JuliaGPU family that allows for the translation of nested data structures. The primary goal is to allow the substitution of Array at the storage level with a GPU array like CuArray from CUDA.jl.

To facilitate this, data structures must be parameterized, so instead of:

struct Container <: Trixi.AbstractContainer
   data::Array{Float64, 2}
end

They must be written as:

struct Container{D<:AbstractArray} <: Trixi.AbstractContainer
   data::D
end

furthermore, we need to define a function that allows for the conversion of storage of our types:

using Adapt

function Adapt.adapt_structure(to, C::Container)
    return Container(adapt(to, C.data))
end

or simply

Adapt.@adapt_structure(Container)

additionally, we must define Adapt.parent_type.

function Adapt.parent_type(::Type{<:Container{D}}) where D
    return D
end

All together we can use this machinery to perform conversions of a container.

julia> C = Container(zeros(3))
Container{Vector{Float64}}([0.0, 0.0, 0.0])

julia> Trixi.storage_type(C)
Array

julia> using CUDA

julia> GPU_C = adapt(CuArray, C)
Container{CuArray{Float64, 1, CUDA.DeviceMemory}}([0.0, 0.0, 0.0])

julia> Trixi.storage_type(C)
CuArray

Element-type conversion with Trixi.trixi_adapt.

We can use Trixi.trixi_adapt to perform both an element-type and a storage-type adoption:

julia> C = Container(zeros(3))
Container{Vector{Float64}}([0.0, 0.0, 0.0])

julia> Trixi.trixi_adapt(Array, Float32, C)
Container{Vector{Float32}}(Float32[0.0, 0.0, 0.0])

julia> Trixi.trixi_adapt(CuArray, Float32, C)
Container{CuArray{Float32, 1, CUDA.DeviceMemory}}(Float32[0.0, 0.0, 0.0])

Writing GPU kernels

Offloading computations to the GPU is done with KernelAbstractions.jl, allowing for vendor-agnostic GPU code.

Example

Given the following Trixi.jl code, which would typically be called from within rhs!:

function trixi_rhs_fct(mesh, equations, solver, cache, args)
    @threaded for element in eachelement(solver, cache)
        # code
    end
end

1. Move the inner code into a new inlined function rhs_fct_per_element.

   @inline function rhs_fct_per_element(..., element, ...)
       ...
   end

   Besides the index element, pass all required fields as arguments, but make sure to @unpack them from their structs in advance.

2. Where trixi_rhs_fct is called, get the backend, i.e., the hardware we are currently running on via trixi_backend(x). This will, e.g., work with u_ode. Internally, KernelAbstractions.jl's get_backend will be called, i.e., KernelAbstractions.jl has to know the type of x.

   backend = trixi_backend(u_ode)

3. Add a new argument backend to trixi_rhs_fct used for dispatch. When backend is nothing, the legacy implementation should be used:

   function trixi_rhs_fct(backend::Nothing, mesh, equations, solver, cache, args)
       @unpack unpacked_args = cache
       @threaded for element in eachelement(solver, cache)
           rhs_fct_per_element(element, unpacked_args, args)
       end
   end

4. When backend is a Backend (a type defined by KernelAbstractions.jl), write a KernelAbstractions.jl kernel:

   function trixi_rhs_fct(backend::Backend, mesh, equations, solver, cache, args)
       nelements(solver, cache) == 0 && return nothing  # return early when there are no elements
       @unpack unpacked_args = cache
       kernel! = rhs_fct_kernel!(backend)
       kernel!(unpacked_args, args,
               ndrange = nelements(solver, cache))
       return nothing
   end
   
   @kernel function rhs_fct_kernel!(unpacked_args, args)
       element = @index(Global)
       rhs_fct_per_element(element, unpacked_args, args)
   end