Skip to content

Use a generated function to create one loop per dimension #149

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
vchuravy wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

Use a generated function to create one loop per dimension #149

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
38bea21
Use a generated function to create one loop per dimension
vchuravy Mar 23, 2026
64d2e59
fixup! Use a generated function to create one loop per dimension
vchuravy Mar 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 18 additions & 7 deletions src/nhs_grid.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -506,20 +506,30 @@ function check_cell_collision(neighbor_cell_::CartesianIndex,
coords[hash] != PointNeighbors.coordinates_flattened(neighbor_cell)
end

using Base.Cartesian: @nloops, @ntuple

macro __unroll()
Expr(:loopinfo, (Symbol("llvm.loop.unroll.full"),))
end

# Specialized version of the function in `neighborhood_search.jl`, which is faster
# than looping over `eachneighbor`.
@inline function foreach_neighbor(f, neighbor_system_coords,
neighborhood_search::GridNeighborhoodSearch,
point, point_coords, search_radius)
@generated function foreach_neighbor(f, neighbor_system_coords,
neighborhood_search::GridNeighborhoodSearch{NDIMS},
point, point_coords, search_radius) where NDIMS
quote
@inline
(; cell_list, periodic_box) = neighborhood_search
cell = cell_coords(point_coords, neighborhood_search)

for neighbor_cell_ in neighboring_cells(cell, neighborhood_search)
neighbor_cell = Tuple(neighbor_cell_)
# Generated explicitly nested for-loops so that LLVM can actually unroll them.
@nloops $NDIMS d d -> -1:1 begin
# @nloops $NDIMS d d -> -1:1 d->nothing d->@__unroll() begin
neighbor_cell = @ntuple $NDIMS i->cell[i] + d_i
neighbors = points_in_cell(neighbor_cell, neighborhood_search)

# Boolean to indicate if this cell has a collision (only with `SpatialHashingCellList`)
cell_collision = check_cell_collision(neighbor_cell_,
cell_collision = check_cell_collision(CartesianIndex(neighbor_cell),
cell_list, neighborhood_search)

for neighbor_ in eachindex(neighbors)
Expand All @@ -543,7 +553,7 @@ end
# If this cell has a collision, check if this point belongs to this cell
# (only with `SpatialHashingCellList`).
if cell_collision &&
check_collision(neighbor_cell_, neighbor_coords, cell_list,
check_collision(CartesianIndex(neighbor_cell), neighbor_coords, cell_list,
neighborhood_search)
continue
end
Expand All @@ -554,6 +564,7 @@ end
end
end
end
end # quote
end

@inline function neighboring_cells(cell, neighborhood_search)
Expand Down
Loading