From aeeec845fe7d1c4478662508aa0e4e2a844dd4fc Mon Sep 17 00:00:00 2001 From: Andrey Prokopenko Date: Wed, 11 Jun 2025 17:26:17 -0400 Subject: [PATCH 1/3] Minor optimization in ExpandHalfToFull No need to hammer local counts with atomics, use local variable instead. --- src/spatial/detail/ArborX_ExpandHalfToFull.hpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/spatial/detail/ArborX_ExpandHalfToFull.hpp b/src/spatial/detail/ArborX_ExpandHalfToFull.hpp index 3e723219c..acb69d4d9 100644 --- a/src/spatial/detail/ArborX_ExpandHalfToFull.hpp +++ b/src/spatial/detail/ArborX_ExpandHalfToFull.hpp @@ -34,12 +34,14 @@ void expandHalfToFull(ExecutionSpace const &space, Offsets &offsets, Kokkos::parallel_for( "ArborX::Experimental::HalfToFull::count", Kokkos::RangePolicy(space, 0, n), KOKKOS_LAMBDA(int i) { + int local_update = 0; for (int j = offsets_orig(i); j < offsets_orig(i + 1); ++j) { + ++local_update; int const k = indices_orig(j); - Kokkos::atomic_inc(&offsets(i)); Kokkos::atomic_inc(&offsets(k)); } + Kokkos::atomic_add(&offsets(i), local_update); }); KokkosExt::exclusive_scan(space, offsets, offsets, 0); From 7f541f434084be9291c68f63223bf2102dff102b Mon Sep 17 00:00:00 2001 From: Andrey Prokopenko Date: Mon, 16 Jun 2025 10:40:20 -0400 Subject: [PATCH 2/3] Only conditionally update locally --- src/spatial/detail/ArborX_ExpandHalfToFull.hpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/spatial/detail/ArborX_ExpandHalfToFull.hpp b/src/spatial/detail/ArborX_ExpandHalfToFull.hpp index acb69d4d9..2cdb4feba 100644 --- a/src/spatial/detail/ArborX_ExpandHalfToFull.hpp +++ b/src/spatial/detail/ArborX_ExpandHalfToFull.hpp @@ -41,7 +41,8 @@ void expandHalfToFull(ExecutionSpace const &space, Offsets &offsets, int const k = indices_orig(j); Kokkos::atomic_inc(&offsets(k)); } - Kokkos::atomic_add(&offsets(i), local_update); + if (local_update > 0) + Kokkos::atomic_add(&offsets(i), local_update); }); KokkosExt::exclusive_scan(space, offsets, offsets, 0); From 1014b92755846131128b237a09d4bf8b165a7052 Mon Sep 17 00:00:00 2001 From: Andrey Prokopenko Date: Wed, 10 Sep 2025 18:54:42 -0400 Subject: [PATCH 3/3] A bit cleaner --- src/spatial/detail/ArborX_ExpandHalfToFull.hpp | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/spatial/detail/ArborX_ExpandHalfToFull.hpp b/src/spatial/detail/ArborX_ExpandHalfToFull.hpp index 2cdb4feba..ce64981bb 100644 --- a/src/spatial/detail/ArborX_ExpandHalfToFull.hpp +++ b/src/spatial/detail/ArborX_ExpandHalfToFull.hpp @@ -34,15 +34,17 @@ void expandHalfToFull(ExecutionSpace const &space, Offsets &offsets, Kokkos::parallel_for( "ArborX::Experimental::HalfToFull::count", Kokkos::RangePolicy(space, 0, n), KOKKOS_LAMBDA(int i) { - int local_update = 0; - for (int j = offsets_orig(i); j < offsets_orig(i + 1); ++j) + auto start = offsets_orig(i); + auto end = offsets_orig(i + 1); + if (start == end) + return; + + Kokkos::atomic_add(&offsets(i), end - start); + for (auto j = start; j < end; ++j) { - ++local_update; - int const k = indices_orig(j); + auto const k = indices_orig(j); Kokkos::atomic_inc(&offsets(k)); } - if (local_update > 0) - Kokkos::atomic_add(&offsets(i), local_update); }); KokkosExt::exclusive_scan(space, offsets, offsets, 0);