Multi-GPU support: P4estMesh3D

NEWS.md

@@ -18,8 +18,9 @@ for human readability.
   Makie, and `trixiheatmap` accepts the attributes of `Makie.mesh`. Plots of slices of
   3D `DGMultiMesh` solutions and plots of 1D `DGMultiMesh` solutions now show the
   correct axis labels ([#3405]).
-- Added multi-GPU support for 2D `P4estMesh`es, i.e., GPU computations with MPI
-  ([#3240]). A local CUDA/ROCm-aware MPI installation is required. For configuration, see
+- Added multi-GPU support for 2D and 3D `P4estMesh`es, i.e., GPU computations with MPI
+  ([#3240], [#3364]). A local CUDA/ROCm-aware MPI installation is required.
+  For configuration, see
   [Parallelization](https://trixi-framework.org/TrixiDocumentation/stable/parallelization/#parallel_system_MPI)
   in the docs. Currently, only hyperbolic problems (no `SemidiscretizationHyperbolicParabolic`)
   on conforming meshes (no mortars, thus also no AMR) are supported.

docs/src/heterogeneous.md

@@ -57,12 +57,12 @@ are supported on conforming (non-adaptive) `P4estMesh{2}`s and `P4estMesh{3}`s a
 see, e.g., `examples/p4est_2d_dgsem/elixir_navierstokes_lid_driven_cavity.jl` and
 `examples/p4est_3d_dgsem/elixir_navierstokes_convergence.jl`.
 
-Hyperbolic problems on conforming `P4estMesh`es in 2D can also be run on multiple GPUs
-with MPI, using one MPI rank per GPU. The setup is the same as above, and the elixir is
-started with MPI as described in [Parallelization](@ref parallel_system_MPI).
+Hyperbolic problems on conforming `P4estMesh`es in 2D and 3D can also be run on
+multiple GPUs with MPI, using one MPI rank per GPU. The setup is the same as above, and
+the elixir is started with MPI as described in [Parallelization](@ref parallel_system_MPI).
 This requires a local CUDA/ROCm-aware MPI installation, since the MPI buffers are stored
-on the GPU and passed directly to MPI. Mixed hyperbolic-parabolic problems, 3D meshes,
-and non-conforming meshes (mortars, AMR) are not supported with multiple GPUs yet.
+on the GPU and passed directly to MPI. Mixed hyperbolic-parabolic problems and
+non-conforming meshes (mortars, AMR) are not supported with multiple GPUs yet.
 See the section on
 [GPU-aware MPI support](https://juliaparallel.org/MPI.jl/stable/usage/#GPU-aware-MPI-support)
 in the MPI.jl documentation for details, including example scripts such as

src/callbacks_step/analysis_dg3d_parallel.jl

@@ -13,6 +13,13 @@ function calc_error_norms(func, u, t, analyzer,
     @unpack node_coordinates, inverse_jacobian = cache.elements
     @unpack u_local, u_tmp1, u_tmp2, x_local, x_tmp1, x_tmp2, jacobian_local, jacobian_tmp1, jacobian_tmp2 = cache_analysis
 
+    # Calculate error norms on the CPU, to ensure the order of summation is the same.
+    if trixi_backend(u) !== nothing
+        node_coordinates = Array(node_coordinates)
+        inverse_jacobian = Array(inverse_jacobian)
+        u = Array(u)
+    end
+
     # Set up data structures
     l2_error = zero_return_type(func, SVector{nvariables(equations), eltype(u)},
                                 typeof(equations))
@@ -50,7 +57,7 @@ function calc_error_norms(func, u, t, analyzer,
     return reduce_error_norms_mpi(l2_error, linf_error, volume)
 end
 
-function integrate_via_indices(func::Func, u,
+function integrate_via_indices(func::Func, backend::Nothing, u,
                                mesh::Union{P4estMeshParallel{3}, T8codeMeshParallel{3}},
                                equations,
                                dg::DGSEM, cache, args...; normalize = true) where {Func}
@@ -74,4 +81,45 @@ function integrate_via_indices(func::Func, u,
 
     return reduce_integral_and_volume_mpi(integral, volume, normalize)
 end
+
+function integrate_via_indices(func::Func, backend::Backend, u,
+                               mesh::Union{P4estMeshParallel{3}, T8codeMeshParallel{3}},
+                               equations,
+                               dg::DGSEM, cache, args...; normalize = true) where {Func}
+    @unpack weights = dg.basis
+    @unpack inverse_jacobian = cache.elements
+
+    function local_plus((integral, total_volume), (integral_element, volume_element))
+        return (integral + integral_element, total_volume + volume_element)
+    end
+
+    # `func(u, 1,1,1,1, equations, dg, args...)` might access GPU memory
+    # so we have to rely on the compiler to correctly infer the type of the integral here.
+    # TODO: Technically we need device_promote_op here that "infers" the function within the context of the GPU.
+    integral0 = zero(Base.promote_op(func, typeof(u), Int, Int, Int, Int,
+                                     typeof(equations), typeof(dg),
+                                     map(typeof, args)...))
+    volume0 = zero(eltype(weights)) * zero(eltype(inverse_jacobian))
+    init = neutral = (integral0, volume0)
+
+    # Use quadrature to numerically integrate over entire domain
+    num_elements = nelements(dg, cache)
+    integral, volume = AcceleratedKernels.mapreduce(local_plus, 1:num_elements,
+                                                    backend; init,
+                                                    neutral) do element
+        # Initialize integral with zeros of the right shape
+        local_integral, local_total_volume = neutral
+
+        for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
+            volume_jacobian = abs(inv(inverse_jacobian[i, j, k, element]))
+            local_integral += volume_jacobian * weights[i] * weights[j] * weights[k] *
+                              func(u, i, j, k, element, equations, dg, args...)
+            local_total_volume += volume_jacobian * weights[i] * weights[j] * weights[k]
+        end
+
+        return (local_integral, local_total_volume)
+    end
+
+    return reduce_integral_and_volume_mpi(integral, volume, normalize)
+end
 end # @muladd

src/solvers/dgsem_p4est/dg_2d_parallel_gpu.jl

@@ -197,42 +197,6 @@ Base.@propagate_inbounds function mpi_buffer_node_offset(mpi_interfaces_buffer_o
     return mpi_interfaces_buffer_offsets[interface] + (i - 1) * nvariables(equations)
 end
 
-Base.@propagate_inbounds function get_mpi_buffer_node_vars(mpi_buffer, equations,
-                                                           offset)
-    # Explicit bounds check, which can be removed by calling this function with `@inbounds`
-    @boundscheck checkbounds(mpi_buffer, offset .+ eachvariable(equations))
-    # Assume inbounds access now
-    return SVector(ntuple(@inline(v->@inbounds mpi_buffer[offset + v]),
-                          Val(nvariables(equations))))
-end
-
-# The MPI send buffer is filled in `prolong2mpiinterfaces!`, and the kernels have
-# finished when this is called, see `rhs_hyperbolic!`.
-function start_mpi_send!(backend::Backend, mpi_cache::P4estMPICache,
-                         mesh::Union{P4estMeshParallel{2}, T8codeMeshParallel{2}},
-                         equations, dg, cache)
-    # TODO GPU: MPI mortars are not yet implemented
-    @assert isempty(eachmpimortar(dg, cache)) "MPI mortars are not yet supported on GPU backends"
-
-    for (index, rank) in enumerate(mpi_cache.mpi_neighbor_ranks)
-        send_buffer = mpi_buffer(get_mpi_send_buffer(mpi_cache, index))
-        mpi_cache.mpi_send_requests[index] = MPI.Isend(send_buffer, rank, mpi_rank(),
-                                                       mpi_comm())
-    end
-
-    return nothing
-end
-
-# The received data is not unpacked but read directly from the MPI receive buffer
-# by `calc_mpi_interface_flux!`. Thus, we can wait for all data at once.
-function finish_mpi_receive!(backend::Backend, mpi_cache::P4estMPICache,
-                             mesh::Union{P4estMeshParallel{2}, T8codeMeshParallel{2}},
-                             equations, dg, cache)
-    MPI.Waitall(mpi_cache.mpi_recv_requests)
-
-    return nothing
-end
-
 # Compute the surface fluxes at all MPI interfaces. In contrast to the CPU version,
 # the local data is taken directly from `u` and the remote data directly from the
 # MPI receive buffer.