From 82350ddcc007e2b157c399e1f0454229d986e4d9 Mon Sep 17 00:00:00 2001 From: Erik Faulhaber <44124897+efaulhaber@users.noreply.github.com> Date: Wed, 8 May 2024 15:41:25 +0200 Subject: [PATCH 1/5] Implement neighborhood search based on CellListMap.jl --- Project.toml | 2 ++ src/PointNeighbors.jl | 1 + src/celllistmap_nhs.jl | 63 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 66 insertions(+) create mode 100644 src/celllistmap_nhs.jl diff --git a/Project.toml b/Project.toml index fc36226..b1cee12 100644 --- a/Project.toml +++ b/Project.toml @@ -4,12 +4,14 @@ authors = ["Erik Faulhaber "] version = "0.2.2-dev" [deps] +CellListMap = "69e1c6dd-3888-40e6-b3c8-31ac5f578864" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [compat] +CellListMap = "0.9" LinearAlgebra = "1" Polyester = "0.7.5" Reexport = "1" diff --git a/src/PointNeighbors.jl b/src/PointNeighbors.jl index 8c18400..0b7d896 100644 --- a/src/PointNeighbors.jl +++ b/src/PointNeighbors.jl @@ -10,6 +10,7 @@ include("util.jl") include("neighborhood_search.jl") include("trivial_nhs.jl") include("grid_nhs.jl") +include("celllistmap_nhs.jl") export for_particle_neighbor export TrivialNeighborhoodSearch, GridNeighborhoodSearch diff --git a/src/celllistmap_nhs.jl b/src/celllistmap_nhs.jl new file mode 100644 index 0000000..3e92726 --- /dev/null +++ b/src/celllistmap_nhs.jl @@ -0,0 +1,63 @@ +mutable struct CellListMapNeighborhoodSearch{CL, B} + cell_list :: CL + box :: B + + function CellListMapNeighborhoodSearch{NDIMS}(search_radius) where {NDIMS} + # Create a cell list with only one particle and resize it later + x = zeros(NDIMS, 1) + box = CellListMap.Box(CellListMap.limits(x, x), search_radius) + cell_list = CellListMap.CellList(x, x, box) + + return new{typeof(cell_list), typeof(box)}(cell_list, box) + end +end + +function initialize!(neighborhood_search::CellListMapNeighborhoodSearch, + x::AbstractMatrix, y::AbstractMatrix) + update!(neighborhood_search, x, y) +end + +function update!(neighborhood_search::CellListMapNeighborhoodSearch, + x::AbstractMatrix, y::AbstractMatrix; + particles_moving = (true, true)) + (; cell_list) = neighborhood_search + + # Resize box + box = CellListMap.Box(CellListMap.limits(x, y), neighborhood_search.box.cutoff) + neighborhood_search.box = box + + # Resize and update cell list + CellListMap.UpdateCellList!(x, y, box, cell_list) + + # Recalculate number of batches for multithreading + CellListMap.set_number_of_batches!(cell_list) + + return neighborhood_search +end + +# The type annotation is to make Julia specialize on the type of the function. +# Otherwise, unspecialized code will cause a lot of allocations +# and heavily impact performance. +# See https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing +function for_particle_neighbor(f::T, system_coords, neighbor_coords, + neighborhood_search::CellListMapNeighborhoodSearch; + particles = axes(system_coords, 2), + parallel = true) where {T} + (; cell_list, box) = neighborhood_search + + # `0` is the returned output, which we don't use + CellListMap.map_pairwise!(0, box, cell_list, + parallel = parallel) do x, y, i, j, d2, output + # Skip all indices not in `particles` + i in particles || return output + + pos_diff = x - y + distance = sqrt(d2) + + @inline f(i, j, pos_diff, distance) + + return output + end + + return nothing +end From f7f48e8c8e6175c7e221d550da375f2d95e1659f Mon Sep 17 00:00:00 2001 From: Erik Faulhaber <44124897+efaulhaber@users.noreply.github.com> Date: Fri, 10 May 2024 09:05:12 +0200 Subject: [PATCH 2/5] using CellListMap --- src/PointNeighbors.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/PointNeighbors.jl b/src/PointNeighbors.jl index 0b7d896..43d20d2 100644 --- a/src/PointNeighbors.jl +++ b/src/PointNeighbors.jl @@ -2,6 +2,7 @@ module PointNeighbors using Reexport: @reexport +using CellListMap: CellListMap using LinearAlgebra: dot using Polyester: @batch @reexport using StaticArrays: SVector From 357ebd252b7682d214e5431f8586d80b59dcab01 Mon Sep 17 00:00:00 2001 From: Erik Faulhaber <44124897+efaulhaber@users.noreply.github.com> Date: Thu, 4 Jul 2024 15:44:00 +0200 Subject: [PATCH 3/5] Add the new NHS as a package extension --- Project.toml | 7 ++++- src/PointNeighbors.jl | 5 ++- src/celllistmap_nhs.jl | 63 ------------------------------------- src/neighborhood_search.jl | 2 ++ test/Project.toml | 2 ++ test/neighborhood_search.jl | 3 ++ test/test_util.jl | 3 ++ 7 files changed, 18 insertions(+), 67 deletions(-) delete mode 100644 src/celllistmap_nhs.jl diff --git a/Project.toml b/Project.toml index 9f486e0..8a41d67 100644 --- a/Project.toml +++ b/Project.toml @@ -6,7 +6,6 @@ version = "0.4.1-dev" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458" -CellListMap = "69e1c6dd-3888-40e6-b3c8-31ac5f578864" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -14,6 +13,12 @@ Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +[weakdeps] +CellListMap = "69e1c6dd-3888-40e6-b3c8-31ac5f578864" + +[extensions] +PointNeighborsCellListMapExt = "CellListMap" + [compat] Adapt = "4" Atomix = "0.1" diff --git a/src/PointNeighbors.jl b/src/PointNeighbors.jl index a6b3075..8314d15 100644 --- a/src/PointNeighbors.jl +++ b/src/PointNeighbors.jl @@ -4,7 +4,6 @@ using Reexport: @reexport using Adapt: Adapt using Atomix: Atomix -using CellListMap: CellListMap using GPUArraysCore: AbstractGPUArray using KernelAbstractions: KernelAbstractions, @kernel, @index using LinearAlgebra: dot @@ -18,11 +17,11 @@ include("nhs_trivial.jl") include("cell_lists/cell_lists.jl") include("nhs_grid.jl") include("nhs_precomputed.jl") -include("celllistmap_nhs.jl") include("gpu.jl") export foreach_point_neighbor, foreach_neighbor -export TrivialNeighborhoodSearch, GridNeighborhoodSearch, PrecomputedNeighborhoodSearch +export TrivialNeighborhoodSearch, GridNeighborhoodSearch, PrecomputedNeighborhoodSearch, + CellListMapNeighborhoodSearch export DictionaryCellList, FullGridCellList export ParallelUpdate, SemiParallelUpdate, SerialUpdate export initialize!, update!, initialize_grid!, update_grid! diff --git a/src/celllistmap_nhs.jl b/src/celllistmap_nhs.jl deleted file mode 100644 index 3e92726..0000000 --- a/src/celllistmap_nhs.jl +++ /dev/null @@ -1,63 +0,0 @@ -mutable struct CellListMapNeighborhoodSearch{CL, B} - cell_list :: CL - box :: B - - function CellListMapNeighborhoodSearch{NDIMS}(search_radius) where {NDIMS} - # Create a cell list with only one particle and resize it later - x = zeros(NDIMS, 1) - box = CellListMap.Box(CellListMap.limits(x, x), search_radius) - cell_list = CellListMap.CellList(x, x, box) - - return new{typeof(cell_list), typeof(box)}(cell_list, box) - end -end - -function initialize!(neighborhood_search::CellListMapNeighborhoodSearch, - x::AbstractMatrix, y::AbstractMatrix) - update!(neighborhood_search, x, y) -end - -function update!(neighborhood_search::CellListMapNeighborhoodSearch, - x::AbstractMatrix, y::AbstractMatrix; - particles_moving = (true, true)) - (; cell_list) = neighborhood_search - - # Resize box - box = CellListMap.Box(CellListMap.limits(x, y), neighborhood_search.box.cutoff) - neighborhood_search.box = box - - # Resize and update cell list - CellListMap.UpdateCellList!(x, y, box, cell_list) - - # Recalculate number of batches for multithreading - CellListMap.set_number_of_batches!(cell_list) - - return neighborhood_search -end - -# The type annotation is to make Julia specialize on the type of the function. -# Otherwise, unspecialized code will cause a lot of allocations -# and heavily impact performance. -# See https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing -function for_particle_neighbor(f::T, system_coords, neighbor_coords, - neighborhood_search::CellListMapNeighborhoodSearch; - particles = axes(system_coords, 2), - parallel = true) where {T} - (; cell_list, box) = neighborhood_search - - # `0` is the returned output, which we don't use - CellListMap.map_pairwise!(0, box, cell_list, - parallel = parallel) do x, y, i, j, d2, output - # Skip all indices not in `particles` - i in particles || return output - - pos_diff = x - y - distance = sqrt(d2) - - @inline f(i, j, pos_diff, distance) - - return output - end - - return nothing -end diff --git a/src/neighborhood_search.jl b/src/neighborhood_search.jl index dacbc10..3299b98 100644 --- a/src/neighborhood_search.jl +++ b/src/neighborhood_search.jl @@ -236,3 +236,5 @@ end @inline function periodic_coords(coords, periodic_box::Nothing) return coords end + +CellListMapNeighborhoodSearch() = error("CellListMap.jl has to be imported to use this") diff --git a/test/Project.toml b/test/Project.toml index f57fdf2..3e2bd50 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,10 +1,12 @@ [deps] BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +CellListMap = "69e1c6dd-3888-40e6-b3c8-31ac5f578864" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] BenchmarkTools = "1" +CellListMap = "0.9" Plots = "1" Test = "1" diff --git a/test/neighborhood_search.jl b/test/neighborhood_search.jl index a49ad2b..c54d62c 100644 --- a/test/neighborhood_search.jl +++ b/test/neighborhood_search.jl @@ -178,6 +178,7 @@ search_radius, backend = Vector{Vector{Int}})), PrecomputedNeighborhoodSearch{NDIMS}(; search_radius, n_points), + CellListMapNeighborhoodSearch(NDIMS, search_radius), ] names = [ @@ -187,6 +188,7 @@ "`GridNeighborhoodSearch` with `FullGridCellList` with `DynamicVectorOfVectors` and `SemiParallelUpdate`", "`GridNeighborhoodSearch` with `FullGridCellList` with `Vector{Vector}`", "`PrecomputedNeighborhoodSearch`", + "`CellListMapNeighborhoodSearch`", ] # Also test copied templates @@ -202,6 +204,7 @@ max_corner, backend = Vector{Vector{Int32}})), PrecomputedNeighborhoodSearch{NDIMS}(), + CellListMapNeighborhoodSearch(NDIMS, 1.0), ] copied_nhs = copy_neighborhood_search.(template_nhs, search_radius, n_points) append!(neighborhood_searches, copied_nhs) diff --git a/test/test_util.jl b/test/test_util.jl index 07a3ac9..a5c1911 100644 --- a/test/test_util.jl +++ b/test/test_util.jl @@ -3,6 +3,9 @@ using Test: @test, @testset, @test_throws using PointNeighbors +# Load `PointNeighborsCellListMapExt` +import CellListMap + """ @trixi_testset "name of the testset" #= code to test #= From c8f1ce2d3e4163faab49e78a69ec57cfa93cb9e1 Mon Sep 17 00:00:00 2001 From: Erik Faulhaber <44124897+efaulhaber@users.noreply.github.com> Date: Thu, 4 Jul 2024 15:57:04 +0200 Subject: [PATCH 4/5] Add missing extension file --- ext/PointNeighborsCellListMapExt.jl | 78 +++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 ext/PointNeighborsCellListMapExt.jl diff --git a/ext/PointNeighborsCellListMapExt.jl b/ext/PointNeighborsCellListMapExt.jl new file mode 100644 index 0000000..617a5a2 --- /dev/null +++ b/ext/PointNeighborsCellListMapExt.jl @@ -0,0 +1,78 @@ +module PointNeighborsCellListMapExt + +using PointNeighbors +using CellListMap: CellListMap + +mutable struct CellListMapNeighborhoodSearch{CL, B} + cell_list :: CL + box :: B + + function PointNeighbors.CellListMapNeighborhoodSearch(NDIMS, search_radius) + # Create a cell list with only one point and resize it later + x = zeros(NDIMS, 1) + box = CellListMap.Box(CellListMap.limits(x, x), search_radius) + cell_list = CellListMap.CellList(x, x, box) + + return new{typeof(cell_list), typeof(box)}(cell_list, box) + end +end + +function PointNeighbors.initialize!(neighborhood_search::CellListMapNeighborhoodSearch, + x::AbstractMatrix, y::AbstractMatrix) + PointNeighbors.update!(neighborhood_search, x, y) +end + +function PointNeighbors.update!(neighborhood_search::CellListMapNeighborhoodSearch, + x::AbstractMatrix, y::AbstractMatrix; + points_moving = (true, true)) + (; cell_list) = neighborhood_search + + # Resize box + box = CellListMap.Box(CellListMap.limits(x, y), neighborhood_search.box.cutoff) + neighborhood_search.box = box + + # Resize and update cell list + CellListMap.UpdateCellList!(x, y, box, cell_list) + + # Recalculate number of batches for multithreading + CellListMap.set_number_of_batches!(cell_list) + + return neighborhood_search +end + +# The type annotation is to make Julia specialize on the type of the function. +# Otherwise, unspecialized code will cause a lot of allocations +# and heavily impact performance. +# See https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing +function PointNeighbors.foreach_point_neighbor(f::T, system_coords, neighbor_coords, + neighborhood_search::CellListMapNeighborhoodSearch; + points = axes(system_coords, 2), + parallel = true) where {T} + (; cell_list, box) = neighborhood_search + + # `0` is the returned output, which we don't use + CellListMap.map_pairwise!(0, box, cell_list, + parallel = parallel) do x, y, i, j, d2, output + # Skip all indices not in `points` + i in points || return output + + pos_diff = x - y + distance = sqrt(d2) + + @inline f(i, j, pos_diff, distance) + + return output + end + + return nothing +end + +function PointNeighbors.copy_neighborhood_search(nhs::CellListMapNeighborhoodSearch, + search_radius, n_points; + eachpoint = 1:n_points) + NDIMS = length(nhs.box.cell_size) + + return PointNeighbors.CellListMapNeighborhoodSearch(NDIMS, search_radius) +end + +end From 1d1c078c34326c8f46e1e8489241c6aab130a279 Mon Sep 17 00:00:00 2001 From: Erik Faulhaber <44124897+efaulhaber@users.noreply.github.com> Date: Fri, 5 Jul 2024 09:41:01 +0200 Subject: [PATCH 5/5] Add functions `search_radius` and `ndims` --- ext/PointNeighborsCellListMapExt.jl | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/ext/PointNeighborsCellListMapExt.jl b/ext/PointNeighborsCellListMapExt.jl index 617a5a2..5ccb7bc 100644 --- a/ext/PointNeighborsCellListMapExt.jl +++ b/ext/PointNeighborsCellListMapExt.jl @@ -17,6 +17,14 @@ mutable struct CellListMapNeighborhoodSearch{CL, B} end end +function PointNeighbors.search_radius(neighborhood_search::CellListMapNeighborhoodSearch) + return neighborhood_search.box.cutoff +end + +function Base.ndims(neighborhood_search::CellListMapNeighborhoodSearch) + return length(neighborhood_search.box.cell_size) +end + function PointNeighbors.initialize!(neighborhood_search::CellListMapNeighborhoodSearch, x::AbstractMatrix, y::AbstractMatrix) PointNeighbors.update!(neighborhood_search, x, y) @@ -70,9 +78,7 @@ end function PointNeighbors.copy_neighborhood_search(nhs::CellListMapNeighborhoodSearch, search_radius, n_points; eachpoint = 1:n_points) - NDIMS = length(nhs.box.cell_size) - - return PointNeighbors.CellListMapNeighborhoodSearch(NDIMS, search_radius) + return PointNeighbors.CellListMapNeighborhoodSearch(ndims(nhs), search_radius) end end