Improve GPU performance of fluid-* interact

Project.toml

@@ -58,7 +58,7 @@ JSON = "1"
 KernelAbstractions = "0.9"
 OrdinaryDiffEq = "6.91"
 OrdinaryDiffEqCore = "2, 3"
-PointNeighbors = "0.6.5"
+PointNeighbors = "0.6.6"
 Polyester = "0.7.10"
 ReadVTK = "0.2"
 RecipesBase = "1"

src/general/corrections.jl

@@ -35,15 +35,16 @@ end
 # `rho_mean` is the mean density of the fluid, which is used to determine correction values near the free surface.
 #  Return a tuple `(viscosity_correction, pressure_correction, surface_tension_correction)` representing the correction terms.
 @inline function free_surface_correction(correction::AkinciFreeSurfaceCorrection,
-                                         particle_system, rho_mean)
+                                         particle_system, rho_a, rho_b)
     # Equation 4 in ref
+    rho_mean = (rho_a + rho_b) / 2
     k = correction.rho0 / rho_mean
 
     # Viscosity, pressure, surface_tension
     return k, 1, k
 end
 
-@inline function free_surface_correction(correction, particle_system, rho_mean)
+@inline function free_surface_correction(correction, particle_system, rho_a, rho_b)
     return 1, 1, 1
 end
 

src/general/neighborhood_search.jl

@@ -17,6 +17,18 @@ end
                                     neighborhood_search, particle)
 end
 
+# We cannot dispatch by `AbstractGPUArray` because this is called from within
+# a kernel, where the arrays are device arrays (like `CuDeviceArray`),
+# which are not `AbstractGPUArray`s.
+@inline function foreach_neighbor(f, system_coords, neighbor_coords, neighborhood_search,
+                                  backend::KernelAbstractions.GPU, particle)
+    # On GPUs, remove all bounds checks for maximum performance.
+    # Note that this is not safe if the neighborhood search was not initialized correctly.
+    # For example, this is unsafe when benchmarking `interact!` with the wrong NHS.
+    PointNeighbors.foreach_neighbor_unsafe(f, system_coords, neighbor_coords,
+                                           neighborhood_search, particle)
+end
+
 # === Compact support selection ===
 # -- Generic
 @inline function compact_support(system, neighbor)

src/schemes/fluid/entropically_damped_sph/rhs.jl

@@ -76,7 +76,7 @@ function interact!(dv, v_particle_system, u_particle_system,
         @inbounds dv_shifting!(dv_particle, shifting_technique(particle_system),
                                particle_system, neighbor_system,
                                v_particle_system, v_neighbor_system,
-                               particle, neighbor, m_a, m_b, rho_a, rho_b,
+                               particle, neighbor, m_a, m_b, rho_a, rho_b, v_a, v_b,
                                pos_diff, distance, grad_kernel, correction)
 
         @inbounds surface_tension_force!(dv_particle, surface_tension_a,

src/schemes/fluid/implicit_incompressible_sph/rhs.jl

@@ -47,17 +47,14 @@ function interact!(dv, v_particle_system, u_particle_system,
         m_a = @inbounds hydrodynamic_mass(particle_system, particle)
         m_b = @inbounds hydrodynamic_mass(neighbor_system, neighbor)
 
+        p_a = @inbounds current_pressure(v_particle_system, particle_system, particle)
         # The following call is equivalent to
-        #     `p_a = particle_pressure(v_particle_system, particle_system, particle)`
-        #     `p_b = particle_pressure(v_neighbor_system, neighbor_system, neighbor)`
-        # Only when the neighbor system is a `WallBoundarySystem` or a `TotalLagrangianSPHSystem`
-        # with the boundary model `PressureMirroring`, this will return `p_b = p_a`, which is
-        # the pressure of the fluid particle.
-        p_a,
-        p_b = @inbounds particle_neighbor_pressure(v_particle_system,
-                                                   v_neighbor_system,
-                                                   particle_system, neighbor_system,
-                                                   particle, neighbor)
+        #     `p_b = current_pressure(v_neighbor_system, neighbor_system, neighbor)`
+        # Only when the neighbor system is a `WallBoundarySystem`
+        # or a `TotalLagrangianSPHSystem` with the boundary model `PressureMirroring`,
+        # this will return `p_b = p_a`, which is the pressure of the fluid particle.
+        p_b = @inbounds neighbor_pressure(v_neighbor_system, neighbor_system,
+                                          neighbor, p_a)
 
         dv_pressure = pressure_acceleration(particle_system, neighbor_system,
                                             particle, neighbor,

src/schemes/fluid/shifting_techniques.jl

@@ -43,7 +43,7 @@ end
 # Additional term in the momentum equation due to the shifting technique
 @inline function dv_shifting!(dv_particle, shifting, system, neighbor_system,
                               v_system, v_neighbor_system, particle, neighbor,
-                              m_a, m_b, rho_a, rho_b, pos_diff, distance,
+                              m_a, m_b, rho_a, rho_b, v_a, v_b, pos_diff, distance,
                               grad_kernel, correction)
     return dv_particle
 end
@@ -329,27 +329,29 @@ struct MomentumEquationTermSun2019 end
                                           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)
+                                          m_a, m_b, rho_a, rho_b, v_a, v_b, pos_diff,
+                                          distance, grad_kernel, correction)
     return dv_shifting!(dv_particle, shifting.momentum_equation_term, system,
                         neighbor_system, v_system, v_neighbor_system,
-                        particle, neighbor, m_a, m_b, rho_a, rho_b,
+                        particle, neighbor, m_a, m_b, rho_a, rho_b, v_a, v_b,
                         pos_diff, distance, grad_kernel, correction)
 end
 
 @propagate_inbounds function dv_shifting!(dv_particle, ::MomentumEquationTermSun2019,
                                           system, neighbor_system,
                                           v_system, v_neighbor_system,
                                           particle, neighbor, m_a, m_b, rho_a, rho_b,
-                                          pos_diff, distance, grad_kernel, correction)
+                                          v_a, v_b, pos_diff, distance,
+                                          grad_kernel, correction)
     delta_v_a = delta_v(system, particle)
     delta_v_b = delta_v(neighbor_system, neighbor)
 
-    v_a = current_velocity(v_system, system, particle)
-    v_b = current_velocity(v_neighbor_system, neighbor_system, neighbor)
-
     tensor_product = v_a * delta_v_a' + v_b * delta_v_b'
-    dv_particle[] += m_b / rho_b *
+
+    # Since this is one of the most performance critical functions, using fast divisions
+    # here gives a significant speedup on GPUs.
+    # See the docs page "Development" for more details on `div_fast`.
+    dv_particle[] += div_fast(m_b, rho_b) *
                      (tensor_product * grad_kernel +
                       v_a * dot(delta_v_a - delta_v_b, grad_kernel))
 
@@ -388,7 +390,10 @@ struct ContinuityEquationTermSun2019 end
                                                              ::ContinuityEquationTermSun2019,
                                                              delta_v_a, delta_v_b,
                                                              rho_a, rho_b)
-    return v_diff + delta_v_a + rho_b / rho_a * delta_v_b
+    # Since this is one of the most performance critical functions, using fast divisions
+    # here gives a significant speedup on GPUs.
+    # See the docs page "Development" for more details on `div_fast`.
+    return v_diff + delta_v_a + div_fast(rho_b, rho_a) * delta_v_b
 end
 
 @inline function second_continuity_equation_term(v_diff, second_continuity_equation_term,
@@ -579,12 +584,9 @@ end
 @propagate_inbounds function dv_shifting!(dv_particle, ::TransportVelocityAdami,
                                           system, neighbor_system,
                                           v_system, v_neighbor_system, particle, neighbor,
-                                          m_a, m_b, rho_a, rho_b, pos_diff, distance,
-                                          grad_kernel, correction)
-    v_a = current_velocity(v_system, system, particle)
+                                          m_a, m_b, rho_a, rho_b, v_a, v_b, pos_diff,
+                                          distance, grad_kernel, correction)
     delta_v_a = delta_v(system, particle)
-
-    v_b = current_velocity(v_neighbor_system, neighbor_system, neighbor)
     delta_v_b = delta_v(neighbor_system, neighbor)
 
     A_a = rho_a * v_a * delta_v_a'