Avoid bounds checking where it is safe to do so (trixi-framework#656)

* Avoid bounds checking where it is safe to do so

* Avoid one more bounds check in density diffusion

* Reformat code

* Use `extract_svector` from PointNeighbors.jl

* Revert 09ab7ba

* Fix tests

src/TrixiParticles.jl

@@ -3,6 +3,7 @@ module TrixiParticles
 using Reexport: @reexport
 
 using Adapt: Adapt
+using Base: @propagate_inbounds
 using CSV: CSV
 using Dates
 using DataFrames: DataFrame

src/general/density_calculators.jl

@@ -23,15 +23,15 @@ difference of the coordinates, ``v_{ab} = v_a - v_b`` of the velocities of parti
 """
 struct ContinuityDensity end
 
-@inline function particle_density(v, system, particle)
+@propagate_inbounds function particle_density(v, system, particle)
     particle_density(v, system.density_calculator, system, particle)
 end
 
-@inline function particle_density(v, ::SummationDensity, system, particle)
+@propagate_inbounds function particle_density(v, ::SummationDensity, system, particle)
     return system.cache.density[particle]
 end
 
-@inline function particle_density(v, ::ContinuityDensity, system, particle)
+@propagate_inbounds function particle_density(v, ::ContinuityDensity, system, particle)
     return v[end, particle]
 end
 

src/general/system.jl

@@ -54,28 +54,31 @@ initialize!(system, neighborhood_search) = system
 @inline active_particles(system, ::Nothing) = eachparticle(system)
 
 # This should not be dispatched by system type. We always expect to get a column of `A`.
-@inline function extract_svector(A, system, i)
+@propagate_inbounds function extract_svector(A, system, i)
     extract_svector(A, Val(ndims(system)), i)
 end
 
 # Return the `i`-th column of the array `A` as an `SVector`.
 @inline function extract_svector(A, ::Val{NDIMS}, i) where {NDIMS}
-    return SVector(ntuple(@inline(dim->A[dim, i]), NDIMS))
+    # Explicit bounds check, which can be removed by calling this function with `@inbounds`
+    @boundscheck checkbounds(A, NDIMS, i)
+
+    # Assume inbounds access now
+    return SVector(ntuple(@inline(dim->@inbounds A[dim, i]), NDIMS))
 end
 
 # Return `A[:, :, i]` as an `SMatrix`.
 @inline function extract_smatrix(A, system, particle)
     # Extract the matrix elements for this particle as a tuple to pass to SMatrix
-    return SMatrix{ndims(system), ndims(system)}(
-                                                 # Convert linear index to Cartesian index
-                                                 ntuple(@inline(i->A[mod(i - 1, ndims(system)) + 1,
-                                                                     div(i - 1, ndims(system)) + 1,
-                                                                     particle]),
-                                                        Val(ndims(system)^2)))
+    return SMatrix{ndims(system),
+                   ndims(system)}(ntuple(@inline(i->A[mod(i - 1, ndims(system)) + 1,
+                                                      div(i - 1, ndims(system)) + 1,
+                                                      particle]),
+                                         Val(ndims(system)^2)))
 end
 
 # Specifically get the current coordinates of a particle for all system types.
-@inline function current_coords(u, system, particle)
+@propagate_inbounds function current_coords(u, system, particle)
     return extract_svector(current_coordinates(u, system), system, particle)
 end
 
@@ -91,7 +94,8 @@ end
 # This can be dispatched by system type.
 @inline initial_coordinates(system) = system.initial_condition.coordinates
 
-@inline current_velocity(v, system, particle) = extract_svector(v, system, particle)
+@propagate_inbounds current_velocity(v, system, particle) = extract_svector(v, system,
+                                                                            particle)
 
 @inline function current_acceleration(system, particle)
     # TODO: Return `dv` of solid particles

src/schemes/fluid/fluid.jl

@@ -13,7 +13,7 @@ function create_cache_density(ic, ::ContinuityDensity)
     return (;)
 end
 
-@inline hydrodynamic_mass(system::FluidSystem, particle) = system.mass[particle]
+@propagate_inbounds hydrodynamic_mass(system::FluidSystem, particle) = system.mass[particle]
 
 function write_u0!(u0, system::FluidSystem)
     (; initial_condition) = system

src/schemes/fluid/viscosity.jl

@@ -1,9 +1,9 @@
 
 # Unpack the neighboring systems viscosity to dispatch on the viscosity type
-@inline function dv_viscosity(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, grad_kernel)
+@propagate_inbounds function dv_viscosity(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, grad_kernel)
     viscosity = viscosity_model(particle_system, neighbor_system)
 
     return dv_viscosity(viscosity, particle_system, neighbor_system,
@@ -12,10 +12,10 @@
                         sound_speed, m_a, m_b, rho_a, rho_b, grad_kernel)
 end
 
-@inline function dv_viscosity(viscosity, 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, grad_kernel)
+@propagate_inbounds function dv_viscosity(viscosity, 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, grad_kernel)
     return viscosity(particle_system, neighbor_system,
                      v_particle_system, v_neighbor_system,
                      particle, neighbor, pos_diff, distance,
@@ -105,12 +105,16 @@ function kinematic_viscosity(system, viscosity::ViscosityMorris)
     return viscosity.nu
 end
 
-@inline function (viscosity::Union{ArtificialViscosityMonaghan,
-                                   ViscosityMorris})(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, grad_kernel)
+@propagate_inbounds function (viscosity::Union{ArtificialViscosityMonaghan,
+                                               ViscosityMorris})(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,
+                                                                 grad_kernel)
     (; smoothing_length) = particle_system
 
     rho_mean = 0.5 * (rho_a + rho_b)
@@ -250,4 +254,6 @@ function kinematic_viscosity(system, viscosity::ViscosityAdami)
     return viscosity.nu
 end
 
-@inline viscous_velocity(v, system, particle) = current_velocity(v, system, particle)
+@propagate_inbounds function viscous_velocity(v, system, particle)
+    return current_velocity(v, system, particle)
+end

src/schemes/fluid/weakly_compressible_sph/density_diffusion.jl

@@ -199,10 +199,10 @@ function update!(density_diffusion::DensityDiffusionAntuono, neighborhood_search
     return density_diffusion
 end
 
-@inline function density_diffusion!(dv, density_diffusion::DensityDiffusion,
-                                    v_particle_system, particle, neighbor, pos_diff,
-                                    distance, m_b, rho_a, rho_b,
-                                    particle_system::FluidSystem, grad_kernel)
+@propagate_inbounds function density_diffusion!(dv, density_diffusion::DensityDiffusion,
+                                                v_particle_system, particle, neighbor,
+                                                pos_diff, distance, m_b, rho_a, rho_b,
+                                                particle_system::FluidSystem, grad_kernel)
     # Density diffusion terms are all zero for distance zero
     distance < sqrt(eps()) && return
 

src/schemes/fluid/weakly_compressible_sph/rhs.jl

@@ -23,40 +23,47 @@ function interact!(dv, v_particle_system, u_particle_system,
     foreach_point_neighbor(particle_system, neighbor_system,
                            system_coords, neighbor_system_coords,
                            neighborhood_search) do particle, neighbor, pos_diff, distance
-        rho_a = particle_density(v_particle_system, particle_system, particle)
-        rho_b = particle_density(v_neighbor_system, neighbor_system, neighbor)
+        # `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.
+        rho_a = @inbounds particle_density(v_particle_system, particle_system, particle)
+        rho_b = @inbounds particle_density(v_neighbor_system, neighbor_system, neighbor)
         rho_mean = 0.5 * (rho_a + rho_b)
 
-        # Determine correction values
-        viscosity_correction, pressure_correction, surface_tension_correction = free_surface_correction(correction,
-                                                                                                        particle_system,
-                                                                                                        rho_mean)
+        # Determine correction factors.
+        # This can be ignored, as these are all 1 when no correction is used.
+        (viscosity_correction, pressure_correction,
+        surface_tension_correction) = free_surface_correction(correction, particle_system,
+                                                              rho_mean)
 
         grad_kernel = smoothing_kernel_grad(particle_system, pos_diff, distance, particle)
 
-        m_a = hydrodynamic_mass(particle_system, particle)
-        m_b = hydrodynamic_mass(neighbor_system, neighbor)
+        m_a = @inbounds hydrodynamic_mass(particle_system, particle)
+        m_b = @inbounds hydrodynamic_mass(neighbor_system, neighbor)
 
         # 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 `BoundarySPHSystem` 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 = particle_neighbor_pressure(v_particle_system, v_neighbor_system,
-                                              particle_system, neighbor_system,
-                                              particle, neighbor)
+        p_a, p_b = @inbounds particle_neighbor_pressure(v_particle_system,
+                                                        v_neighbor_system,
+                                                        particle_system, neighbor_system,
+                                                        particle, neighbor)
 
         dv_pressure = pressure_correction *
                       pressure_acceleration(particle_system, neighbor_system, neighbor,
                                             m_a, m_b, p_a, p_b, rho_a, rho_b, pos_diff,
                                             distance, grad_kernel, correction)
 
+        # Propagate `@inbounds` to the viscosity function, which accesses particle data
         dv_viscosity_ = viscosity_correction *
-                        dv_viscosity(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, grad_kernel)
+                        @inbounds dv_viscosity(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,
+                                               grad_kernel)
 
         dv_surface_tension = surface_tension_correction *
                              surface_tension_force(surface_tension_a, surface_tension_b,
@@ -66,17 +73,19 @@ function interact!(dv, v_particle_system, u_particle_system,
         dv_adhesion = adhesion_force(surface_tension, particle_system, neighbor_system,
                                      particle, neighbor, pos_diff, distance)
 
-        @inbounds for i in 1:ndims(particle_system)
-            dv[i, particle] += dv_pressure[i] + dv_viscosity_[i] + dv_surface_tension[i] +
-                               dv_adhesion[i]
+        for i in 1:ndims(particle_system)
+            @inbounds dv[i, particle] += dv_pressure[i] + dv_viscosity_[i] +
+                                         dv_surface_tension[i] + dv_adhesion[i]
             # Debug example
             # debug_array[i, particle] += dv_pressure[i]
         end
 
         # TODO If variable smoothing_length is used, this should use the neighbor smoothing length
-        continuity_equation!(dv, density_calculator, v_particle_system, v_neighbor_system,
-                             particle, neighbor, pos_diff, distance, m_b, rho_a, rho_b,
-                             particle_system, neighbor_system, grad_kernel)
+        # Propagate `@inbounds` to the continuity equation, which accesses particle data
+        @inbounds continuity_equation!(dv, density_calculator, v_particle_system,
+                                       v_neighbor_system, particle, neighbor,
+                                       pos_diff, distance, m_b, rho_a, rho_b,
+                                       particle_system, neighbor_system, grad_kernel)
     end
     # Debug example
     # periodic_box = neighborhood_search.periodic_box
@@ -98,12 +107,12 @@ end
 end
 
 # This formulation was chosen to be consistent with the used pressure_acceleration formulations.
-@inline function continuity_equation!(dv, density_calculator::ContinuityDensity,
-                                      v_particle_system, v_neighbor_system,
-                                      particle, neighbor, pos_diff, distance,
-                                      m_b, rho_a, rho_b,
-                                      particle_system::WeaklyCompressibleSPHSystem,
-                                      neighbor_system, grad_kernel)
+@propagate_inbounds function continuity_equation!(dv, density_calculator::ContinuityDensity,
+                                                  v_particle_system, v_neighbor_system,
+                                                  particle, neighbor, pos_diff, distance,
+                                                  m_b, rho_a, rho_b,
+                                                  particle_system::WeaklyCompressibleSPHSystem,
+                                                  neighbor_system, grad_kernel)
     (; density_diffusion) = particle_system
 
     vdiff = current_velocity(v_particle_system, particle_system, particle) -
@@ -121,9 +130,10 @@ end
     end
 end
 
-@inline function particle_neighbor_pressure(v_particle_system, v_neighbor_system,
-                                            particle_system, neighbor_system,
-                                            particle, neighbor)
+@propagate_inbounds function particle_neighbor_pressure(v_particle_system,
+                                                        v_neighbor_system,
+                                                        particle_system, neighbor_system,
+                                                        particle, neighbor)
     p_a = particle_pressure(v_particle_system, particle_system, particle)
     p_b = particle_pressure(v_neighbor_system, neighbor_system, neighbor)
 

src/schemes/fluid/weakly_compressible_sph/system.jl

@@ -211,7 +211,8 @@ end
     return ndims(system) + 1
 end
 
-@inline function particle_pressure(v, system::WeaklyCompressibleSPHSystem, particle)
+@propagate_inbounds function particle_pressure(v, system::WeaklyCompressibleSPHSystem,
+                                               particle)
     return system.pressure[particle]
 end
 

test/systems/solid_system.jl

@@ -131,7 +131,7 @@
                 Base.ntuple(f, ::Symbol) = ntuple(f, 2) # Make `extract_svector` work
                 function TrixiParticles.current_coords(system::Val{:mock_system_tensor},
                                                        particle)
-                    return TrixiParticles.extract_svector(current_coordinates[i], system,
+                    return TrixiParticles.extract_svector(current_coordinates[i], Val(2),
                                                           particle)
                 end