Improve GPU performance of TLSPH

src/callbacks/split_integration.jl

@@ -203,9 +203,8 @@ function kick_split!(dv_ode_split, v_ode_split, u_ode_split, p, t)
                                   v_ode_split, u_ode_split, semi_split)
     end
 
-    @trixi_timeit timer() "source terms" add_source_terms!(dv_ode_split, v_ode_split,
-                                                           u_ode_split, semi, t;
-                                                           semi_wrap=semi_split)
+    add_source_terms!(dv_ode_split, v_ode_split, u_ode_split, semi, t;
+                      semi_wrap=semi_split)
 end
 
 function drift_split!(du_ode, v_ode, u_ode, p, t)

src/general/semidiscretization.jl

@@ -416,60 +416,69 @@ end
 
 function drift!(du_ode, v_ode, u_ode, semi, t)
     @trixi_timeit timer() "drift!" begin
-        @trixi_timeit timer() "reset ∂u/∂t" set_zero!(du_ode)
-
-        @trixi_timeit timer() "velocity" begin
-            # Set velocity and add acceleration for each system
-            foreach_system_wrapped(semi, v_ode, u_ode) do system, v, u
-                du = wrap_u(du_ode, system, semi)
-
-                integrate_tlsph = semi.integrate_tlsph[]
-                @threaded semi for particle in each_integrated_particle(system)
-                    # This can be dispatched per system
-                    add_velocity!(du, v, u, particle, system, integrate_tlsph, t)
-                end
-            end
+        foreach_system(semi) do system
+            du = wrap_u(du_ode, system, semi)
+            v = wrap_v(v_ode, system, semi)
+            u = wrap_u(u_ode, system, semi)
+
+            set_velocity!(du, v, u, system, semi, t)
         end
     end
 
     return du_ode
 end
 
-@inline function add_velocity!(du, v, u, particle, system, integrate_tlsph, t)
-    add_velocity!(du, v, u, particle, system, t)
-end
-
-@inline function add_velocity!(du, v, u, particle, system::TotalLagrangianSPHSystem,
-                               integrate_tlsph, t)
-    # Only add velocity for TLSPH systems if they are integrated
-    if integrate_tlsph
-        add_velocity!(du, v, u, particle, system, t)
-    end
+# Generic fallback for all systems that don't define this function
+function set_velocity!(du, v, u, system, semi, t)
+    set_velocity_default!(du, v, u, system, semi, t)
 end
 
-@inline function add_velocity!(du, v, u, particle, system, t)
-    # Generic fallback for all systems that don't define this function
-    for i in 1:ndims(system)
-        @inbounds du[i, particle] = v[i, particle]
+# Only set velocity for TLSPH systems if they are integrated
+function set_velocity!(du, v, u, system::TotalLagrangianSPHSystem, semi, t)
+    if semi.integrate_tlsph[]
+        set_velocity_default!(du, v, u, system, semi, t)
+    else
+        set_zero!(du)
     end
 
     return du
 end
 
 # Solid wall boundary system doesn't integrate the particle positions
-@inline add_velocity!(du, v, u, particle, system::WallBoundarySystem, t) = du
+function set_velocity!(du, v, u, system::WallBoundarySystem, semi, t)
+    # Note that `du` is of length zero, so we don't have to set it to zero
+    return du
+end
 
-@inline function add_velocity!(du, v, u, particle, system::AbstractFluidSystem, t)
-    # This is zero unless a shifting technique is used
-    delta_v_ = delta_v(system, particle)
+# Fluid systems integrate the particle positions and can have a shifting velocity
+function set_velocity!(du, v, u, system::AbstractFluidSystem, semi, t)
+    @threaded semi for particle in each_integrated_particle(system)
+        delta_v_ = @inbounds delta_v(system, particle)
 
-    for i in 1:ndims(system)
-        @inbounds du[i, particle] = v[i, particle] + delta_v_[i]
+        for i in 1:ndims(system)
+            @inbounds du[i, particle] = v[i, particle] + delta_v_[i]
+        end
     end
 
     return du
 end
 
+function set_velocity_default!(du, v, u, system, semi, t)
+    @threaded semi for particle in each_integrated_particle(system)
+        for i in 1:ndims(system)
+            @inbounds du[i, particle] = v[i, particle]
+        end
+    end
+
+    return du
+end
+
+# This defaults to optimized GPU copy that is about 4x faster than the threaded version above
+function set_velocity_default!(du::AbstractGPUArray, v, u, system, semi, t)
+    indices = CartesianIndices(du)
+    copyto!(du, indices, v, indices)
+end
+
 function kick!(dv_ode, v_ode, u_ode, semi, t)
     @trixi_timeit timer() "kick!" begin
         # Check that the `UpdateCallback` is used if required
@@ -483,8 +492,7 @@ function kick!(dv_ode, v_ode, u_ode, semi, t)
         @trixi_timeit timer() "system interaction" system_interaction!(dv_ode, v_ode, u_ode,
                                                                        semi)
 
-        @trixi_timeit timer() "source terms" add_source_terms!(dv_ode, v_ode, u_ode,
-                                                               semi, t)
+        add_source_terms!(dv_ode, v_ode, u_ode, semi, t)
     end
 
     return dv_ode
@@ -541,6 +549,7 @@ end
 
 # The `SplitIntegrationCallback` overwrites `semi_wrap` to use a different
 # semidiscretization for wrapping arrays.
+# `semi_wrap` is the small semidiscretization, `semi` is the large semidiscretization.
 # TODO `semi` is not used yet, but will be used when the source terms API is modified
 # to match the custom quantities API.
 function add_source_terms!(dv_ode, v_ode, u_ode, semi, t; semi_wrap=semi)
@@ -551,54 +560,63 @@ function add_source_terms!(dv_ode, v_ode, u_ode, semi, t; semi_wrap=semi)
 
         # `integrate_tlsph` is extracted from the `semi_wrap`, so that this function
         # can be used in the `SplitIntegrationCallback` as well.
-        integrate_tlsph = semi_wrap.integrate_tlsph[]
-
-        @threaded semi for particle in each_integrated_particle(system)
-            # Dispatch by system type to exclude boundary systems
-            add_acceleration!(dv, particle, system, integrate_tlsph)
-            add_source_terms_inner!(dv, v, u, particle, system, source_terms(system), t,
-                                    integrate_tlsph)
-        end
+        # In this case, `semi_wrap` will be the small sub-integration semidiscretization.
+        add_source_terms!(dv, v, u, system, semi, t, semi_wrap.integrate_tlsph[])
     end
 
     return dv_ode
 end
 
-@inline source_terms(system) = nothing
-@inline source_terms(system::Union{AbstractFluidSystem, AbstractStructureSystem}) = system.source_terms
-
-@inline function add_acceleration!(dv, particle, system, integrate_tlsph)
-    add_acceleration!(dv, particle, system)
+# This is a no-op by default but can be dispatched by system type
+function add_source_terms!(dv, v, u, system, semi, t, integrate_tlsph)
+    return dv
 end
 
-@inline function add_acceleration!(dv, particle, system::TotalLagrangianSPHSystem,
-                                   integrate_tlsph)
-    integrate_tlsph && add_acceleration!(dv, particle, system)
+function add_source_terms!(dv, v, u,
+                           system::Union{AbstractFluidSystem, AbstractStructureSystem},
+                           semi, t, integrate_tlsph)
+    add_source_terms_inner!(dv, v, u, system, semi, t)
 end
 
-@inline add_acceleration!(dv, particle, system) = dv
+function add_source_terms!(dv, v, u, system::TotalLagrangianSPHSystem,
+                           semi, t, integrate_tlsph)
+    if integrate_tlsph
+        add_source_terms_inner!(dv, v, u, system, semi, t)
+    end
 
-@propagate_inbounds function add_acceleration!(dv, particle,
-                                               system::Union{AbstractFluidSystem,
-                                                             AbstractStructureSystem})
-    (; acceleration) = system
+    return dv
+end
 
-    for i in 1:ndims(system)
-        dv[i, particle] += acceleration[i]
+function add_source_terms_inner!(dv, v, u,
+                                 system::Union{AbstractFluidSystem,
+                                               AbstractStructureSystem},
+                                 semi, t)
+    if iszero(system.acceleration) && isnothing(source_terms(system))
+        # Nothing to do
+        return dv
+    end
+
+    @trixi_timeit timer() "source terms" begin
+        @threaded semi for particle in each_integrated_particle(system)
+            add_acceleration!(dv, system, particle)
+            add_source_terms_inner!(dv, v, u, particle, system, source_terms(system), t)
+        end
     end
 
     return dv
 end
 
-@inline function add_source_terms_inner!(dv, v, u, particle, system, source_terms_, t,
-                                         integrate_tlsph)
-    add_source_terms_inner!(dv, v, u, particle, system, source_terms_, t)
-end
+@inline source_terms(system) = nothing
+@inline source_terms(system::Union{AbstractFluidSystem, AbstractStructureSystem}) = system.source_terms
 
-@inline function add_source_terms_inner!(dv, v, u, particle,
-                                         system::TotalLagrangianSPHSystem,
-                                         source_terms_, t, integrate_tlsph)
-    integrate_tlsph && add_source_terms_inner!(dv, v, u, particle, system, source_terms_, t)
+@inline function add_acceleration!(dv, system, particle)
+    (; acceleration) = system
+
+    for i in 1:ndims(system)
+        @inbounds dv[i, particle] += acceleration[i]
+    end
+
+    return dv
 end
 
 @propagate_inbounds function add_source_terms_inner!(dv, v, u, particle,

src/preprocessing/particle_packing/system.jl

@@ -259,7 +259,13 @@ function kinetic_energy(system::ParticlePackingSystem, v_ode, u_ode, semi, t)
 end
 
 @inline source_terms(system::ParticlePackingSystem) = nothing
-@inline add_acceleration!(dv, particle, system::ParticlePackingSystem) = dv
+
+# Special case for `ParticlePackingSystem`, which is an `AbstractFluidSystem`
+# but doesn't have source terms or gravity/acceleration.
+function add_source_terms!(dv, v, u, system::ParticlePackingSystem,
+                           semi, t, integrate_tlsph)
+    return dv
+end
 
 # Update from `UpdateCallback` (between time steps)
 update_particle_packing(system, v_ode, u_ode, semi, integrator) = system
@@ -390,15 +396,19 @@ end
 end
 
 # Skip for fixed systems
-@inline function add_velocity!(du, v, u, particle,
-                               system::ParticlePackingSystem{<:Any, true}, t)
+@inline function set_velocity!(du, v, u,
+                               system::ParticlePackingSystem{<:Any, true},
+                               semi, t)
+    # Note that `du` is of size `(0, n_particles)`, so we don't have to set it to zero
     return du
 end
 
 # Add advection velocity
-@inline function add_velocity!(du, v, u, particle, system::ParticlePackingSystem, t)
-    for i in 1:ndims(system)
-        du[i, particle] = system.advection_velocity[i, particle]
+@inline function set_velocity!(du, v, u, system::ParticlePackingSystem, semi, t)
+    @threaded semi for particle in each_integrated_particle(system)
+        for i in 1:ndims(system)
+            @inbounds du[i, particle] = system.advection_velocity[i, particle]
+        end
     end
 
     return du

src/schemes/boundary/open_boundary/system.jl

@@ -293,17 +293,20 @@ function calculate_dt(v_ode, u_ode, cfl_number, system::OpenBoundarySystem, semi
     return Inf
 end
 
-@inline function add_velocity!(du, v, u, particle, system::OpenBoundarySystem, t)
-    boundary_zone = current_boundary_zone(system, particle)
+@inline function set_velocity!(du, v, u, system::OpenBoundarySystem, semi, t)
+    @threaded semi for particle in each_integrated_particle(system)
+        boundary_zone = current_boundary_zone(system, particle)
 
-    pos = current_coords(u, system, particle)
-    v_particle = reference_velocity(boundary_zone, v, system, particle, pos, t)
+        pos = @inbounds current_coords(u, system, particle)
+        v_particle = @inbounds reference_velocity(boundary_zone, v, system,
+                                                  particle, pos, t)
 
-    # This is zero unless a shifting technique is used
-    delta_v_ = delta_v(system, particle)
+        # This is zero unless a shifting technique is used
+        delta_v_ = @inbounds delta_v(system, particle)
 
-    for i in 1:ndims(system)
-        @inbounds du[i, particle] = v_particle[i] + delta_v_[i]
+        for i in 1:ndims(system)
+            @inbounds du[i, particle] = v_particle[i] + delta_v_[i]
+        end
     end
 
     return du

src/schemes/fluid/shifting_techniques.jl

@@ -411,7 +411,9 @@ function update_shifting_from_callback!(system,
     v = wrap_v(v_ode, system, semi)
     u = wrap_u(u_ode, system, 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, <:Any, <:Any, true}`

src/schemes/structure/total_lagrangian_sph/system.jl

@@ -422,6 +422,17 @@ function update_tlsph_positions!(system::TotalLagrangianSPHSystem, u, semi)
     return system
 end
 
+# This defaults to optimized GPU copy that is about 4x faster than the threaded version above
+function update_tlsph_positions!(system::TotalLagrangianSPHSystem,
+                                 u::AbstractGPUArray, semi)
+    (; current_coordinates) = system
+
+    indices = CartesianIndices(u)
+    copyto!(current_coordinates, indices, u, indices)
+
+    return system
+end
+
 function apply_prescribed_motion!(system::TotalLagrangianSPHSystem,
                                   prescribed_motion::PrescribedMotion, semi, t)
     (; clamped_particles_moving, current_coordinates, cache) = system