Optimize neighbor list scratch usage

src/SPHCellList.jl

@@ -18,7 +18,9 @@ using ..OpenExternalPrograms
 using ..SPHKernels
 using ..SPHViscosityModels
 using ..SPHDensityDiffusionModels
-using ..SPHNeighborList: BuildNeighborCellLists!, ComputeCellNeighborCounts, ComputeCellParticleCounts, ConstructStencil, ExtractCells!, MapFloor, UpdateNeighbors!, UpdateΔx!
+using ..SPHNeighborList: BuildNeighborCellLists!, ComputeCellNeighborCounts!,
+    ComputeCellParticleCounts!, ConstructStencil, ExtractCells!, MapFloor,
+    NeighborListScratch, UpdateCellCounts!, UpdateNeighbors!, UpdateΔx!
 
 using StaticArrays
 using StructArrays: StructArray, foreachfield
@@ -710,7 +712,7 @@ using LinearAlgebra
                                       SimConstants, SimParticles, FullStencil,
                                       ParticleRanges, UniqueCells, CellDict,
                                       SortingScratchSpace,
-                                      NeighborCellLists, dρdtI, Velocityₙ⁺,
+                                      NeighborCellLists, NeighborScratch, dρdtI, Velocityₙ⁺,
                                       Positionₙ⁺, ρₙ⁺, ∇Cᵢ, ∇◌rᵢ,
                                       MotionDefinition::Union{
                                           Nothing,
@@ -761,7 +763,8 @@ using LinearAlgebra
                         @timeit SimMetaData.HourGlass "01a Actual Calculate IndexCounter" SimMetaData.IndexCounter = UpdateNeighbors!(SimParticles, SimKernel.H⁻¹, SortingScratchSpace,  ParticleRanges, UniqueCells, CellDict)
                         SimMetaData.Δx    = zero(eltype(dρdtI))
                         UniqueCellsView   = view(UniqueCells, 1:SimMetaData.IndexCounter)
-                        BuildNeighborCellLists!(NeighborCellLists, FullStencil, UniqueCellsView, ParticleRanges, CellDict)
+                        UpdateCellCounts!(NeighborScratch, ParticleRanges, SimMetaData.IndexCounter)
+                        BuildNeighborCellLists!(NeighborCellLists, FullStencil, UniqueCellsView, CellDict, NeighborScratch.CellOccupied)
                     end
                 end
 
@@ -862,6 +865,7 @@ using LinearAlgebra
         CellDict               = Dict{CartesianIndex{Dimensions}, Int}()
         FullStencil            = ConstructStencil(Val(Dimensions))
         NeighborCellLists      = [Int[] for _ in 1:length(UniqueCells)]
+        NeighborScratch        = NeighborListScratch()
         _, SortingScratchSpace = Base.Sort.make_scratch(nothing, eltype(SimParticles), NumberOfPoints)
 
         output = SetupVTKOutput(SimMetaData, SimParticles, SimKernel, Dimensions)
@@ -874,15 +878,15 @@ using LinearAlgebra
             cell_particle_counts = nothing
             cell_neighbor_counts = nothing
             if SimMetaData.ExportGridCellParticleCounts
-                cell_particle_counts = ComputeCellParticleCounts(
-                    ParticleRanges,
-                    SimMetaData.IndexCounter,
-                )
-                cell_neighbor_counts = ComputeCellNeighborCounts(
-                    ParticleRanges,
+                UpdateCellCounts!(NeighborScratch, ParticleRanges, SimMetaData.IndexCounter)
+                ComputeCellNeighborCounts!(
+                    NeighborScratch.NeighborCounts,
+                    NeighborScratch.CellCounts,
                     NeighborCellLists,
                     SimMetaData.IndexCounter,
                 )
+                cell_particle_counts = NeighborScratch.CellCounts
+                cell_neighbor_counts = NeighborScratch.NeighborCounts
             end
             output.enqueue_grid(
                 SimMetaData.OutputIterationCounter,
@@ -901,7 +905,7 @@ using LinearAlgebra
                 SimDensityDiffusion, SimViscosity, SimKernel, SimMetaData,
                 SimConstants, SimParticles, FullStencil, ParticleRanges,
                 UniqueCells, CellDict, SortingScratchSpace,
-                NeighborCellLists, dρdtI, Velocityₙ⁺, Positionₙ⁺, ρₙ⁺,
+                NeighborCellLists, NeighborScratch, dρdtI, Velocityₙ⁺, Positionₙ⁺, ρₙ⁺,
                 ∇Cᵢ, ∇◌rᵢ, MotionDefinition,
             )
             push!(SimMetaData.TimeSteps, SimMetaData.CurrentTimeStep)
@@ -914,15 +918,15 @@ using LinearAlgebra
             cell_particle_counts = nothing
             cell_neighbor_counts = nothing
             if SimMetaData.ExportGridCellParticleCounts
-                cell_particle_counts = ComputeCellParticleCounts(
-                    ParticleRanges,
-                    length(UniqueCellsView),
-                )
-                cell_neighbor_counts = ComputeCellNeighborCounts(
-                    ParticleRanges,
+                UpdateCellCounts!(NeighborScratch, ParticleRanges, length(UniqueCellsView))
+                ComputeCellNeighborCounts!(
+                    NeighborScratch.NeighborCounts,
+                    NeighborScratch.CellCounts,
                     NeighborCellLists,
                     length(UniqueCellsView),
                 )
+                cell_particle_counts = NeighborScratch.CellCounts
+                cell_neighbor_counts = NeighborScratch.NeighborCounts
             end
 
             @timeit SimMetaData.HourGlass "13 Save Particle Data"  begin

src/SPHExample.jl

@@ -57,7 +57,9 @@ module SPHExample
     export SimulationConstants
 
     using .SPHNeighborList
-    export ConstructStencil, ExtractCells!, UpdateNeighbors!, BuildNeighborCellLists!, ComputeCellParticleCounts, ComputeCellNeighborCounts
+    export ConstructStencil, ExtractCells!, UpdateNeighbors!, BuildNeighborCellLists!,
+           NeighborListScratch, UpdateCellCounts!, ComputeCellParticleCounts!,
+           ComputeCellNeighborCounts!, ComputeCellParticleCounts, ComputeCellNeighborCounts
 
     using .SPHCellList
     export NeighborLoop!, ComputeInteractions!, RunSimulation

src/SPHNeighborList.jl

@@ -1,14 +1,35 @@
 module SPHNeighborList
 
-export ConstructStencil, ExtractCells!, UpdateNeighbors!, BuildNeighborCellLists!, ComputeCellParticleCounts, ComputeCellNeighborCounts, UpdateΔx!
+export ConstructStencil, ExtractCells!, UpdateNeighbors!, BuildNeighborCellLists!,
+       NeighborListScratch, UpdateCellCounts!, ComputeCellParticleCounts!,
+       ComputeCellNeighborCounts!, ComputeCellParticleCounts, ComputeCellNeighborCounts,
+       UpdateΔx!
 
 using StaticArrays
+using Base.Threads
+
+mutable struct NeighborListScratch
+    CellCounts::Vector{Int}
+    CellOccupied::Vector{Bool}
+    NeighborCounts::Vector{Int}
+
+    function NeighborListScratch()
+        new(Int[], Bool[], Int[])
+    end
+end
+
+@inline function ResizeNeighborListScratch!(Scratch::NeighborListScratch, CellCount)
+    resize!(Scratch.CellCounts, CellCount)
+    resize!(Scratch.CellOccupied, CellCount)
+    resize!(Scratch.NeighborCounts, CellCount)
+    return nothing
+end
 
 function ConstructStencil(V::Val{d}) where d
     return CartesianIndices(ntuple(_ -> -1:1, V))
 end
 
-function BuildNeighborCellLists!(NeighborCellLists, FullStencil, UniqueCellsView, ParticleRanges, CellDict)
+function BuildNeighborCellLists!(NeighborCellLists, FullStencil, UniqueCellsView, CellDict, CellOccupied)
     TargetLen   = length(UniqueCellsView)
     OriginalLen = length(NeighborCellLists)
     resize!(NeighborCellLists, TargetLen)
@@ -19,17 +40,32 @@ function BuildNeighborCellLists!(NeighborCellLists, FullStencil, UniqueCellsView
         end
     end
 
-    @inbounds for CellIndex in eachindex(UniqueCellsView)
-        Neighbors = NeighborCellLists[CellIndex]
-        empty!(Neighbors)
-        Cell = UniqueCellsView[CellIndex]
-        for Offset in FullStencil
-            NeighborCell = Cell + Offset
-            NeighborIndex = get(CellDict, NeighborCell, 1)
-            StartIndex = ParticleRanges[NeighborIndex]
-            EndIndex = ParticleRanges[NeighborIndex + 1] - 1
-            if StartIndex <= EndIndex && NeighborIndex != CellIndex
-                push!(Neighbors, NeighborIndex)
+    if nthreads() > 1
+        @inbounds Threads.@threads for CellIndex in eachindex(UniqueCellsView)
+            Neighbors = NeighborCellLists[CellIndex]
+            empty!(Neighbors)
+            sizehint!(Neighbors, length(FullStencil))
+            Cell = UniqueCellsView[CellIndex]
+            for Offset in FullStencil
+                NeighborCell = Cell + Offset
+                NeighborIndex = get(CellDict, NeighborCell, 1)
+                if NeighborIndex != CellIndex && CellOccupied[NeighborIndex]
+                    push!(Neighbors, NeighborIndex)
+                end
+            end
+        end
+    else
+        @inbounds for CellIndex in eachindex(UniqueCellsView)
+            Neighbors = NeighborCellLists[CellIndex]
+            empty!(Neighbors)
+            sizehint!(Neighbors, length(FullStencil))
+            Cell = UniqueCellsView[CellIndex]
+            for Offset in FullStencil
+                NeighborCell = Cell + Offset
+                NeighborIndex = get(CellDict, NeighborCell, 1)
+                if NeighborIndex != CellIndex && CellOccupied[NeighborIndex]
+                    push!(Neighbors, NeighborIndex)
+                end
             end
         end
     end
@@ -102,24 +138,48 @@ function UpdateNeighbors!(Particles, InverseCutOff, SortingScratchSpace,
     return IndexCounter
 end
 
-function ComputeCellParticleCounts(ParticleRanges, CellCount)
-    Counts = Vector{Int}(undef, CellCount)
+function UpdateCellCounts!(Scratch::NeighborListScratch, ParticleRanges, CellCount)
+    ResizeNeighborListScratch!(Scratch, CellCount)
+    Counts = Scratch.CellCounts
+    Occupied = Scratch.CellOccupied
+    @inbounds for Index in 1:CellCount
+        Count = ParticleRanges[Index + 1] - ParticleRanges[Index]
+        Counts[Index] = Count
+        Occupied[Index] = Count > 0
+    end
+    return nothing
+end
+
+function ComputeCellParticleCounts!(Counts, ParticleRanges, CellCount)
+    resize!(Counts, CellCount)
     @inbounds for Index in 1:CellCount
         Counts[Index] = ParticleRanges[Index + 1] - ParticleRanges[Index]
     end
+    return nothing
+end
+
+function ComputeCellParticleCounts(ParticleRanges, CellCount)
+    Counts = Vector{Int}(undef, CellCount)
+    ComputeCellParticleCounts!(Counts, ParticleRanges, CellCount)
     return Counts
 end
 
-function ComputeCellNeighborCounts(ParticleRanges, NeighborCellLists, CellCount)
-    Counts = ComputeCellParticleCounts(ParticleRanges, CellCount)
-    Neighbors = Vector{Int}(undef, CellCount)
+function ComputeCellNeighborCounts!(Neighbors, Counts, NeighborCellLists, CellCount)
+    resize!(Neighbors, CellCount)
     @inbounds for Index in 1:CellCount
         NeighborTotal = 0
         for NeighborIndex in NeighborCellLists[Index]
             NeighborTotal += Counts[NeighborIndex]
         end
         Neighbors[Index] = max(Counts[Index] - 1, 0) + NeighborTotal
     end
+    return nothing
+end
+
+function ComputeCellNeighborCounts(ParticleRanges, NeighborCellLists, CellCount)
+    Counts = ComputeCellParticleCounts(ParticleRanges, CellCount)
+    Neighbors = Vector{Int}(undef, CellCount)
+    ComputeCellNeighborCounts!(Neighbors, Counts, NeighborCellLists, CellCount)
     return Neighbors
 end