Keep support arrays aligned on neighbor rebuild

src/SPHCellList.jl

@@ -15,7 +15,7 @@ using ..OpenExternalPrograms
 using ..SPHKernels
 using ..SPHViscosityModels
 using ..SPHDensityDiffusionModels
-using ..SPHNeighborList: BuildNeighborCellLists!, ComputeCellNeighborCounts, ComputeCellParticleCounts, ConstructStencil, ExtractCells!, FindCellIndex, MapFloor, UpdateNeighbors!, UpdateΔx!
+using ..SPHNeighborList: BuildNeighborCellLists!, ComputeCellNeighborCounts, ComputeCellParticleCounts, ConstructStencil, ExtractCells!, FindCellIndex, MapFloor, UpdateNeighbors!, UpdateNeighborsAligned!, UpdateΔx!
 
 using Base.Threads: @threads
 using Bumper: @alloc, @no_escape
@@ -726,7 +726,7 @@ using TimerOutputs: @timeit
             AccelerationMax = @alloc(FloatType, length(SimParticles.Position))
             dt₂ = dt * 0.5
 
-            SimMetaData.IndexCounter = UpdateNeighbors!(SimParticles, SimKernel.H⁻¹, SortingScratchSpace, ParticleRanges, UniqueCells, CellListIndices)
+            SimMetaData.IndexCounter = UpdateNeighborsAligned!(SimParticles, SimKernel.H⁻¹, SortingScratchSpace, ParticleRanges, UniqueCells, CellListIndices, dρdtI, Velocityₙ⁺, Positionₙ⁺, ρₙ⁺, ∇Cᵢ, ∇◌rᵢ)
             UniqueCellsView = view(UniqueCells, 1:SimMetaData.IndexCounter)
             BuildNeighborCellLists!(NeighborCellLists, FullStencil, UniqueCellsView, ParticleRanges)
 
@@ -756,7 +756,7 @@ using TimerOutputs: @timeit
                     # Remove if statement logic if you want to update each iteration
                     # if mod(SimMetaData.Iteration, ceil(Int, SimKernel.H / (SimConstants.c₀ * dt * (1/SimConstants.CFL)) )) == 0 || SimMetaData.Iteration == 1
                     if ShouldRebuild
-                        @timeit SimMetaData.HourGlass "01a Actual Calculate IndexCounter" SimMetaData.IndexCounter = UpdateNeighbors!(SimParticles, SimKernel.H⁻¹, SortingScratchSpace,  ParticleRanges, UniqueCells, CellListIndices)
+                        @timeit SimMetaData.HourGlass "01a Actual Calculate IndexCounter" SimMetaData.IndexCounter = UpdateNeighborsAligned!(SimParticles, SimKernel.H⁻¹, SortingScratchSpace, ParticleRanges, UniqueCells, CellListIndices, dρdtI, Velocityₙ⁺, Positionₙ⁺, ρₙ⁺, ∇Cᵢ, ∇◌rᵢ)
                         SimMetaData.Δx    = zero(eltype(dρdtI))
                         UniqueCellsView   = view(UniqueCells, 1:SimMetaData.IndexCounter)
                         BuildNeighborCellLists!(NeighborCellLists, FullStencil, UniqueCellsView, ParticleRanges)

src/SPHExample.jl

@@ -58,7 +58,7 @@ module SPHExample
     export SimulationConstants
 
     using .SPHNeighborList
-    export ConstructStencil, ExtractCells!, UpdateNeighbors!, BuildNeighborCellLists!, ComputeCellParticleCounts, ComputeCellNeighborCounts
+    export ConstructStencil, ExtractCells!, UpdateNeighbors!, UpdateNeighborsAligned!, BuildNeighborCellLists!, ComputeCellParticleCounts, ComputeCellNeighborCounts
 
     using .SPHCellList
     export NeighborLoop!, ComputeInteractions!, RunSimulation

src/SPHNeighborList.jl

@@ -1,6 +1,6 @@
 module SPHNeighborList
 
-export ConstructStencil, ExtractCells!, UpdateNeighbors!, BuildNeighborCellLists!, ComputeCellParticleCounts, ComputeCellNeighborCounts, UpdateΔx!, FindCellIndex
+export ConstructStencil, ExtractCells!, UpdateNeighbors!, UpdateNeighborsAligned!, BuildNeighborCellLists!, ComputeCellParticleCounts, ComputeCellNeighborCounts, UpdateΔx!, FindCellIndex
 
 using StaticArrays
 
@@ -92,6 +92,16 @@ function UpdateNeighbors!(Particles, InverseCutOff, SortingScratchSpace,
     ExtractCells!(Particles, InverseCutOff)
 
     sort!(Particles, by = p -> p.Cells; scratch=SortingScratchSpace)
+    return RebuildCellRanges!(Particles, ParticleRanges, UniqueCells, CellListIndices)
+end
+
+@inline function ApplyPermutation!(Array, Permutation)
+    Copy = Array[Permutation]
+    copyto!(Array, Copy)
+    return nothing
+end
+
+function RebuildCellRanges!(Particles, ParticleRanges, UniqueCells, CellListIndices)
     Cells = @views Particles.Cells
     UniqueCells[1] = MinCell(eltype(Cells))
     @. ParticleRanges             = zero(eltype(ParticleRanges))
@@ -114,6 +124,32 @@ function UpdateNeighbors!(Particles, InverseCutOff, SortingScratchSpace,
     return IndexCounter
 end
 
+"""
+    UpdateNeighborsAligned!(Particles, InverseCutOff, SortingScratchSpace,
+                            ParticleRanges, UniqueCells, CellListIndices,
+                            AlignedArrays...)
+
+Update particle cell assignments, sort `Particles` by cell, and apply the same
+particle permutation to every array in `AlignedArrays`. Use this rebuild path
+when separate per-particle work arrays must remain indexed with `Particles`
+after the neighbor list is rebuilt.
+"""
+function UpdateNeighborsAligned!(Particles, InverseCutOff, SortingScratchSpace,
+                                 ParticleRanges, UniqueCells, CellListIndices,
+                                 AlignedArrays...)
+    ExtractCells!(Particles, InverseCutOff)
+
+    Permutation = sortperm(eachindex(Particles), by = Index -> Particles.Cells[Index])
+    ApplyPermutation!(Particles, Permutation)
+    for Array in AlignedArrays
+        if length(Array) == length(Particles)
+            ApplyPermutation!(Array, Permutation)
+        end
+    end
+
+    return RebuildCellRanges!(Particles, ParticleRanges, UniqueCells, CellListIndices)
+end
+
 function ComputeCellParticleCounts(ParticleRanges, CellCount)
     Counts = Vector{Int}(undef, CellCount)
     @inbounds for Index in 1:CellCount

test/runtests.jl

@@ -2,19 +2,48 @@ using Test
 using SPHExample
 using StaticArrays
 using StructArrays
+using LinearAlgebra
 
 @testset "time stepping" begin
     pos = [SVector{2,Float64}(0.0, 0.0), SVector{2,Float64}(1.0, 0.0)]
     vel = [SVector{2,Float64}(0.0, 0.0), SVector{2,Float64}(0.0, 0.0)]
     acc = [SVector{2,Float64}(0.0, 0.0), SVector{2,Float64}(0.0, -9.81)]
     sc = SimulationConstants{Float64}()
     ker = SPHKernelInstance{2, Float64}(WendlandC2(); dx=sc.dx)
-    dt  = Δt(pos, vel, acc, sc, ker)
+    dt  = Δt(maximum(norm, acc), sc, ker)
     @test dt > 0
-    alloc = @allocated Δt(pos, vel, acc, sc, ker)
+    alloc = @allocated Δt(maximum(norm, acc), sc, ker)
     @test alloc == 0
 end
 
+@testset "aligned neighbor rebuild" begin
+    D = 2
+    T = Float64
+    positions = [
+        SVector{D,T}(2.0, 0.0),
+        SVector{D,T}(0.0, 0.0),
+        SVector{D,T}(1.0, 0.0),
+    ]
+    particles = StructArray((
+        Cells = fill(CartesianIndex(0, 0), 3),
+        Position = copy(positions),
+        ID = [20, 0, 10],
+    ))
+    aligned = ["id-20", "id-0", "id-10"]
+    particle_ranges = zeros(Int, length(particles) + 2)
+    unique_cells = zeros(CartesianIndex{D}, length(particles) + 1)
+    cell_list_indices = zeros(Int, length(particles))
+    _, sorting_scratch = Base.Sort.make_scratch(nothing, eltype(particles), length(particles))
+
+    SPHExample.SPHNeighborList.UpdateNeighborsAligned!(
+        particles, one(T), sorting_scratch, particle_ranges, unique_cells,
+        cell_list_indices, aligned,
+    )
+
+    @test particles.ID == [0, 10, 20]
+    @test aligned == ["id-0", "id-10", "id-20"]
+end
+
 @testset "isolated particle" begin
     D = 2
     T = Float64
@@ -51,15 +80,14 @@ end
 
     for _ in 1:1000
         ResetArrays!(dρdtI, particles.Acceleration)
-        dt = Δt(particles.Position, particles.Velocity, particles.Acceleration,
-                 sc, ker)
+        dt = Δt(maximum(norm, particles.Acceleration), sc, ker)
         dt2 = dt / 2
 
         SPHExample.SPHCellList.HalfTimeStep(meta, sc, particles, pos_n, vel_n,
                                            ρ_n, dρdtI, dt2)
         LimitDensityAtBoundary!(ρ_n, sc.ρ₀, particles.MotionLimiter)
         Pressure!(press, ρ_n, sc)
-        SPHExample.SPHCellList.FullTimeStep(meta, ker, sc, particles, ∇C, ∇r, dt)
+        SPHExample.SPHCellList.FullTimeStep(meta, ker, sc, particles, vel_n, ∇C, ∇r, dt)
         DensityEpsi!(dens, dρdtI, ρ_n, dt)
         LimitDensityAtBoundary!(dens, sc.ρ₀, particles.MotionLimiter)
         SPHExample.SPHCellList.UpdateMetaData!(meta, dt)