From cd347182d9176825b5fba3b172c6ff5852ad77e7 Mon Sep 17 00:00:00 2001 From: Andrey Prokopenko Date: Wed, 11 Jun 2025 17:26:17 -0400 Subject: [PATCH 1/2] Minor optimization in ExpandHalfToFull No need to hammer local counts with atomics, use local variable instead. --- .../detail/ArborX_ExpandHalfToFull.hpp | 27 ++++++++++++------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/src/spatial/detail/ArborX_ExpandHalfToFull.hpp b/src/spatial/detail/ArborX_ExpandHalfToFull.hpp index 3e723219c..d6e4d0e8b 100644 --- a/src/spatial/detail/ArborX_ExpandHalfToFull.hpp +++ b/src/spatial/detail/ArborX_ExpandHalfToFull.hpp @@ -34,10 +34,15 @@ void expandHalfToFull(ExecutionSpace const &space, Offsets &offsets, Kokkos::parallel_for( "ArborX::Experimental::HalfToFull::count", Kokkos::RangePolicy(space, 0, n), KOKKOS_LAMBDA(int i) { - for (int j = offsets_orig(i); j < offsets_orig(i + 1); ++j) + auto const start = offsets_orig(i); + auto const end = offsets_orig(i + 1); + if (start == end) + return; + + Kokkos::atomic_add(&offsets(i), end - start); + for (auto j = start; j < end; ++j) { - int const k = indices_orig(j); - Kokkos::atomic_inc(&offsets(i)); + auto const k = indices_orig(j); Kokkos::atomic_inc(&offsets(k)); } }); @@ -55,13 +60,17 @@ void expandHalfToFull(ExecutionSpace const &space, Offsets &offsets, typename Kokkos::TeamPolicy::member_type const &member) { auto const i = member.league_rank(); - auto const first = offsets_orig(i); - auto const last = offsets_orig(i + 1); + auto const start = offsets_orig(i); + auto const end = offsets_orig(i + 1); + if (start == end) + return; + + auto const offset = offsets(i); Kokkos::parallel_for( - Kokkos::TeamVectorRange(member, last - first), [&](int j) { - int const k = indices_orig(first + j); - indices(Kokkos::atomic_fetch_inc(&counts(i))) = k; - indices(Kokkos::atomic_fetch_inc(&counts(k))) = i; + Kokkos::TeamVectorRange(member, end - start), [&](int j) { + auto const k = indices_orig(start + j); + indices(offset + j) = k; + indices(Kokkos::atomic_dec_fetch(&counts(k + 1))) = i; }); }); Kokkos::Profiling::popRegion(); From 0613c5ffede07f3173873681fcf7659305a28566 Mon Sep 17 00:00:00 2001 From: Andrey Prokopenko Date: Mon, 3 Nov 2025 11:37:51 -0500 Subject: [PATCH 2/2] Change the order of stored neighbors to improve traversal --- src/spatial/detail/ArborX_NeighborList.hpp | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/spatial/detail/ArborX_NeighborList.hpp b/src/spatial/detail/ArborX_NeighborList.hpp index 094dc0dfd..17764fddd 100644 --- a/src/spatial/detail/ArborX_NeighborList.hpp +++ b/src/spatial/detail/ArborX_NeighborList.hpp @@ -81,8 +81,8 @@ void findHalfNeighborList(ExecutionSpace const &space, Kokkos::deep_copy(space, offsets, 0); HalfTraversal( space, bvh, - KOKKOS_LAMBDA(Value const &, Value const &value) { - Kokkos::atomic_inc(&offsets(value.index)); + KOKKOS_LAMBDA(Value const &value1, Value const &) { + Kokkos::atomic_inc(&offsets(value1.index)); }, NeighborListPredicateGetter{radius}); KokkosExt::exclusive_scan(space, offsets, offsets, 0); @@ -98,7 +98,7 @@ void findHalfNeighborList(ExecutionSpace const &space, HalfTraversal( space, bvh, KOKKOS_LAMBDA(Value const &value1, Value const &value2) { - indices(Kokkos::atomic_fetch_inc(&counts(value2.index))) = value1.index; + indices(Kokkos::atomic_fetch_inc(&counts(value1.index))) = value2.index; }, NeighborListPredicateGetter{radius}); @@ -159,7 +159,7 @@ void findFullNeighborList(ExecutionSpace const &space, HalfTraversal( space, bvh, KOKKOS_LAMBDA(Value const &value1, Value const &value2) { - indices(Kokkos::atomic_fetch_inc(&counts(value2.index))) = value1.index; + indices(Kokkos::atomic_fetch_inc(&counts(value1.index))) = value2.index; }, NeighborListPredicateGetter{radius}); @@ -174,11 +174,10 @@ void findFullNeighborList(ExecutionSpace const &space, typename Kokkos::TeamPolicy::member_type const &member) { auto const i = member.league_rank(); - auto const first = offsets(i); - auto const last = counts_copy(i); Kokkos::parallel_for( - Kokkos::TeamVectorRange(member, last - first), [&](int j) { - int const k = indices(first + j); + Kokkos::TeamVectorRange(member, offsets(i), counts_copy(i)), + [&](int j) { + int const k = indices(j); indices(Kokkos::atomic_fetch_inc(&counts(k))) = i; }); });