Heterogeneous computing

Support for heterogeneous computing is currently being worked on.

The 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. Put the inner code in a new function rhs_fct_per_element. 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