Use SIMD vectorization in TLSPH RHS

Project.toml

@@ -63,7 +63,7 @@ LinearAlgebra = "1"
 OrdinaryDiffEqLowStorageRK = "3"
 OrdinaryDiffEqCore = "4"
 OrdinaryDiffEqSymplecticRK = "2"
-PointNeighbors = "0.6.6"
+PointNeighbors = "0.6.7"
 Polyester = "0.7.10"
 Printf = "1"
 Random = "1"

src/general/neighborhood_search.jl

@@ -58,7 +58,7 @@ 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.
+# which are not `AbstractGPUArray`s. Hence, we need to pass the backend to this function.
 @inline function foreach_neighbor(f, system_coords, neighbor_coords, neighborhood_search,
                                   backend::KernelAbstractions.GPU, particle)
     # On GPUs, remove all bounds checks for maximum performance.
@@ -68,6 +68,47 @@ end
                                            neighborhood_search, particle)
 end
 
+@propagate_inbounds function mapreduce_neighbor(f, op, system_coords, neighbor_coords,
+                                                neighborhood_search, backend, particle;
+                                                init, simd=Val(false))
+    # When using SIMD, remove all bounds checks, as these would prevent vectorization.
+    # 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.
+    # Hence, we leave the bounds checks in the non-SIMD case where it has minimal overhead.
+    unsafe = simd
+    _mapreduce_neighbor(f, op, system_coords, neighbor_coords, neighborhood_search,
+                        backend, particle, init, unsafe, simd)
+end
+
+@propagate_inbounds function mapreduce_neighbor(f, op, system_coords, neighbor_coords,
+                                                neighborhood_search,
+                                                backend::KernelAbstractions.GPU, particle;
+                                                init, simd=Val(false))
+    # 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.
+    unsafe = Val(true)
+
+    # SIMD vectorization only works on the CPU.
+    simd_ = Val(false)
+    _mapreduce_neighbor(f, op, system_coords, neighbor_coords, neighborhood_search,
+                        backend, particle, init, unsafe, simd_)
+end
+
+@propagate_inbounds function _mapreduce_neighbor(f, op, system_coords, neighbor_coords,
+                                                 neighborhood_search, backend, particle,
+                                                 init, unsafe::Val{true}, simd)
+    PointNeighbors.mapreduce_neighbor_unsafe(f, op, system_coords, neighbor_coords,
+                                             neighborhood_search, particle; init, simd)
+end
+
+@propagate_inbounds function _mapreduce_neighbor(f, op, system_coords, neighbor_coords,
+                                                 neighborhood_search, backend, particle,
+                                                 init, unsafe::Val{false}, simd)
+    PointNeighbors.mapreduce_neighbor(f, op, system_coords, neighbor_coords,
+                                      neighborhood_search, particle; init)
+end
+
 # === Compact support selection ===
 # -- Generic
 @inline function compact_support(system, neighbor)

src/io/write_vtk.jl

@@ -351,13 +351,13 @@ function write2vtk!(vtk, v, u, t, system::AbstractFluidSystem)
             rho_b = current_density(v, system, neighbor)
             grad_kernel = smoothing_kernel_grad(system, pos_diff, distance, particle)
 
-            dv_surface_tension = Ref(zero(pos_diff))
-            surface_tension_force!(dv_surface_tension,
-                                   surface_tension_a, surface_tension_b,
-                                   system, system, particle, neighbor,
-                                   pos_diff, distance, rho_a, rho_b, grad_kernel, 1)
+            dv_surface_tension = surface_tension_force(zero(pos_diff),
+                                                       surface_tension_a, surface_tension_b,
+                                                       system, system, particle, neighbor,
+                                                       pos_diff, distance, rho_a, rho_b,
+                                                       grad_kernel, 1)
 
-            surface_tension[1:ndims(system), particle] .+= dv_surface_tension[]
+            surface_tension[1:ndims(system), particle] .+= dv_surface_tension
         end
         vtk["surface_tension"] = surface_tension
 

src/schemes/boundary/open_boundary/rhs.jl

@@ -61,17 +61,16 @@ function interact!(dv, v_particle_system, u_particle_system,
         dv_pressure_boundary = 2 * p_boundary * (m_b / (rho_a * rho_b)) * grad_kernel
 
         # Propagate `@inbounds` to the viscosity function, which accesses particle data
-        dv_viscosity_ = Ref(zero(pos_diff))
-        @inbounds dv_viscosity!(dv_viscosity_,
-                                viscosity_model(fluid_system,
-                                                neighbor_system),
-                                particle_system, neighbor_system,
-                                v_particle_system, v_neighbor_system,
-                                particle, neighbor, pos_diff, distance,
-                                sound_speed, m_a, m_b, rho_a, rho_b,
-                                v_a, v_b, grad_kernel)
-
-        dv_particle = dv_pressure + dv_viscosity_[] + dv_pressure_boundary
+        dv_viscosity_ = @inbounds dv_viscosity(zero(pos_diff),
+                                               viscosity_model(fluid_system,
+                                                               neighbor_system),
+                                               particle_system, neighbor_system,
+                                               v_particle_system, v_neighbor_system,
+                                               particle, neighbor, pos_diff, distance,
+                                               sound_speed, m_a, m_b, rho_a, rho_b,
+                                               v_a, v_b, grad_kernel)
+
+        dv_particle = dv_pressure + dv_viscosity_ + dv_pressure_boundary
 
         for i in 1:ndims(particle_system)
             @inbounds dv[i, particle] += dv_particle[i]
@@ -82,12 +81,12 @@ function interact!(dv, v_particle_system, u_particle_system,
         v_diff = v_a - v_b
 
         # Propagate `@inbounds` to the continuity equation, which accesses particle data
-        drho_particle = Ref(zero(rho_a))
-        @inbounds continuity_equation!(drho_particle,
-                                       particle_system, neighbor_system,
-                                       particle, neighbor, pos_diff, distance,
-                                       m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
-        dv[end, particle] += drho_particle[]
+        drho_particle = @inbounds continuity_equation(zero(rho_a),
+                                                      particle_system, neighbor_system,
+                                                      particle, neighbor, pos_diff,
+                                                      distance, m_b, rho_a, rho_b, v_a, v_b,
+                                                      grad_kernel)
+        dv[end, particle] += drho_particle
 
         # Open boundary pressure evolution matches the corresponding fluid system:
         # - EDAC: Compute pressure evolution like the fluid system

src/schemes/boundary/wall_boundary/rhs.jl

@@ -48,11 +48,12 @@ function interact!(dv, v_particle_system, u_particle_system,
         v_a = current_velocity(v_particle_system, particle_system, particle)
         v_b = current_velocity(v_neighbor_system, neighbor_system, neighbor)
 
-        drho_particle = Ref(zero(rho_a))
-        continuity_equation!(drho_particle, density_calculator(neighbor_system),
-                             m_b, rho_a, rho_b, v_a, v_b, grad_kernel, particle)
+        drho_particle = continuity_equation(zero(rho_a),
+                                            density_calculator(neighbor_system),
+                                            m_b, rho_a, rho_b, v_a, v_b, grad_kernel,
+                                            particle)
 
-        dv[end, particle] += drho_particle[]
+        dv[end, particle] += drho_particle
     end
 
     return dv
@@ -61,19 +62,15 @@ end
 # This is the derivative of the density summation, which is compatible with the
 # `SummationDensity` pressure acceleration.
 # Energy preservation tests will fail with the other formulation.
-@propagate_inbounds function continuity_equation!(drho_particle, ::SummationDensity,
-                                                  m_b, rho_a, rho_b, v_a, v_b,
-                                                  grad_kernel, particle)
-    drho_particle[] += m_b * dot(v_a - v_b, grad_kernel)
-
-    return drho_particle
+@propagate_inbounds function continuity_equation(drho_particle, ::SummationDensity,
+                                                 m_b, rho_a, rho_b, v_a, v_b,
+                                                 grad_kernel, particle)
+    return drho_particle + m_b * dot(v_a - v_b, grad_kernel)
 end
 
 # This is identical to the continuity equation of the fluid
-@propagate_inbounds function continuity_equation!(drho_particle, ::ContinuityDensity,
-                                                  m_b, rho_a, rho_b, v_a, v_b,
-                                                  grad_kernel, particle)
-    drho_particle[] += rho_a / rho_b * m_b * dot(v_a - v_b, grad_kernel)
-
-    return drho_particle
+@propagate_inbounds function continuity_equation(drho_particle, ::ContinuityDensity,
+                                                 m_b, rho_a, rho_b, v_a, v_b,
+                                                 grad_kernel, particle)
+    return drho_particle + rho_a / rho_b * m_b * dot(v_a - v_b, grad_kernel)
 end

src/schemes/fluid/entropically_damped_sph/rhs.jl

@@ -65,32 +65,34 @@ function interact!(dv, v_particle_system, u_particle_system,
                                             rho_b, pos_diff, distance, grad_kernel,
                                             correction)
 
-        dv_particle = Ref(dv_pressure)
-        @inbounds dv_viscosity!(dv_particle, particle_system, neighbor_system,
-                                v_particle_system, v_neighbor_system,
-                                particle, neighbor, pos_diff, distance,
-                                sound_speed, m_a, m_b, rho_a, rho_b,
-                                v_a, v_b, grad_kernel)
+        dv_particle = @inbounds dv_viscosity(dv_pressure, particle_system, neighbor_system,
+                                             v_particle_system, v_neighbor_system,
+                                             particle, neighbor, pos_diff, distance,
+                                             sound_speed, m_a, m_b, rho_a, rho_b,
+                                             v_a, v_b, grad_kernel)
 
         # Extra terms in the momentum equation when using a shifting technique
-        @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, v_a, v_b,
-                               pos_diff, distance, grad_kernel, correction)
-
-        @inbounds surface_tension_force!(dv_particle, surface_tension_a,
-                                         surface_tension_b,
-                                         particle_system, neighbor_system,
-                                         particle, neighbor, pos_diff, distance,
-                                         rho_a, rho_b, grad_kernel, 1)
-
-        @inbounds adhesion_force!(dv_particle, surface_tension_a, particle_system,
-                                  neighbor_system,
-                                  particle, neighbor, pos_diff, distance)
+        dv_particle = @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,
+                                            v_a, v_b,
+                                            pos_diff, distance, grad_kernel, correction)
+
+        dv_particle = @inbounds surface_tension_force(dv_particle, surface_tension_a,
+                                                      surface_tension_b,
+                                                      particle_system, neighbor_system,
+                                                      particle, neighbor, pos_diff,
+                                                      distance,
+                                                      rho_a, rho_b, grad_kernel, 1)
+
+        dv_particle = @inbounds adhesion_force(dv_particle, surface_tension_a,
+                                               particle_system, neighbor_system,
+                                               particle, neighbor, pos_diff, distance)
 
         for i in 1:ndims(particle_system)
-            @inbounds dv[i, particle] += dv_particle[][i]
+            @inbounds dv[i, particle] += dv_particle[i]
         end
 
         v_a = current_velocity(v_particle_system, particle_system, particle)
@@ -101,14 +103,15 @@ function interact!(dv, v_particle_system, u_particle_system,
                             particle, neighbor, pos_diff, distance,
                             sound_speed, m_a, m_b, p_a, p_b, rho_a, rho_b, nu_edac)
 
-        drho_particle = Ref(zero(rho_a))
+        drho_particle = zero(rho_a)
 
         # TODO If variable smoothing_length is used, this should use the neighbor smoothing length
         # Propagate `@inbounds` to the continuity equation, which accesses particle data
-        @inbounds continuity_equation!(drho_particle, density_calculator,
-                                       particle_system, neighbor_system,
-                                       particle, neighbor, pos_diff, distance,
-                                       m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
+        drho_particle = @inbounds continuity_equation(drho_particle, density_calculator,
+                                                      particle_system, neighbor_system,
+                                                      particle, neighbor, pos_diff,
+                                                      distance, m_b, rho_a, rho_b, v_a, v_b,
+                                                      grad_kernel)
 
         @inbounds write_drho_particle!(dv, density_calculator, drho_particle, particle)
     end

src/schemes/fluid/fluid.jl

@@ -138,29 +138,29 @@ function compute_density!(system, u, u_ode, semi, ::SummationDensity)
 end
 
 # With 'SummationDensity', density is calculated in wcsph/system.jl:compute_density!
-@inline function continuity_equation!(drho_particle, ::SummationDensity,
-                                      particle_system, neighbor_system,
-                                      particle, neighbor, pos_diff, distance,
-                                      m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
+@inline function continuity_equation(drho_particle, ::SummationDensity,
+                                     particle_system, neighbor_system,
+                                     particle, neighbor, pos_diff, distance,
+                                     m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
     return drho_particle
 end
 
 # This formulation was chosen to be consistent with the used pressure_acceleration formulations
-@propagate_inbounds function continuity_equation!(drho_particle,
-                                                  ::ContinuityDensity,
-                                                  particle_system::AbstractFluidSystem,
-                                                  neighbor_system,
-                                                  particle, neighbor, pos_diff, distance,
-                                                  m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
-    continuity_equation!(drho_particle, particle_system, neighbor_system,
-                         particle, neighbor, pos_diff, distance,
-                         m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
-end
-
-@propagate_inbounds function continuity_equation!(drho_particle,
-                                                  particle_system, neighbor_system,
-                                                  particle, neighbor, pos_diff, distance,
-                                                  m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
+@propagate_inbounds function continuity_equation(drho_particle,
+                                                 ::ContinuityDensity,
+                                                 particle_system::AbstractFluidSystem,
+                                                 neighbor_system,
+                                                 particle, neighbor, pos_diff, distance,
+                                                 m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
+    continuity_equation(drho_particle, particle_system, neighbor_system,
+                        particle, neighbor, pos_diff, distance,
+                        m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
+end
+
+@propagate_inbounds function continuity_equation(drho_particle,
+                                                 particle_system, neighbor_system,
+                                                 particle, neighbor, pos_diff, distance,
+                                                 m_b, rho_a, rho_b, v_a, v_b, grad_kernel)
     v_diff = v_a - v_b
 
     v_diff = continuity_equation_shifting_term(v_diff,
@@ -171,15 +171,16 @@ end
     # 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`.
-    drho_particle[] += div_fast(rho_a, rho_b) * m_b * dot(v_diff, grad_kernel)
+    drho_particle += div_fast(rho_a, rho_b) * m_b * dot(v_diff, grad_kernel)
 
     # Artificial density diffusion should only be applied to systems representing a fluid
     # with the same physical properties i.e. density and viscosity.
     # TODO: shouldn't be applied to particles on the interface (depends on PR #539)
     if particle_system === neighbor_system
-        density_diffusion!(drho_particle, density_diffusion(particle_system),
-                           particle_system, particle, neighbor,
-                           pos_diff, distance, m_b, rho_a, rho_b, grad_kernel)
+        drho_particle = density_diffusion(drho_particle, density_diffusion(particle_system),
+                                          particle_system, particle, neighbor,
+                                          pos_diff, distance, m_b, rho_a, rho_b,
+                                          grad_kernel)
     end
 
     return drho_particle
@@ -191,7 +192,7 @@ end
 
 @propagate_inbounds function write_drho_particle!(dv, ::ContinuityDensity,
                                                   drho_particle, particle)
-    dv[end, particle] += drho_particle[]
+    dv[end, particle] += drho_particle
 
     return dv
 end