Implement expand half to full (neighbor list)

src/details/ArborX_DetailsExpandHalfToFull.hpp

@@ -0,0 +1,67 @@
+/****************************************************************************
+ * Copyright (c) 2017-2023 by the ArborX authors                            *
+ * All rights reserved.                                                     *
+ *                                                                          *
+ * This file is part of the ArborX library. ArborX is                       *
+ * distributed under a BSD 3-clause license. For the licensing terms see    *
+ * the LICENSE file in the top-level directory.                             *
+ *                                                                          *
+ * SPDX-License-Identifier: BSD-3-Clause                                    *
+ ****************************************************************************/
+
+#ifndef ARBORX_DETAILS_EXPAND_HALF_TO_FULL_HPP
+#define ARBORX_DETAILS_EXPAND_HALF_TO_FULL_HPP
+
+#include <ArborX_DetailsKokkosExtViewHelpers.hpp>
+#include <ArborX_DetailsUtils.hpp>
+
+#include <Kokkos_Core.hpp>
+
+namespace ArborX::Details
+{
+
+template <class ExecutionSpace, class Offsets, class Indices>
+void expandHalfToFull(ExecutionSpace const &space, Offsets &offsets,
+                      Indices &indices)
+{
+  Kokkos::Profiling::pushRegion("ArborX::Experimental::HalfToFull");
+  typename Offsets::const_type const offsets_orig = offsets;
+  typename Indices::const_type const indices_orig = indices;
+
+  auto const n = offsets.extent(0) - 1;
+  offsets = KokkosExt::cloneWithoutInitializingNorCopying(space, offsets_orig);
+  Kokkos::deep_copy(space, offsets, 0);
+  Kokkos::parallel_for(
+      "ArborX::Experimental::HalfToFull::count",
+      Kokkos::RangePolicy<ExecutionSpace>(space, 0, n), KOKKOS_LAMBDA(int i) {
+        for (int j = offsets_orig(i); j < offsets_orig(i + 1); ++j)
+        {
+          int const k = indices_orig(j);
+          Kokkos::atomic_increment(&offsets(i));
+          Kokkos::atomic_increment(&offsets(k));
+        }
+      });
+  exclusivePrefixSum(space, offsets);
+
+  auto const m = KokkosExt::lastElement(space, offsets);
+  KokkosExt::reallocWithoutInitializing(space, indices, m);
+
+  Offsets counts(
+      Kokkos::view_alloc(space, "ArborX::Experimental::HalfToFull::count"), n);
+  Kokkos::parallel_for(
+      "ArborX::Experimental::HalfToFull::rewrite",
+      Kokkos::RangePolicy<ExecutionSpace>(space, 0, n), KOKKOS_LAMBDA(int i) {
+        auto const offsets_i = offsets(i);
+        for (int j = offsets_orig(i); j < offsets_orig(i + 1); ++j)
+        {
+          int const k = indices_orig(j);
+          indices(offsets_i + Kokkos::atomic_fetch_inc(&counts(i))) = k;
+          indices(offsets(k) + Kokkos::atomic_fetch_inc(&counts(k))) = i;
+        }
+      });
+  Kokkos::Profiling::popRegion();
+}
+
+} // namespace ArborX::Details
+
+#endif

test/CMakeLists.txt

@@ -205,6 +205,7 @@ add_test(NAME ArborX_Test_DetailsClusteringHelpers COMMAND ./ArborX_Test_Details
 
 add_executable(ArborX_Test_SpecializedTraversals.exe
   tstDetailsHalfTraversal.cpp
+  tstDetailsExpandHalfToFull.cpp
   utf_main.cpp
 )
 target_link_libraries(ArborX_Test_SpecializedTraversals.exe PRIVATE ArborX Boost::unit_test_framework)

test/boost_ext/CompressedStorageComparison.hpp

@@ -13,6 +13,7 @@
 #define ARBORX_BOOST_TEST_COMPRESSED_STORAGE_COMPARISON_HPP
 
 #include <boost/test/tools/detail/print_helper.hpp>
+#include <boost/test/utils/is_forward_iterable.hpp>
 
 #include <iosfwd>
 #include <set>

test/tstDetailsExpandHalfToFull.cpp

@@ -0,0 +1,78 @@
+/****************************************************************************
+ * Copyright (c) 2017-2023 by the ArborX authors                            *
+ * All rights reserved.                                                     *
+ *                                                                          *
+ * This file is part of the ArborX library. ArborX is                       *
+ * distributed under a BSD 3-clause license. For the licensing terms see    *
+ * the LICENSE file in the top-level directory.                             *
+ *                                                                          *
+ * SPDX-License-Identifier: BSD-3-Clause                                    *
+ ****************************************************************************/
+
+// clang-format off
+#include "boost_ext/CompressedStorageComparison.hpp"
+// clang-format on
+
+#include "ArborXTest_StdVectorToKokkosView.hpp"
+#include "ArborX_EnableDeviceTypes.hpp" // ARBORX_DEVICE_TYPES
+#include "ArborX_EnableViewComparison.hpp"
+#include <ArborX_DetailsExpandHalfToFull.hpp>
+
+#include "BoostTest_CUDA_clang_workarounds.hpp"
+#include <boost/test/unit_test.hpp>
+
+namespace Test
+{
+using ArborXTest::toView;
+
+template <class ExecutionSpace>
+auto expand(ExecutionSpace space, std::vector<int> const &offsets_host,
+            std::vector<int> const &indices_host)
+{
+  auto offsets = toView<ExecutionSpace>(offsets_host, "Test::offsets");
+  auto indices = toView<ExecutionSpace>(indices_host, "Test::indices");
+  ArborX::Details::expandHalfToFull(space, offsets, indices);
+
+  return make_compressed_storage(
+      Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, offsets),
+      Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, indices));
+}
+
+#define ARBORX_TEST_EXPAND_HALF_TO_FULL(exec_space, offsets_in, indices_in,    \
+                                        offsets_out, indices_out)              \
+  BOOST_TEST(Test::expand(exec_space, offsets_in, indices_in) ==               \
+                 make_compressed_storage(offsets_out, indices_out),            \
+             boost::test_tools::per_element())
+
+} // namespace Test
+
+BOOST_AUTO_TEST_CASE_TEMPLATE(expand_half_to_full, DeviceType,
+                              ARBORX_DEVICE_TYPES)
+{
+  using ExecutionSpace = typename DeviceType::execution_space;
+  ExecutionSpace exec_space;
+
+  ARBORX_TEST_EXPAND_HALF_TO_FULL(exec_space, (std::vector<int>{0}),
+                                  (std::vector<int>{}), (std::vector<int>{0}),
+                                  (std::vector<int>{}));
+
+  // (0,0) -> (0,0)(0,0)
+  ARBORX_TEST_EXPAND_HALF_TO_FULL(
+      exec_space, (std::vector<int>{0, 1}), (std::vector<int>{0}),
+      (std::vector<int>{0, 2}), (std::vector<int>{0, 0}));
+
+  // (0,0)(0,0) -> (0,0)(0,0)(0,0)(0,0)
+  ARBORX_TEST_EXPAND_HALF_TO_FULL(
+      exec_space, (std::vector<int>{0, 2}), (std::vector<int>{0, 0}),
+      (std::vector<int>{0, 4}), (std::vector<int>{0, 0, 0, 0}));
+
+  // (0,0)(0,1) -> (0,0)(0,1)(0,0)(1,0)
+  ARBORX_TEST_EXPAND_HALF_TO_FULL(
+      exec_space, (std::vector<int>{0, 2, 2}), (std::vector<int>{0, 1}),
+      (std::vector<int>{0, 3, 4}), (std::vector<int>{0, 1, 0, 0}));
+
+  // (1,0) -> (0,1)(1,0)
+  ARBORX_TEST_EXPAND_HALF_TO_FULL(
+      exec_space, (std::vector<int>{0, 0, 1, 1, 1}), (std::vector<int>{0}),
+      (std::vector<int>{0, 1, 2, 2, 2}), (std::vector<int>{1, 0}));
+}