Refactor DGMulti to use a shared array container

src/callbacks_step/analysis_dgmulti.jl

@@ -10,7 +10,7 @@ function calc_error_norms(func, u, t, analyzer,
                           dg::DGMulti{NDIMS}, cache, cache_analysis) where {NDIMS}
     rd = dg.basis
     md = mesh.md
-    @unpack u_values = cache
+    (; u_values) = cache.common_arrays
 
     # interpolate u to quadrature points
     apply_to_each_field(mul_by!(rd.Vq), u_values, u)
@@ -31,7 +31,7 @@ function integrate(func::Func, u, mesh::DGMultiMesh,
                    equations, dg::DGMulti, cache; normalize = true) where {Func}
     rd = dg.basis
     md = mesh.md
-    @unpack u_values = cache
+    (; u_values) = cache.common_arrays
 
     # interpolate u to quadrature points
     apply_to_each_field(mul_by!(rd.Vq), u_values, u)
@@ -47,7 +47,7 @@ function analyze(::typeof(entropy_timederivative), du, u, t,
                  mesh::DGMultiMesh, equations, dg::DGMulti, cache)
     rd = dg.basis
     md = mesh.md
-    @unpack u_values = cache
+    (; u_values) = cache.common_arrays
 
     # interpolate u, du to quadrature points
     du_values = similar(u_values) # Todo: DGMulti. Can we move this to the analysis cache somehow?

src/solvers/dgmulti/dg.jl

@@ -42,7 +42,7 @@ end
     end
 end
 
-@inline nelements(dg::DGMulti, cache) = size(cache.u_values)[end]
+@inline nelements(dg::DGMulti, cache) = size(cache.common_arrays.u_values)[end]
 
 """
     eachdim(mesh)
@@ -161,6 +161,30 @@ function set_zero!(du, dg::DGMulti, other_args...)
     return nothing
 end
 
+# Holds arrays shared across most DGMulti cache types:
+# solution values at volume/face quadrature points and thread-local scratch storage.
+struct DGMultiCommonArrays{uType, ufType, ffType, lType}
+    u_values::uType
+    u_face_values::ufType
+    flux_face_values::ffType
+    local_values_threaded::lType
+end
+
+# Allocates arrays shared across most DGMulti cache types.
+function allocate_common_dgmulti_arrays(mesh::DGMultiMesh, equations, dg::DGMulti,
+                                        uEltype)
+    rd = dg.basis
+    md = mesh.md
+    nvars = nvariables(equations)
+    u_values = allocate_nested_array(uEltype, nvars, size(md.xq), dg)
+    u_face_values = allocate_nested_array(uEltype, nvars, size(md.xf), dg)
+    flux_face_values = allocate_nested_array(uEltype, nvars, size(md.xf), dg)
+    local_values_threaded = [allocate_nested_array(uEltype, nvars, (rd.Nq,), dg)
+                             for _ in 1:Threads.maxthreadid()]
+    return DGMultiCommonArrays(u_values, u_face_values, flux_face_values,
+                               local_values_threaded)
+end
+
 # Constructs cache variables for both affine and non-affine (curved) DGMultiMeshes
 function create_cache(mesh::DGMultiMesh{NDIMS}, equations, dg::DGMultiWeakForm, RealT,
                       uEltype) where {NDIMS}
@@ -173,23 +197,13 @@ function create_cache(mesh::DGMultiMesh{NDIMS}, equations, dg::DGMultiWeakForm,
     # ∫f(u) * dv/dx_i = ∑_j (Vq*Drst[i])'*diagm(wq)*(rstxyzJ[i,j].*f(Vq*u))
     weak_differentiation_matrices = map(D -> -M \ ((Vq * D)' * Diagonal(wq)), Drst)
 
-    nvars = nvariables(equations)
-
-    # storage for volume quadrature values, face quadrature values, flux values
-    u_values = allocate_nested_array(uEltype, nvars, size(md.xq), dg)
-    u_face_values = allocate_nested_array(uEltype, nvars, size(md.xf), dg)
-    flux_face_values = allocate_nested_array(uEltype, nvars, size(md.xf), dg)
     if typeof(rd.approximation_type) <:
        Union{SBP, AbstractNonperiodicDerivativeOperator}
         lift_scalings = rd.wf ./ rd.wq[rd.Fmask] # lift scalings for diag-norm SBP operators
     else
         lift_scalings = nothing
     end
 
-    # local storage for volume integral and source computations
-    local_values_threaded = [allocate_nested_array(uEltype, nvars, (rd.Nq,), dg)
-                             for _ in 1:Threads.maxthreadid()]
-
     # For curved meshes, we interpolate geometric terms from nodal points to quadrature points.
     # For affine meshes, we just access one element of this interpolated data.
     dxidxhatj = map(x -> rd.Vq * x, md.rstxyzJ)
@@ -198,21 +212,23 @@ function create_cache(mesh::DGMultiMesh{NDIMS}, equations, dg::DGMultiWeakForm,
     invJ = inv.(rd.Vq * md.J)
 
     # for scaling by curved geometric terms (not used by affine DGMultiMesh)
+    nvars = nvariables(equations)
     flux_threaded = [[allocate_nested_array(uEltype, nvars, (rd.Nq,), dg)
                       for _ in 1:NDIMS] for _ in 1:Threads.maxthreadid()]
     rotated_flux_threaded = [allocate_nested_array(uEltype, nvars, (rd.Nq,), dg)
                              for _ in 1:Threads.maxthreadid()]
 
+    common_arrays = allocate_common_dgmulti_arrays(mesh, equations, dg, uEltype)
+
     return (; md, weak_differentiation_matrices, lift_scalings, invJ, dxidxhatj,
-            u_values, u_face_values, flux_face_values,
-            local_values_threaded, flux_threaded, rotated_flux_threaded)
+            common_arrays, flux_threaded, rotated_flux_threaded)
 end
 
 function compute_coefficients!(::Nothing, u, initial_condition, t,
                                mesh::DGMultiMesh, equations, dg::DGMulti, cache)
     md = mesh.md
     rd = dg.basis
-    @unpack u_values = cache
+    (; u_values) = cache.common_arrays
 
     # evaluate the initial condition at quadrature points
     @threaded for i in each_quad_node_global(mesh, dg, cache)
@@ -345,7 +361,7 @@ end
 function prolong2interfaces!(cache, u,
                              mesh::DGMultiMesh, equations, dg::DGMulti)
     rd = dg.basis
-    @unpack u_face_values = cache
+    (; u_face_values) = cache.common_arrays
     apply_to_each_field(mul_by!(rd.Vf), u_face_values, u)
 
     return nothing
@@ -358,7 +374,8 @@ end
                                          have_nonconservative_terms::False, equations,
                                          volume_integral::VolumeIntegralWeakForm,
                                          dg::DGMulti, cache)
-    @unpack weak_differentiation_matrices, dxidxhatj, u_values, local_values_threaded = cache
+    @unpack weak_differentiation_matrices, dxidxhatj = cache
+    (; u_values, local_values_threaded) = cache.common_arrays
 
     flux_values = local_values_threaded[Threads.threadid()]
     for i in eachdim(mesh)
@@ -386,7 +403,8 @@ end
                                          have_nonconservative_terms::False, equations,
                                          volume_integral::VolumeIntegralWeakForm,
                                          dg::DGMulti, cache) where {NDIMS}
-    (; weak_differentiation_matrices, dxidxhatj, u_values) = cache
+    (; weak_differentiation_matrices, dxidxhatj) = cache
+    (; u_values) = cache.common_arrays
 
     flux_values = cache.flux_threaded[Threads.threadid()]
     for i in eachdim(mesh)
@@ -427,7 +445,7 @@ function calc_volume_integral!(du, u, mesh::DGMultiMesh,
                                volume_integral::VolumeIntegralWeakForm, dg::DGMulti,
                                cache)
     rd = dg.basis
-    (; u_values) = cache
+    (; u_values) = cache.common_arrays
     # interpolate to quadrature points
     apply_to_each_field(mul_by!(rd.Vq), u_values, u)
 
@@ -447,7 +465,7 @@ function calc_interface_flux!(cache, surface_integral::SurfaceIntegralWeakForm,
     @unpack surface_flux = surface_integral
     md = mesh.md
     @unpack mapM, mapP, nxyzJ, Jf = md
-    @unpack u_face_values, flux_face_values = cache
+    (; u_face_values, flux_face_values) = cache.common_arrays
 
     @threaded for face_node_index in each_face_node_global(mesh, dg, cache)
 
@@ -469,7 +487,7 @@ function calc_interface_flux!(cache, surface_integral::SurfaceIntegralWeakForm,
     flux_conservative, flux_nonconservative = surface_integral.surface_flux
     md = mesh.md
     @unpack mapM, mapP, nxyzJ, Jf = md
-    @unpack u_face_values, flux_face_values = cache
+    (; u_face_values, flux_face_values) = cache.common_arrays
 
     @threaded for face_node_index in each_face_node_global(mesh, dg, cache)
 
@@ -503,7 +521,8 @@ function calc_surface_integral!(du, u, mesh::DGMultiMesh, equations,
                                 surface_integral::SurfaceIntegralWeakForm,
                                 dg::DGMulti, cache)
     rd = dg.basis
-    apply_to_each_field(mul_by_accum!(rd.LIFT), du, cache.flux_face_values)
+    apply_to_each_field(mul_by_accum!(rd.LIFT), du,
+                        cache.common_arrays.flux_face_values)
 
     return nothing
 end
@@ -513,7 +532,7 @@ function prolong2interfaces!(cache, u,
                              mesh::DGMultiMesh, equations, dg::DGMultiSBP)
     rd = dg.basis
     @unpack Fmask = rd
-    @unpack u_face_values = cache
+    (; u_face_values) = cache.common_arrays
     @threaded for e in eachelement(mesh, dg, cache)
         for (i, fid) in enumerate(Fmask)
             u_face_values[i, e] = u[fid, e]
@@ -529,7 +548,8 @@ function calc_surface_integral!(du, u, mesh::DGMultiMesh, equations,
                                 surface_integral::SurfaceIntegralWeakForm,
                                 dg::DGMultiSBP, cache)
     rd = dg.basis
-    @unpack flux_face_values, lift_scalings = cache
+    (; flux_face_values) = cache.common_arrays
+    @unpack lift_scalings = cache
 
     @threaded for e in eachelement(mesh, dg, cache)
         for i in each_face_node(mesh, dg, cache)
@@ -563,7 +583,7 @@ function calc_single_boundary_flux!(cache, t, boundary_condition, boundary_key,
                                     dg::DGMulti{NDIMS}) where {NDIMS}
     rd = dg.basis
     md = mesh.md
-    @unpack u_face_values, flux_face_values = cache
+    (; u_face_values, flux_face_values) = cache.common_arrays
     @unpack xyzf, nxyzJ, Jf = md
     @unpack surface_flux = dg.surface_integral
 
@@ -598,8 +618,8 @@ function calc_single_boundary_flux!(cache, t, boundary_condition, boundary_key,
         end
     end
 
-    # Note: modifying the values of the reshaped array modifies the values of cache.flux_face_values.
-    # However, we don't have to re-reshape, since cache.flux_face_values still retains its original shape.
+    # Note: modifying the values of the reshaped array modifies the values of cache.common_arrays.flux_face_values.
+    # However, we don't have to re-reshape, since cache.common_arrays.flux_face_values still retains its original shape.
 
     return nothing
 end
@@ -622,8 +642,8 @@ function calc_single_boundary_flux!(cache, t, boundary_condition, boundary_key,
     # https://github.com/JuliaLang/julia/issues/36313#issuecomment-782336300.
     reshape_by_face(u) = Base.ReshapedArray(u, (num_pts_per_face, num_faces_total), ())
 
-    u_face_values = reshape_by_face(cache.u_face_values)
-    flux_face_values = reshape_by_face(cache.flux_face_values)
+    u_face_values = reshape_by_face(cache.common_arrays.u_face_values)
+    flux_face_values = reshape_by_face(cache.common_arrays.flux_face_values)
     Jf = reshape_by_face(md.Jf)
     nxyzJ, xyzf = reshape_by_face.(md.nxyzJ), reshape_by_face.(md.xyzf) # broadcast over nxyzJ::NTuple{NDIMS,Matrix}
 
@@ -648,8 +668,8 @@ function calc_single_boundary_flux!(cache, t, boundary_condition, boundary_key,
         end
     end
 
-    # Note: modifying the values of the reshaped array modifies the values of cache.flux_face_values.
-    # However, we don't have to re-reshape, since cache.flux_face_values still retains its original shape.
+    # Note: modifying the values of the reshaped array modifies the values of cache.common_arrays.flux_face_values.
+    # However, we don't have to re-reshape, since cache.common_arrays.flux_face_values still retains its original shape.
 
     return nothing
 end
@@ -675,7 +695,8 @@ function invert_jacobian!(du, mesh::DGMultiMesh{NDIMS, <:NonAffine}, equations,
                           dg::DGMulti, cache; scaling = -1) where {NDIMS}
     # Vq = interpolation matrix to quadrature points, Pq = quadrature-based L2 projection matrix
     (; Pq, Vq) = dg.basis
-    (; local_values_threaded, invJ) = cache
+    (; invJ) = cache
+    (; local_values_threaded) = cache.common_arrays
 
     @threaded for e in eachelement(mesh, dg, cache)
         du_at_quad_points = local_values_threaded[Threads.threadid()]
@@ -711,7 +732,7 @@ function calc_sources!(du, u, t, source_terms,
     rd = dg.basis
     md = mesh.md
     @unpack Pq = rd
-    @unpack u_values, local_values_threaded = cache
+    (; u_values, local_values_threaded) = cache.common_arrays
     @threaded for e in eachelement(mesh, dg, cache)
         source_values = local_values_threaded[Threads.threadid()]
 

src/solvers/dgmulti/flux_differencing.jl

@@ -302,25 +302,18 @@ function create_cache(mesh::DGMultiMesh, equations, dg::DGMultiFluxDiffSBP,
     # for use with flux differencing schemes
     Qrst_skew = compute_flux_differencing_SBP_matrices(dg)
 
-    # Todo: DGMulti. Factor common storage into a struct (MeshDataCache?) for reuse across solvers?
-    # storage for volume quadrature values, face quadrature values, flux values
-    nvars = nvariables(equations)
-    u_values = allocate_nested_array(uEltype, nvars, size(md.xq), dg)
-    u_face_values = allocate_nested_array(uEltype, nvars, size(md.xf), dg)
-    flux_face_values = allocate_nested_array(uEltype, nvars, size(md.xf), dg)
     lift_scalings = rd.wf ./ rd.wq[rd.Fmask] # lift scalings for diag-norm SBP operators
 
-    local_values_threaded = [allocate_nested_array(uEltype, nvars, (rd.Nq,), dg)
-                             for _ in 1:Threads.maxthreadid()]
-
+    nvars = nvariables(equations)
     # Use an array of SVectors (chunks of `nvars` are contiguous in memory) to speed up flux differencing
     fluxdiff_local_threaded = [zeros(SVector{nvars, uEltype}, rd.Nq)
                                for _ in 1:Threads.maxthreadid()]
 
+    common_arrays = allocate_common_dgmulti_arrays(mesh, equations, dg, uEltype)
+
     return (; md, Qrst_skew, dxidxhatj = md.rstxyzJ,
             invJ = inv.(md.J), lift_scalings, inv_wq = inv.(rd.wq),
-            u_values, u_face_values, flux_face_values,
-            local_values_threaded, fluxdiff_local_threaded)
+            common_arrays, fluxdiff_local_threaded)
 end
 
 # most general create_cache: works for `DGMultiFluxDiff{<:Polynomial}`
@@ -372,20 +365,23 @@ function create_cache(mesh::DGMultiMesh, equations, dg::DGMultiFluxDiff, RealT,
     interpolated_geometric_terms = map(x -> [Vq; Vf] * x, mesh.md.rstxyzJ)
     J = rd.Vq * md.J
 
+    common_arrays = DGMultiCommonArrays(u_values, u_face_values, flux_face_values,
+                                        local_values_threaded)
+
     return (; md, Qrst_skew, VhP, Ph,
             invJ = inv.(J), dxidxhatj = interpolated_geometric_terms,
             entropy_var_values, projected_entropy_var_values,
             entropy_projected_u_values,
-            u_values, u_face_values, flux_face_values,
-            local_values_threaded, fluxdiff_local_threaded, rhs_local_threaded)
+            common_arrays, fluxdiff_local_threaded, rhs_local_threaded)
 end
 
 # TODO: DGMulti. Address hard-coding of `entropy2cons!` and `cons2entropy!` for this function.
 function entropy_projection!(cache, u, mesh::DGMultiMesh, equations, dg::DGMulti)
     rd = dg.basis
     @unpack Vq = rd
-    @unpack VhP, entropy_var_values, u_values = cache
+    @unpack VhP, entropy_var_values = cache
     @unpack projected_entropy_var_values, entropy_projected_u_values = cache
+    (; u_values) = cache.common_arrays
 
     apply_to_each_field(mul_by!(Vq), u_values, u)
 

src/solvers/dgmulti/flux_differencing_gauss_sbp.jl

@@ -454,9 +454,10 @@ function entropy_projection!(cache, u, mesh::DGMultiMesh, equations,
                              dg::DGMultiFluxDiff{<:GaussSBP})
     rd = dg.basis
     @unpack Vq = rd
-    @unpack VhP, entropy_var_values, u_values = cache
+    @unpack VhP, entropy_var_values = cache
     @unpack projected_entropy_var_values, entropy_projected_u_values = cache
     @unpack interp_matrix_lobatto_to_gauss, interp_matrix_gauss_to_face = cache
+    (; u_values) = cache.common_arrays
 
     @threaded for e in eachelement(mesh, dg, cache)
         apply_to_each_field(mul_by!(interp_matrix_lobatto_to_gauss),
@@ -501,7 +502,7 @@ function calc_surface_integral!(du, u, mesh::DGMultiMesh, equations,
         # applies LIFT matrix, output is stored at Gauss nodes
         gauss_volume_local = gauss_volume_local_threaded[Threads.threadid()]
         apply_to_each_field(mul_by!(gauss_LIFT), gauss_volume_local,
-                            view(cache.flux_face_values, :, e))
+                            view(cache.common_arrays.flux_face_values, :, e))
 
         for i in eachindex(gauss_volume_local)
             du[i, e] = du[i, e] + gauss_volume_local[i]

src/solvers/dgmulti/sbp.jl

@@ -179,10 +179,8 @@ function create_cache(mesh::DGMultiMesh, equations,
                       dg::DGMultiFluxDiffPeriodicFDSBP, RealT, uEltype)
     md = mesh.md
 
-    # storage for volume quadrature values, face quadrature values, flux values
-    nvars = nvariables(equations)
-    u_values = allocate_nested_array(uEltype, nvars, size(md.xq), dg)
-    return (; u_values, invJ = inv.(md.J))
+    common_arrays = allocate_common_dgmulti_arrays(mesh, equations, dg, uEltype)
+    return (; common_arrays, invJ = inv.(md.J))
 end
 
 # Specialize calc_volume_integral for periodic SBP operators (assumes the operator is sparse).