Use Adapt.jl to change storage and element type

In order to eventually support GPU computation we need
to use Adapt.jl to allow GPU backend packages to swap
out host-array types like `CuArray` with device-side types
like `CuDeviceArray`.

Additionally this will allow us to change the element type
of a simulation by using `adapt(Array{Float32}`.

Co-authored-by: Lars Christmann <[email protected]>
Co-authored-by: Benedict Geihe <[email protected]>

Author: 3 people

Project.toml

@@ -4,6 +4,7 @@ authors = ["Michael Schlottke-Lakemper <[email protected]>", "Gregor
 version = "0.9.15-DEV"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 CodeTracking = "da1fd8a2-8d9e-5ec2-8556-3022fb5608a2"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
@@ -64,6 +65,7 @@ TrixiMakieExt = "Makie"
 TrixiNLsolveExt = "NLsolve"
 
 [compat]
+Adapt = "3.7, 4.0"
 Accessors = "0.1.12"
 CodeTracking = "1.0.5"
 ConstructionBase = "1.3"

src/Trixi.jl

@@ -43,6 +43,7 @@ import SciMLBase: get_du, get_tmp_cache, u_modified!,
 
 using DelimitedFiles: readdlm
 using Downloads: Downloads
+import Adapt
 using CodeTracking: CodeTracking
 using ConstructionBase: ConstructionBase
 using DiffEqCallbacks: PeriodicCallback, PeriodicCallbackAffect

src/auxiliary/vector_of_arrays.jl

@@ -0,0 +1,22 @@
+# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
+# Since these FMAs can increase the performance of many numerical algorithms,
+# we need to opt-in explicitly.
+# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
+@muladd begin
+#! format: noindent
+
+# Wraps a Vector of Arrays, forwards `getindex` to the underlying Vector.
+# Implements `Adapt.adapt_structure` to allow offloading to the GPU which is
+# not possible for a plain Vector of Arrays.
+struct VecOfArrays{T <: AbstractArray}
+    arrays::Vector{T}
+end
+Base.getindex(v::VecOfArrays, i::Int) = Base.getindex(v.arrays, i)
+Base.IndexStyle(v::VecOfArrays) = Base.IndexStyle(v.arrays)
+Base.size(v::VecOfArrays) = Base.size(v.arrays)
+Base.length(v::VecOfArrays) = Base.length(v.arrays)
+Base.eltype(v::VecOfArrays{T}) where {T} = T
+function Adapt.adapt_structure(to, v::VecOfArrays)
+    return [Adapt.adapt(to, arr) for arr in v.arrays] |> VecOfArrays
+end
+end # @muladd

src/semidiscretization/semidiscretization_hyperbolic.jl

@@ -30,19 +30,19 @@ mutable struct SemidiscretizationHyperbolic{Mesh, Equations, InitialCondition,
 
     function SemidiscretizationHyperbolic{Mesh, Equations, InitialCondition,
                                           BoundaryConditions, SourceTerms, Solver,
-                                          Cache}(mesh::Mesh, equations::Equations,
+                                          Cache}(mesh::Mesh,
+                                                 equations::Equations,
                                                  initial_condition::InitialCondition,
                                                  boundary_conditions::BoundaryConditions,
                                                  source_terms::SourceTerms,
                                                  solver::Solver,
-                                                 cache::Cache) where {Mesh, Equations,
+                                                 cache::Cache,
+                                                 performance_counter::PerformanceCounter) where {Mesh, Equations,
                                                                       InitialCondition,
                                                                       BoundaryConditions,
                                                                       SourceTerms,
                                                                       Solver,
                                                                       Cache}
-        performance_counter = PerformanceCounter()
-
         new(mesh, equations, initial_condition, boundary_conditions, source_terms,
             solver, cache, performance_counter)
     end
@@ -74,14 +74,16 @@ function SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver
 
     check_periodicity_mesh_boundary_conditions(mesh, _boundary_conditions)
 
+    performance_counter = PerformanceCounter()
+
     SemidiscretizationHyperbolic{typeof(mesh), typeof(equations),
                                  typeof(initial_condition),
                                  typeof(_boundary_conditions), typeof(source_terms),
                                  typeof(solver), typeof(cache)}(mesh, equations,
                                                                 initial_condition,
                                                                 _boundary_conditions,
                                                                 source_terms, solver,
-                                                                cache)
+                                                                cache, performance_counter)
 end
 
 # Create a new semidiscretization but change some parameters compared to the input.
@@ -103,6 +105,30 @@ function remake(semi::SemidiscretizationHyperbolic; uEltype = real(semi.solver),
                                  source_terms, boundary_conditions, uEltype)
 end
 
+function Adapt.adapt_structure(to, semi::SemidiscretizationHyperbolic)
+    if !(typeof(semi.mesh) <: P4estMesh)
+        error("Adapt.adapt is only supported for semidiscretizations based on P4estMesh")
+    end
+
+    mesh = semi.mesh
+    equations = Adapt.adapt_structure(to, semi.equations)
+    initial_condition = Adapt.adapt_structure(to, semi.initial_condition)
+    boundary_conditions = Adapt.adapt_structure(to, semi.boundary_conditions)
+    source_terms = Adapt.adapt_structure(to, semi.source_terms)
+    solver = Adapt.adapt_structure(to, semi.solver)
+    cache = Adapt.adapt_structure(to, semi.cache)
+    performance_counter = semi.performance_counter
+
+    SemidiscretizationHyperbolic{typeof(mesh), typeof(equations),
+                                 typeof(initial_condition),
+                                 typeof(boundary_conditions), typeof(source_terms),
+                                 typeof(solver), typeof(cache)}(mesh, equations,
+                                                                initial_condition,
+                                                                boundary_conditions,
+                                                                source_terms, solver,
+                                                                cache, performance_counter)
+end
+
 # general fallback
 function digest_boundary_conditions(boundary_conditions, mesh, solver, cache)
     boundary_conditions

src/solvers/dgsem/basis_lobatto_legendre.jl

@@ -122,6 +122,31 @@ In particular, not the nodes themselves are returned.
 
 @inline get_nodes(basis::LobattoLegendreBasis) = basis.nodes
 
+function Adapt.adapt_structure(to, basis::LobattoLegendreBasis)
+    # Do not adapt SVector fields, i.e. nodes, weights and inverse_weights
+    (; nodes, weights, inverse_weights) = basis
+    inverse_vandermonde_legendre = Adapt.adapt_structure(to,
+                                                         basis.inverse_vandermonde_legendre)
+    boundary_interpolation = basis.boundary_interpolation
+    derivative_matrix = Adapt.adapt_structure(to, basis.derivative_matrix)
+    derivative_split = Adapt.adapt_structure(to, basis.derivative_split)
+    derivative_split_transpose = Adapt.adapt_structure(to,
+                                                       basis.derivative_split_transpose)
+    derivative_dhat = Adapt.adapt_structure(to, basis.derivative_dhat)
+    return LobattoLegendreBasis{real(basis), nnodes(basis), typeof(basis.nodes),
+                                typeof(inverse_vandermonde_legendre),
+                                typeof(boundary_interpolation),
+                                typeof(derivative_matrix)}(nodes,
+                                                           weights,
+                                                           inverse_weights,
+                                                           inverse_vandermonde_legendre,
+                                                           boundary_interpolation,
+                                                           derivative_matrix,
+                                                           derivative_split,
+                                                           derivative_split_transpose,
+                                                           derivative_dhat)
+end
+
 """
     integrate(f, u, basis::LobattoLegendreBasis)
 
@@ -209,6 +234,16 @@ end
 
 @inline polydeg(mortar::LobattoLegendreMortarL2) = nnodes(mortar) - 1
 
+function Adapt.adapt_structure(to, mortar::LobattoLegendreMortarL2)
+    forward_upper = Adapt.adapt_structure(to, mortar.forward_upper)
+    forward_lower = Adapt.adapt_structure(to, mortar.forward_lower)
+    reverse_upper = Adapt.adapt_structure(to, mortar.reverse_upper)
+    reverse_lower = Adapt.adapt_structure(to, mortar.reverse_lower)
+    return LobattoLegendreMortarL2{real(mortar), nnodes(mortar), typeof(forward_upper),
+                                   typeof(reverse_upper)}(forward_upper, forward_lower,
+                                                          reverse_upper, reverse_lower)
+end
+
 # TODO: We can create EC mortars along the lines of the following implementation.
 # abstract type AbstractMortarEC{RealT} <: AbstractMortar{RealT} end