Combine interact! kernels to reduce GPU latency

Author: efaulhaber

src/general/abstract_system.jl

@@ -26,13 +26,20 @@ vtkname(system::AbstractBoundarySystem) = "boundary"
 # Number of particles in the system
 @inline nparticles(system) = length(system.mass)
 
-# Number of particles in the system whose positions are to be integrated (corresponds to the size of u and du)
+# Number of particles in the system whose positions are to be integrated
+# (corresponds to the size of u and du).
+# See comment below for `each_integrated_particle` for more details.
 @inline n_integrated_particles(system) = nparticles(system)
 
+# Iterator over all particle indices in the system
 @inline eachparticle(system::AbstractSystem) = each_active_particle(system)
 @inline eachparticle(initial_condition) = Base.OneTo(nparticles(initial_condition))
 
-# Wrapper for systems with `SystemBuffer`
+# Iterator over all particle indices in the system whose positions are to be integrated
+# (corresponds to the size of u and du).
+# For most systems, this is the same as `eachparticle`, but for the
+# `TotalLagrangianSPHSystem`, the clamped particles are included in `eachparticle`
+# but not in `each_integrated_particle`.
 @inline each_integrated_particle(system) = each_integrated_particle(system, buffer(system))
 @inline function each_integrated_particle(system, ::Nothing)
     return Base.OneTo(n_integrated_particles(system))

src/general/semidiscretization.jl

@@ -55,6 +55,7 @@ struct Semidiscretization{BACKEND, S, RU, RV, NS, UCU, IT}
     ranges_v                :: RV
     neighborhood_searches   :: NS
     parallelization_backend :: BACKEND
+    interaction_timers      :: TIMERS
     update_callback_used    :: UCU
     integrate_tlsph         :: IT # `false` if TLSPH integration is decoupled
 
@@ -64,19 +65,20 @@ struct Semidiscretization{BACKEND, S, RU, RV, NS, UCU, IT}
     # and by Adapt.jl.
     function Semidiscretization(systems::Tuple, ranges_u, ranges_v, neighborhood_searches,
                                 parallelization_backend::PointNeighbors.ParallelizationBackend,
-                                update_callback_used, integrate_tlsph)
-        new{typeof(parallelization_backend), typeof(systems), typeof(ranges_u),
+                                interaction_timers, update_callback_used, integrate_tlsph)
+        new{typeof(interaction_timers), typeof(systems), typeof(ranges_u),
             typeof(ranges_v), typeof(neighborhood_searches),
-            typeof(update_callback_used),
-            typeof(integrate_tlsph)}(systems, ranges_u, ranges_v,
-                                     neighborhood_searches, parallelization_backend,
+            typeof(parallelization_backend), typeof(update_callback_used),
+            typeof(integrate_tlsph)}(systems, ranges_u, ranges_v, neighborhood_searches,
+                                     parallelization_backend, interaction_timers,
                                      update_callback_used, integrate_tlsph)
     end
 end
 
 function Semidiscretization(systems::Union{AbstractSystem, Nothing}...;
                             neighborhood_search=GridNeighborhoodSearch{ndims(first(systems))}(),
-                            parallelization_backend=PolyesterBackend())
+                            parallelization_backend=PolyesterBackend(),
+                            interaction_timers=IndividualTimers())
     systems = filter(system -> !isnothing(system), systems)
 
     if isempty(systems)
@@ -120,10 +122,13 @@ function Semidiscretization(systems::Union{AbstractSystem, Nothing}...;
     integrate_tlsph = Ref(true)
 
     return Semidiscretization(systems, ranges_u, ranges_v, searches,
-                              parallelization_backend, update_callback_used,
-                              integrate_tlsph)
+                              parallelization_backend, interaction_timers,
+                              update_callback_used, integrate_tlsph)
 end
 
+struct IndividualTimers end
+struct CombinedTimers end
+
 # Inline show function e.g. Semidiscretization(neighborhood_search=...)
 function Base.show(io::IO, semi::Semidiscretization)
     @nospecialize semi # reduce precompilation time
@@ -649,8 +654,8 @@ end
     return -damping_coefficient * velocity
 end
 
-function system_interaction!(dv_ode, v_ode, u_ode, semi)
-    # Call `interact!` for each pair of systems
+function system_interaction!(dv_ode, v_ode, u_ode,
+                             semi::Semidiscretization{IndividualTimers})
     foreach_system(semi) do system
         foreach_system(semi) do neighbor
             # Construct string for the interactions timer.
@@ -663,7 +668,34 @@ function system_interaction!(dv_ode, v_ode, u_ode, semi)
                 timer_str = ""
             end
 
-            interact!(dv_ode, v_ode, u_ode, system, neighbor, semi, timer_str=timer_str)
+            # Call individual `interact!` for this pair of systems.
+            # On GPUs, this is fully synchronized, so we get a separate timer for each
+            # pair of systems.
+            @trixi_timeit timer() timer_str begin
+                interact!(dv_ode, v_ode, u_ode, system, neighbor, semi)
+            end
+        end
+    end
+
+    return dv_ode
+end
+
+function system_interaction!(dv_ode, v_ode, u_ode,
+                             semi::Semidiscretization{CombinedTimers})
+    foreach_system(semi) do system
+        # Construct string for the interactions timer.
+        # Avoid allocations from string construction when no timers are used.
+        if timeit_debug_enabled()
+            system_index = system_indices(system, semi)
+            timer_str = "$(timer_name(system))$system_index-*"
+        else
+            timer_str = ""
+        end
+
+        # Call a combined `interact!` for all interactions of this system with other systems.
+        # On GPUs, this is fully synchronized, so we get a separate timer for each system.
+        @trixi_timeit timer() timer_str begin
+            interact_combined!(dv_ode, v_ode, u_ode, system, semi)
         end
     end
 
@@ -674,16 +706,51 @@ end
 # One can benchmark, e.g. the fluid-fluid interaction, with:
 # dv_ode, du_ode = copy(sol.u[end]).x; v_ode, u_ode = copy(sol.u[end]).x;
 # @btime TrixiParticles.interact!($dv_ode, $v_ode, $u_ode, $fluid_system, $fluid_system, $semi);
-@inline function interact!(dv_ode, v_ode, u_ode, system, neighbor, semi; timer_str="")
+function interact!(dv_ode, v_ode, u_ode, system, neighbor, semi)
     dv = wrap_v(dv_ode, system, semi)
     v_system = wrap_v(v_ode, system, semi)
     u_system = wrap_u(u_ode, system, semi)
 
     v_neighbor = wrap_v(v_ode, neighbor, semi)
     u_neighbor = wrap_u(u_ode, neighbor, semi)
 
-    @trixi_timeit timer() timer_str begin
-        interact!(dv, v_system, u_system, v_neighbor, u_neighbor, system, neighbor, semi)
+    nhs = get_neighborhood_search(system, neighbor, semi)
+
+    # Loop over all particles that are integrated for this system, i.e., all particles
+    # for which `dv` has entries.
+    @threaded semi for particle in each_integrated_particle(system)
+        interact!(dv, v_system, u_system, v_neighbor, u_neighbor,
+                  system, neighbor, semi, nhs, particle)
+    end
+end
+
+# Benchmark the combined interaction for a system with:
+# dv_ode, du_ode = copy(sol.u[end]).x; v_ode, u_ode = copy(sol.u[end]).x;
+# @btime TrixiParticles.interact_combined!($dv_ode, $v_ode, $u_ode, $system, $semi);
+function interact_combined!(dv_ode, v_ode, u_ode, system, semi; synchronize=true)
+    dv = wrap_v(dv_ode, system, semi)
+    v_system = wrap_v(v_ode, system, semi)
+    u_system = wrap_u(u_ode, system, semi)
+
+    # Create an iterator combining systems with their wrapped arrays and NHS.
+    # Note that we cannot use `get_neighborhood_search` because this will return
+    # a NHS of a different type for TLSPH-TLSPH interactions, making the iterator tuple
+    # type-unstable. The different self-interaction NHS for TLSPH is handled
+    # by the `interact!` function for TLSPH.
+    system_index = system_indices(system, semi)
+    f(neighbor) = (neighbor, wrap_v(v_ode, neighbor, semi), wrap_u(u_ode, neighbor, semi),
+                   semi.neighborhood_searches[system_index][system_indices(neighbor, semi)])
+    iterator = map(f, semi.systems)
+
+    # Loop over all particles that are integrated for this system, i.e., all particles
+    # for which `dv` has entries.
+    @threaded semi for particle in each_integrated_particle(system)
+        # Now loop over all neighbor systems to avoid separate loops/kernels
+        # for each pair of systems.
+        foreach_noalloc(iterator) do (neighbor, v_neighbor, u_neighbor, nhs)
+            interact!(dv, v_system, u_system, v_neighbor, u_neighbor,
+                      system, neighbor, semi, nhs, particle)
+        end
     end
 end
 

src/schemes/boundary/open_boundary/mirroring.jl

@@ -165,10 +165,8 @@ function extrapolate_values!(system,
                                                             ndims(system) + 1,
                                                             eltype(system)}))
 
-        # TODO: Not public API
-        PointNeighbors.foreach_neighbor(fluid_coords, nhs, particle, ghost_node_position,
-                                        nhs.search_radius) do particle, neighbor, pos_diff,
-                                                              distance
+        foreach_neighbor(fluid_coords, nhs, particle, ghost_node_position,
+                         nhs.search_radius) do particle, neighbor, pos_diff, distance
             m_b = hydrodynamic_mass(fluid_system, neighbor)
             rho_b = current_density(v_fluid, fluid_system, neighbor)
             pressure_b = current_pressure(v_fluid, fluid_system, neighbor)
@@ -305,10 +303,8 @@ function extrapolate_values!(system, mirror_method::ZerothOrderMirroring,
         interpolated_pressure = Ref(zero(eltype(system)))
         interpolated_velocity = Ref(zero(particle_coords))
 
-        # TODO: Not public API
-        PointNeighbors.foreach_neighbor(fluid_coords, nhs, particle, ghost_node_position,
-                                        nhs.search_radius) do particle, neighbor, pos_diff,
-                                                              distance
+        foreach_neighbor(fluid_coords, nhs, particle, ghost_node_position,
+                         nhs.search_radius) do particle, neighbor, pos_diff, distance
             m_b = hydrodynamic_mass(fluid_system, neighbor)
             rho_b = current_density(v_fluid, fluid_system, neighbor)
             volume_b = m_b / rho_b

src/schemes/boundary/wall_boundary/rhs.jl

@@ -2,7 +2,7 @@
 function interact!(dv, v_particle_system, u_particle_system,
                    v_neighbor_system, u_neighbor_system,
                    particle_system::Union{AbstractBoundarySystem, OpenBoundarySystem},
-                   neighbor_system, semi)
+                   neighbor_system, semi, nhs, particle)
     # TODO Solids and moving boundaries should be considered in the continuity equation
     return dv
 end

src/schemes/fluid/weakly_compressible_sph/rhs.jl

@@ -2,9 +2,10 @@
 # in `neighbor_system` and updates `dv` accordingly.
 # It takes into account pressure forces, viscosity, and for `ContinuityDensity` updates
 # the density using the continuity equation.
-function interact!(dv, v_particle_system, u_particle_system,
-                   v_neighbor_system, u_neighbor_system,
-                   particle_system::WeaklyCompressibleSPHSystem, neighbor_system, semi)
+@inline function interact!(dv, v_particle_system, u_particle_system,
+                           v_neighbor_system, u_neighbor_system,
+                           particle_system::WeaklyCompressibleSPHSystem, neighbor_system,
+                           semi, nhs, particle)
     (; density_calculator, correction) = particle_system
 
     sound_speed = system_sound_speed(particle_system)
@@ -20,12 +21,8 @@ function interact!(dv, v_particle_system, u_particle_system,
     # debug_array = zeros(ndims(particle_system), nparticles(particle_system))
 
     # Loop over all pairs of particles and neighbors within the kernel cutoff
-    foreach_point_neighbor(particle_system, neighbor_system,
-                           system_coords, neighbor_system_coords, semi;
-                           points=each_integrated_particle(particle_system)) do particle,
-                                                                                neighbor,
-                                                                                pos_diff,
-                                                                                distance
+    foreach_neighbor(system_coords, neighbor_system_coords,
+                     nhs, particle) do particle, neighbor, pos_diff, distance
         # `foreach_point_neighbor` makes sure that `particle` and `neighbor` are
         # in bounds of the respective system. For performance reasons, we use `@inbounds`
         # in this hot loop to avoid bounds checking when extracting particle quantities.

src/schemes/structure/total_lagrangian_sph/rhs.jl

@@ -2,30 +2,34 @@
 function interact!(dv, v_particle_system, u_particle_system,
                    v_neighbor_system, u_neighbor_system,
                    particle_system::TotalLagrangianSPHSystem,
-                   neighbor_system::TotalLagrangianSPHSystem, semi;
+                   neighbor_system::TotalLagrangianSPHSystem,
+                   semi, nhs, particle;
                    integrate_tlsph=semi.integrate_tlsph[])
     # Different structures do not interact with each other (yet)
     particle_system === neighbor_system || return dv
 
     # Skip interaction if TLSPH systems are integrated separately
     integrate_tlsph || return dv
 
-    interact_structure_structure!(dv, v_particle_system, particle_system, semi)
+    interact_structure_structure!(dv, v_particle_system, particle_system, semi, particle)
 end
 
 # Function barrier without dispatch for unit testing
-@inline function interact_structure_structure!(dv, v_system, system, semi)
+@inline function interact_structure_structure!(dv, v_system, system, semi, particle)
     (; penalty_force) = system
 
     # Everything here is done in the initial coordinates
     system_coords = initial_coordinates(system)
 
+    # Specialized neighborhood search for structure-structure interaction that is
+    # optimized for finding neighbors in the initial configuration.
+    nhs = system.self_interaction_nhs
+
     # Loop over all pairs of particles and neighbors within the kernel cutoff.
     # For structure-structure interaction, this has to happen in the initial coordinates.
-    foreach_point_neighbor(system, system, system_coords, system_coords, semi;
-                           points=each_integrated_particle(system)) do particle, neighbor,
-                                                                       initial_pos_diff,
-                                                                       initial_distance
+    foreach_neighbor(system_coords, system_coords, nhs,
+                     particle) do particle, neighbor,
+                                  initial_pos_diff, initial_distance
         # Only consider particles with a distance > 0.
         # See `src/general/smoothing_kernels.jl` for more details.
         initial_distance^2 < eps(initial_smoothing_length(system)^2) && return
@@ -76,7 +80,7 @@ end
 function interact!(dv, v_particle_system, u_particle_system,
                    v_neighbor_system, u_neighbor_system,
                    particle_system::TotalLagrangianSPHSystem,
-                   neighbor_system::AbstractFluidSystem, semi;
+                   neighbor_system::AbstractFluidSystem, semi, nhs, particle;
                    integrate_tlsph=semi.integrate_tlsph[])
     # Skip interaction if TLSPH systems are integrated separately
     integrate_tlsph || return dv
@@ -87,12 +91,8 @@ function interact!(dv, v_particle_system, u_particle_system,
     neighbor_coords = current_coordinates(u_neighbor_system, neighbor_system)
 
     # Loop over all pairs of particles and neighbors within the kernel cutoff
-    foreach_point_neighbor(particle_system, neighbor_system, system_coords, neighbor_coords,
-                           semi;
-                           points=each_integrated_particle(particle_system)) do particle,
-                                                                                neighbor,
-                                                                                pos_diff,
-                                                                                distance
+    foreach_neighbor(system_coords, neighbor_coords,
+                     nhs, particle) do particle, neighbor, pos_diff, distance
         # Only consider particles with a distance > 0.
         # See `src/general/smoothing_kernels.jl` for more details.
         distance^2 < eps(initial_smoothing_length(particle_system)^2) && return
@@ -183,7 +183,8 @@ end
 function interact!(dv, v_particle_system, u_particle_system,
                    v_neighbor_system, u_neighbor_system,
                    particle_system::TotalLagrangianSPHSystem,
-                   neighbor_system::Union{WallBoundarySystem, OpenBoundarySystem}, semi;
+                   neighbor_system::Union{WallBoundarySystem, OpenBoundarySystem},
+                   semi, nhs, particle;
                    integrate_tlsph=semi.integrate_tlsph[])
     # TODO continuity equation?
     return dv