start working on test and use custom adaptor

vchuravy · vchuravy · commit a3e4530c8260 · 2025-04-20T21:59:52.000+02:00
diff --git a/src/Trixi.jl b/src/Trixi.jl
@@ -44,7 +44,7 @@ import SciMLBase: get_du, get_tmp_cache, u_modified!,
 
 using DelimitedFiles: readdlm
 using Downloads: Downloads
-using Adapt: Adapt
+using Adapt: Adapt, adapt
 using CodeTracking: CodeTracking
 using ConstructionBase: ConstructionBase
 using DiffEqBase: DiffEqBase, get_tstops, get_tstops_array
diff --git a/src/auxiliary/containers.jl b/src/auxiliary/containers.jl
@@ -346,4 +346,34 @@ function unsafe_wrap_or_alloc(to, vec, size)
         return unsafe_wrap(to, pointer(vec), size)
     end
 end
+
+struct TrixiAdaptor{Storage, Real} end
+
+function trixi_adapt(storage, real, x)
+    adapt(TrixiAdaptor{storage, real}(), x)
+end
+
+# Custom rules
+# 1. handling of StaticArrays
+function Adapt.adapt_storage(::TrixiAdaptor{<:Any, Real}, x::StaticArrays.StaticArray{S, T, N}) where {Real,S,T,N}
+    StaticArrays.similar_type(x, Real)(x)
+end
+
+# 2. Handling of Arrays
+function Adapt.adapt_storage(::TrixiAdaptor{Storage, Real}, x::AbstractArray{T}) where{Storage, Real, T<:AbstractFloat}
+    adapt(Storage{Real}, x)
+end
+
+function Adapt.adapt_storage(::TrixiAdaptor{Storage, Real}, x::AbstractArray{T}) where {Storage, Real,T<:StaticArrays.StaticArray}
+    adapt(Storage{StaticArrays.similar_type(T, Real)},x)
+end
+
+function Adapt.adapt_storage(::TrixiAdaptor{Storage, Real}, x::AbstractArray) where{Storage, Real}
+    adapt(Storage, x)
+end
+
+# 3. TODO: Should we have a fallback? But that would imply implementing things for NamedTuple again
+
+unsafe_wrap_or_alloc(::TrixiAdaptor{Storage}, vec, size) where {Storage} = unsafe_wrap_or_alloc(Storage, vec, size)
+
 end # @muladd
diff --git a/src/semidiscretization/semidiscretization_hyperbolic.jl b/src/semidiscretization/semidiscretization_hyperbolic.jl
@@ -27,28 +27,6 @@ mutable struct SemidiscretizationHyperbolic{Mesh, Equations, InitialCondition,
     solver::Solver
     cache::Cache
     performance_counter::PerformanceCounter
-
-    function SemidiscretizationHyperbolic{Mesh, Equations, InitialCondition,
-                                          BoundaryConditions, SourceTerms, Solver,
-                                          Cache}(mesh::Mesh,
-                                                 equations::Equations,
-                                                 initial_condition::InitialCondition,
-                                                 boundary_conditions::BoundaryConditions,
-                                                 source_terms::SourceTerms,
-                                                 solver::Solver,
-                                                 cache::Cache,
-                                                 performance_counter::PerformanceCounter) where {
-                                                                                                 Mesh,
-                                                                                                 Equations,
-                                                                                                 InitialCondition,
-                                                                                                 BoundaryConditions,
-                                                                                                 SourceTerms,
-                                                                                                 Solver,
-                                                                                                 Cache
-                                                                                                 }
-        new(mesh, equations, initial_condition, boundary_conditions, source_terms,
-            solver, cache, performance_counter)
-    end
 end
 
 """
@@ -90,6 +68,31 @@ function SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver
                                                                 performance_counter)
 end
 
+@eval Adapt.@adapt_structure(SemidiscretizationHyperbolic)
+# 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(to, semi.equations)
+#     initial_condition = adapt(to, semi.initial_condition)
+#     boundary_conditions = adapt(to, semi.boundary_conditions)
+#     source_terms = adapt(to, semi.source_terms)
+#     solver = adapt(to, semi.solver)
+#     cache = adapt(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
+
 # Create a new semidiscretization but change some parameters compared to the input.
 # `Base.similar` follows a related concept but would require us to `copy` the `mesh`,
 # which would impact the performance. Instead, `SciMLBase.remake` has exactly the
@@ -109,9 +112,6 @@ function remake(semi::SemidiscretizationHyperbolic; uEltype = real(semi.solver),
                                  source_terms, boundary_conditions, uEltype)
 end
 
-# @eval due to @muladd
-@eval Adapt.@adapt_structure(SemidiscretizationHyperbolic)
-
 # general fallback
 function digest_boundary_conditions(boundary_conditions, mesh, solver, cache)
     boundary_conditions
diff --git a/src/solvers/dg.jl b/src/solvers/dg.jl
@@ -400,6 +400,8 @@ struct DG{Basis, Mortar, SurfaceIntegral, VolumeIntegral}
     volume_integral::VolumeIntegral
 end
 
+@eval Adapt.@adapt_structure(DG)
+
 function Base.show(io::IO, dg::DG)
     @nospecialize dg # reduce precompilation time
 
diff --git a/src/solvers/dgsem/basis_lobatto_legendre.jl b/src/solvers/dgsem/basis_lobatto_legendre.jl
@@ -34,6 +34,32 @@ struct LobattoLegendreBasis{RealT <: Real, NNODES,
     # negative adjoint wrt the SBP dot product
 end
 
+function Adapt.adapt_structure(to, basis::LobattoLegendreBasis)
+    inverse_vandermonde_legendre = adapt(to, basis.inverse_vandermonde_legendre)
+    RealT = eltype(inverse_vandermonde_legendre)
+
+    nodes = SVector{<:Any, RealT}(basis.nodes)
+    weights = SVector{<:Any, RealT}(basis.weights)
+    inverse_weights = SVector{<:Any, RealT}(basis.inverse_weights)
+    boundary_interpolation = adapt(to, basis.boundary_interpolation)
+    derivative_matrix = adapt(to, basis.derivative_matrix)
+    derivative_split = adapt(to, basis.derivative_split)
+    derivative_split_transpose = adapt(to,basis.derivative_split_transpose)
+    derivative_dhat = adapt(to, basis.derivative_dhat)
+    return LobattoLegendreBasis{RealT, nnodes(basis), typeof(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
+
 function LobattoLegendreBasis(RealT, polydeg::Integer)
     nnodes_ = polydeg + 1
 
@@ -122,9 +148,6 @@ In particular, not the nodes themselves are returned.
 
 @inline get_nodes(basis::LobattoLegendreBasis) = basis.nodes
 
-# @eval due to @muladd
-@eval Adapt.@adapt_structure(LobattoLegendreBasis)
-
 """
     integrate(f, u, basis::LobattoLegendreBasis)
 
@@ -158,6 +181,16 @@ struct LobattoLegendreMortarL2{RealT <: Real, NNODES,
     reverse_lower::ReverseMatrix
 end
 
+function Adapt.adapt_structure(to, mortar::LobattoLegendreMortarL2)
+    forward_upper = adapt(to, mortar.forward_upper)
+    forward_lower = adapt(to, mortar.forward_lower)
+    reverse_upper = adapt(to, mortar.reverse_upper)
+    reverse_lower = adapt(to, mortar.reverse_lower)
+    return LobattoLegendreMortarL2{eltype(forward_upper), nnodes(mortar), typeof(forward_upper),
+                                   typeof(reverse_upper)}(forward_upper, forward_lower,
+                                                          reverse_upper, reverse_lower)
+end
+
 function MortarL2(basis::LobattoLegendreBasis)
     RealT = real(basis)
     nnodes_ = nnodes(basis)
@@ -208,9 +241,6 @@ end
 
 @inline polydeg(mortar::LobattoLegendreMortarL2) = nnodes(mortar) - 1
 
-# @eval due to @muladd
-@eval Adapt.@adapt_structure(LobattoLegendreMortarL2)
-
 # TODO: We can create EC mortars along the lines of the following implementation.
 # abstract type AbstractMortarEC{RealT} <: AbstractMortar{RealT} end
 
diff --git a/src/solvers/dgsem_p4est/containers.jl b/src/solvers/dgsem_p4est/containers.jl
@@ -145,17 +145,16 @@ end
 
 # Manual adapt_structure since we have aliasing memory
 function Adapt.adapt_structure(to,
-                               elements::P4estElementContainer{NDIMS, RealT, uEltype}) where {
-                                                                                              NDIMS,
-                                                                                              RealT,
-                                                                                              uEltype
-                                                                                              }
+                               elements::P4estElementContainer{NDIMS}) where {NDIMS}
     # Adapt underlying storage
-    _node_coordinates = Adapt.adapt_structure(to, elements._node_coordinates)
-    _jacobian_matrix = Adapt.adapt_structure(to, elements._jacobian_matrix)
-    _contravariant_vectors = Adapt.adapt_structure(to, elements._contravariant_vectors)
-    _inverse_jacobian = Adapt.adapt_structure(to, elements._inverse_jacobian)
-    _surface_flux_values = Adapt.adapt_structure(to, elements._surface_flux_values)
+    _node_coordinates = adapt(to, elements._node_coordinates)
+    _jacobian_matrix = adapt(to, elements._jacobian_matrix)
+    _contravariant_vectors = adapt(to, elements._contravariant_vectors)
+    _inverse_jacobian = adapt(to, elements._inverse_jacobian)
+    _surface_flux_values = adapt(to, elements._surface_flux_values)
+
+    RealT = eltype(_inverse_jacobian)
+    uEltype = eltype(_surface_flux_values)
 
     # Wrap arrays again
     node_coordinates = unsafe_wrap_or_alloc(to, _node_coordinates,
@@ -180,8 +179,7 @@ function Adapt.adapt_structure(to,
                        typeof(jacobian_matrix), # ArrayNDIMSP3
                        typeof(_node_coordinates), # VectorRealT
                        typeof(_surface_flux_values), # VectoruEltype
-                       to,
-                       true)
+                       to)
     return P4estElementContainer{new_type_params...}(node_coordinates,
                                                      jacobian_matrix,
                                                      contravariant_vectors,
@@ -296,9 +294,9 @@ end
 # Manual adapt_structure since we have aliasing memory
 function Adapt.adapt_structure(to, interfaces::P4estInterfaceContainer)
     # Adapt underlying storage
-    _u = Adapt.adapt_structure(to, interfaces._u)
-    _neighbor_ids = Adapt.adapt_structure(to, interfaces._neighbor_ids)
-    _node_indices = Adapt.adapt_structure(to, interfaces._node_indices)
+    _u = adapt(to, interfaces._u)
+    _neighbor_ids = adapt(to, interfaces._neighbor_ids)
+    _node_indices = adapt(to, interfaces._node_indices)
     # Wrap arrays again
     u = unsafe_wrap_or_alloc(to, _u, size(interfaces.u))
     neighbor_ids = unsafe_wrap_or_alloc(to, _neighbor_ids,
@@ -308,12 +306,11 @@ function Adapt.adapt_structure(to, interfaces::P4estInterfaceContainer)
 
     NDIMS = ndims(interfaces)
     new_type_params = (NDIMS,
-                       eltype(interfaces),
+                       eltype(_u),
                        NDIMS + 2,
                        typeof(u), typeof(neighbor_ids), typeof(node_indices),
                        typeof(_u), typeof(_neighbor_ids), typeof(_node_indices),
-                       to,
-                       true)
+                       to)
     return P4estInterfaceContainer{new_type_params...}(u, neighbor_ids, node_indices,
                                                        _u, _neighbor_ids, _node_indices)
 end
@@ -439,16 +436,16 @@ end
 
 # Manual adapt_structure since we have aliasing memory
 function Adapt.adapt_structure(to, boundaries::P4estBoundaryContainer)
-    _u = Adapt.adapt_structure(to, boundaries._u)
+    _u = adapt(to, boundaries._u)
     u = unsafe_wrap_or_alloc(to, _u, size(boundaries.u))
-    neighbor_ids = Adapt.adapt_structure(to, boundaries.neighbor_ids)
-    node_indices = Adapt.adapt_structure(to, boundaries.node_indices)
+    neighbor_ids = adapt(to, boundaries.neighbor_ids)
+    node_indices = adapt(to, boundaries.node_indices)
     name = boundaries.name
 
     NDIMS = ndims(boundaries)
-    return P4estBoundaryContainer{NDIMS, eltype(boundaries), NDIMS + 1, typeof(u),
+    return P4estBoundaryContainer{NDIMS, eltype(_u), NDIMS + 1, typeof(u),
                                   typeof(neighbor_ids), typeof(node_indices),
-                                  typeof(_u), to, true}(u, neighbor_ids, node_indices,
+                                  typeof(_u), to}(u, neighbor_ids, node_indices,
                                                         name, _u)
 end
 
@@ -578,9 +575,9 @@ end
 # Manual adapt_structure since we have aliasing memory
 function Adapt.adapt_structure(to, mortars::P4estMortarContainer)
     # Adapt underlying storage
-    _u = Adapt.adapt_structure(to, mortars._u)
-    _neighbor_ids = Adapt.adapt_structure(to, mortars._neighbor_ids)
-    _node_indices = Adapt.adapt_structure(to, mortars._node_indices)
+    _u = adapt(to, mortars._u)
+    _neighbor_ids = adapt(to, mortars._neighbor_ids)
+    _node_indices = adapt(to, mortars._node_indices)
 
     # Wrap arrays again
     u = unsafe_wrap_or_alloc(to, _u, size(mortars.u))
@@ -589,13 +586,12 @@ function Adapt.adapt_structure(to, mortars::P4estMortarContainer)
 
     NDIMS = ndims(mortars)
     new_type_params = (NDIMS,
-                       eltype(mortars),
+                       eltype(_u),
                        NDIMS + 1,
                        NDIMS + 3,
                        typeof(u), typeof(neighbor_ids), typeof(node_indices),
                        typeof(_u), typeof(_neighbor_ids), typeof(_node_indices),
-                       to,
-                       true)
+                       to)
     return P4estMortarContainer{new_type_params...}(u, neighbor_ids, node_indices,
                                                     _u, _neighbor_ids, _node_indices)
 end
diff --git a/src/solvers/dgsem_p4est/containers_parallel.jl b/src/solvers/dgsem_p4est/containers_parallel.jl
@@ -93,10 +93,10 @@ end
 # Manual adapt_structure since we have aliasing memory
 function Adapt.adapt_structure(to, mpi_interfaces::P4estMPIInterfaceContainer)
     # Adapt Vectors and underlying storage
-    _u = Adapt.adapt_structure(to, mpi_interfaces._u)
-    local_neighbor_ids = Adapt.adapt_structure(to, mpi_interfaces.local_neighbor_ids)
-    node_indices = Adapt.adapt_structure(to, mpi_interfaces.node_indices)
-    local_sides = Adapt.adapt_structure(to, mpi_interfaces.local_sides)
+    _u = adapt(to, mpi_interfaces._u)
+    local_neighbor_ids = adapt(to, mpi_interfaces.local_neighbor_ids)
+    node_indices = adapt(to, mpi_interfaces.node_indices)
+    local_sides = adapt(to, mpi_interfaces.local_sides)
 
     # Wrap array again
     u = unsafe_wrap_or_alloc(to, _u, size(mpi_interfaces.u))
@@ -224,7 +224,7 @@ function Adapt.adapt_structure(to, mpi_mortars::P4estMPIMortarContainer)
     # must be redesigned. This skeleton implementation here just exists just
     # for compatibility with the rest of the KA.jl solver code
 
-    _u = Adapt.adapt_structure(to, mpi_mortars._u)
+    _u = adapt(to, mpi_mortars._u)
     _node_indices = mpi_mortars._node_indices
     _normal_directions = mpi_mortars._normal_directions
 
diff --git a/test/Project.toml b/test/Project.toml
@@ -1,5 +1,6 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Convex = "f65535da-76fb-5f13-bab9-19810c17039a"
diff --git a/test/test_p4est_2d.jl b/test/test_p4est_2d.jl
@@ -27,6 +27,11 @@ isdir(outdir) && rm(outdir, recursive = true)
         du_ode = similar(u_ode)
         @test (@allocated Trixi.rhs!(du_ode, u_ode, semi, t)) < 1000
     end
+    semi32 = Trixi.trixi_adapt(Array, Float32, semi)
+    @test real(semi32.solver) == Float32
+    @test real(semi32.solver.basis) == Float32
+    @test real(semi32.solver.mortar) == Float32
+    @test real(semi32.mesh) == Float32
 end
 
 @trixi_testset "elixir_advection_nonconforming_flag.jl" begin
diff --git a/test/test_unstructured_2d.jl b/test/test_unstructured_2d.jl
@@ -2,6 +2,7 @@ module TestExamplesUnstructuredMesh2D
 
 using Test
 using Trixi
+using Adapt
 
 include("test_trixi.jl")
 
@@ -32,6 +33,11 @@ isdir(outdir) && rm(outdir, recursive = true)
         du_ode = similar(u_ode)
         @test (@allocated Trixi.rhs!(du_ode, u_ode, semi, t)) < 1000
     end
+    semi32 = Trixi.trixi_adapt(Array, Float32, semi)
+    @test real(semi32.solver) == Float32
+    @test real(semi32.solver.basis) == Float32
+    @test real(semi32.solver.mortar) == Float32
+    @test_broken real(semi32.mesh) == Float32
 end
 
 @trixi_testset "elixir_euler_free_stream.jl" begin
