Improve GPU performance with shifting (trixi-framework#974)

* Improve GPU performance with shifting

* Improve update callback timers

* Fix TVF

* Add timer

* Fix

---------

Co-authored-by: Sven Berger <berger.sven@gmail.com>

Author: efaulhaber

src/callbacks/update.jl

@@ -88,20 +88,20 @@ function (update_callback!::UpdateCallback)(integrator)
         end
 
         # Update open boundaries first, since particles might be activated or deactivated
-        @trixi_timeit timer() "update open boundary" foreach_system(semi) do system
+        foreach_system(semi) do system
             update_open_boundary_eachstep!(system, v_ode, u_ode, semi, t, integrator)
         end
 
-        @trixi_timeit timer() "update particle packing" foreach_system(semi) do system
+        foreach_system(semi) do system
             update_particle_packing(system, v_ode, u_ode, semi, integrator)
         end
 
         # This is only used by the particle packing system and should be removed in the future
-        @trixi_timeit timer() "update TVF" foreach_system(semi) do system
+        foreach_system(semi) do system
             update_transport_velocity!(system, v_ode, semi, integrator)
         end
 
-        @trixi_timeit timer() "particle shifting" foreach_system(semi) do system
+        foreach_system(semi) do system
             particle_shifting_from_callback!(u_ode, shifting_technique(system), system,
                                              v_ode, semi, integrator)
         end

src/preprocessing/particle_packing/system.jl

@@ -269,7 +269,7 @@ function update_particle_packing(system::ParticlePackingSystem, v_ode, u_ode,
                                  semi, integrator)
     u = wrap_u(u_ode, system, semi)
 
-    update_position!(u, system, semi)
+    @trixi_timeit timer() "update particle packing" update_position!(u, system, semi)
 
     # Tell OrdinaryDiffEq that `integrator.u` has been modified
     u_modified!(integrator, true)
@@ -365,14 +365,16 @@ end
 # Update from `UpdateCallback` between time steps (skip for fixed systems)
 @inline function update_transport_velocity!(system::ParticlePackingSystem{<:Any, false},
                                             v_ode, semi, integrator)
-    v = wrap_v(v_ode, system, semi)
-    @threaded semi for particle in each_integrated_particle(system)
-        for i in 1:ndims(system)
-            system.advection_velocity[i, particle] = v[i, particle]
-
-            # The particle velocity is set to zero at the beginning of each time step to
-            # achieve a fully stationary state.
-            v[i, particle] = zero(eltype(system))
+    @trixi_timeit timer() "update packing TVF" begin
+        v = wrap_v(v_ode, system, semi)
+        @threaded semi for particle in each_integrated_particle(system)
+            for i in 1:ndims(system)
+                system.advection_velocity[i, particle] = v[i, particle]
+
+                # The particle velocity is set to zero at the beginning of each time step to
+                # achieve a fully stationary state.
+                v[i, particle] = zero(eltype(system))
+            end
         end
     end
 

src/schemes/boundary/open_boundary/system.jl

@@ -301,16 +301,18 @@ function update_open_boundary_eachstep!(system::OpenBoundarySystem, v_ode, u_ode
                                         semi, t, integrator)
     (; boundary_model) = system
 
-    u = wrap_u(u_ode, system, semi)
-    v = wrap_v(v_ode, system, semi)
+    @trixi_timeit timer() "update open boundary" begin
+        u = wrap_u(u_ode, system, semi)
+        v = wrap_v(v_ode, system, semi)
 
-    @trixi_timeit timer() "check domain" check_domain!(system, v, u, v_ode, u_ode, semi)
+        @trixi_timeit timer() "check domain" check_domain!(system, v, u, v_ode, u_ode, semi)
 
-    update_pressure_model!(system, v, u, semi, integrator.dt)
+        update_pressure_model!(system, v, u, semi, integrator.dt)
 
-    # Update density, pressure and velocity based on the specific boundary model
-    @trixi_timeit timer() "update boundary quantities" begin
-        update_boundary_quantities!(system, boundary_model, v, u, v_ode, u_ode, semi, t)
+        # Update density, pressure and velocity based on the specific boundary model
+        @trixi_timeit timer() "update boundary quantities" begin
+            update_boundary_quantities!(system, boundary_model, v, u, v_ode, u_ode, semi, t)
+        end
     end
 
     # Tell OrdinaryDiffEq that `integrator.u` has been modified

src/schemes/fluid/fluid.jl

@@ -146,6 +146,10 @@ end
     vdiff = current_velocity(v_particle_system, particle_system, particle) -
             current_velocity(v_neighbor_system, neighbor_system, neighbor)
 
+    vdiff += continuity_equation_shifting_term(shifting_technique(particle_system),
+                                               particle_system, neighbor_system,
+                                               particle, neighbor, rho_a, rho_b)
+
     dv[end, particle] += rho_a / rho_b * m_b * dot(vdiff, grad_kernel)
 
     # Artificial density diffusion should only be applied to systems representing a fluid
@@ -157,10 +161,6 @@ end
                            pos_diff, distance, m_b, rho_a, rho_b, particle_system,
                            grad_kernel)
     end
-
-    continuity_equation_shifting!(dv, shifting_technique(particle_system),
-                                  particle_system, neighbor_system,
-                                  particle, neighbor, grad_kernel, rho_a, rho_b, m_b)
 end
 
 function calculate_dt(v_ode, u_ode, cfl_number, system::AbstractFluidSystem, semi)

src/schemes/fluid/shifting_techniques.jl

@@ -49,10 +49,10 @@ end
 end
 
 # Additional term(s) in the continuity equation due to the shifting technique
-function continuity_equation_shifting!(dv, shifting,
-                                       particle_system, neighbor_system,
-                                       particle, neighbor, grad_kernel, rho_a, rho_b, m_b)
-    return dv
+@inline function continuity_equation_shifting_term(shifting, particle_system,
+                                                   neighbor_system,
+                                                   particle, neighbor, rho_a, rho_b)
+    return zero(SVector{ndims(particle_system), eltype(particle_system)})
 end
 
 @doc raw"""
@@ -324,10 +324,11 @@ See [`ParticleShiftingTechnique`](@ref).
 struct MomentumEquationTermSun2019 end
 
 # Additional term in the momentum equation due to the shifting technique
-@inline function dv_shifting(shifting::ParticleShiftingTechnique, system, neighbor_system,
-                             v_system, v_neighbor_system, particle, neighbor,
-                             m_a, m_b, rho_a, rho_b, pos_diff, distance,
-                             grad_kernel, correction)
+@propagate_inbounds function dv_shifting(shifting::ParticleShiftingTechnique, system,
+                                         neighbor_system,
+                                         v_system, v_neighbor_system, particle, neighbor,
+                                         m_a, m_b, rho_a, rho_b, pos_diff, distance,
+                                         grad_kernel, correction)
     return dv_shifting(shifting.momentum_equation_term, system, neighbor_system,
                        v_system, v_neighbor_system, particle, neighbor,
                        m_a, m_b, rho_a, rho_b, pos_diff, distance,
@@ -351,21 +352,19 @@ end
 end
 
 # `ParticleShiftingTechnique{<:Any, <:Any, true}` means `modify_continuity_equation=true`
-function continuity_equation_shifting!(dv,
-                                       shifting::ParticleShiftingTechnique{<:Any, <:Any,
-                                                                           true},
-                                       system, neighbor_system,
-                                       particle, neighbor, grad_kernel, rho_a, rho_b, m_b)
-    delta_v_diff = delta_v(system, particle) -
-                   delta_v(neighbor_system, neighbor)
-
-    dv[end, particle] += rho_a / rho_b * m_b * dot(delta_v_diff, grad_kernel)
-
-    second_continuity_equation_term!(dv, shifting.second_continuity_equation_term,
-                                     system, neighbor_system,
-                                     particle, neighbor, grad_kernel, rho_a, rho_b, m_b)
+@propagate_inbounds function continuity_equation_shifting_term(shifting::ParticleShiftingTechnique{<:Any,
+                                                                                                   <:Any,
+                                                                                                   true},
+                                                               system, neighbor_system,
+                                                               particle, neighbor,
+                                                               rho_a, rho_b)
+    delta_v_a = delta_v(system, particle)
+    delta_v_b = delta_v(neighbor_system, neighbor)
+    delta_v_diff = delta_v_a - delta_v_b
 
-    return dv
+    second_term = second_continuity_equation_term(shifting.second_continuity_equation_term,
+                                                  delta_v_a, delta_v_b, rho_a, rho_b)
+    return delta_v_diff + second_term
 end
 
 """
@@ -378,29 +377,23 @@ See [`ParticleShiftingTechnique`](@ref).
 """
 struct ContinuityEquationTermSun2019 end
 
-@inline function second_continuity_equation_term!(dv,
-                                                  ::ContinuityEquationTermSun2019,
-                                                  system, neighbor_system,
-                                                  particle, neighbor, grad_kernel,
-                                                  rho_a, rho_b, m_b)
-    rho_v = rho_a * delta_v(system, particle) + rho_b * delta_v(neighbor_system, neighbor)
-
-    dv[end, particle] += m_b / rho_b * dot(rho_v, grad_kernel)
-
-    return dv
+@propagate_inbounds function second_continuity_equation_term(::ContinuityEquationTermSun2019,
+                                                             delta_v_a, delta_v_b,
+                                                             rho_a, rho_b)
+    return delta_v_a + rho_b / rho_a * delta_v_b
 end
 
-@inline function second_continuity_equation_term!(dv, second_continuity_equation_term,
-                                                  system, neighbor_system,
-                                                  particle, neighbor, grad_kernel,
-                                                  rho_a, rho_b, m_b)
-    return dv
+@inline function second_continuity_equation_term(second_continuity_equation_term,
+                                                 delta_v_a, delta_v_b, rho_a, rho_b)
+    return zero(delta_v_a)
 end
 
 # `ParticleShiftingTechnique{<:Any, true}` means `update_everystage=true`
 function update_shifting!(system, shifting::ParticleShiftingTechnique{<:Any, true},
                           v, u, v_ode, u_ode, semi)
-    update_shifting_inner!(system, shifting, v, u, v_ode, u_ode, semi)
+    @trixi_timeit timer() "update shifting" begin
+        update_shifting_inner!(system, shifting, v, u, v_ode, u_ode, semi)
+    end
 end
 
 # `ParticleShiftingTechnique{<:Any, false}` means `update_everystage=false`
@@ -506,10 +499,10 @@ end
 function particle_shifting_from_callback!(u_ode,
                                           shifting::ParticleShiftingTechnique{<:Any, false},
                                           system, v_ode, semi, integrator)
-    @trixi_timeit timer() "particle shifting" begin
-        # Update the shifting velocity
-        update_shifting_from_callback!(system, shifting, v_ode, u_ode, semi)
+    # Update the shifting velocity
+    update_shifting_from_callback!(system, shifting, v_ode, u_ode, semi)
 
+    @trixi_timeit timer() "apply particle shifting" begin
         # Update the particle positions with the shifting velocity
         apply_particle_shifting!(u_ode, shifting, system, semi, integrator.dt)
     end
@@ -604,18 +597,18 @@ end
 # The function above misuses the pressure acceleration function by passing a Matrix as `p_a`.
 # This doesn't work with `tensile_instability_control`, so we disable TIC in this case.
 @inline function tensile_instability_control(m_a, m_b, rho_a, rho_b, p_a::SMatrix, p_b, W_a)
-    return pressure_acceleration_continuity_density(m_a, m_b, rho_a, rho_b, p_a, A_b, W_a)
+    return pressure_acceleration_continuity_density(m_a, m_b, rho_a, rho_b, p_a, p_b, W_a)
 end
 
-function continuity_equation_shifting!(dv, shifting::TransportVelocityAdami{true},
-                                       particle_system, neighbor_system,
-                                       particle, neighbor, grad_kernel, rho_a, rho_b, m_b)
+@propagate_inbounds function continuity_equation_shifting_term(::TransportVelocityAdami{true},
+                                                               particle_system,
+                                                               neighbor_system,
+                                                               particle, neighbor,
+                                                               rho_a, rho_b)
     delta_v_diff = delta_v(particle_system, particle) -
                    delta_v(neighbor_system, neighbor)
 
-    dv[end, particle] += rho_a / rho_b * m_b * dot(delta_v_diff, grad_kernel)
-
-    return dv
+    return delta_v_diff
 end
 
 function update_shifting!(system, shifting::TransportVelocityAdami, v, u, v_ode,