Add ArborX neighbor lists

Currently compile-time switch which always uses ArborX if found

CMakeLists.txt

@@ -6,4 +6,7 @@ include(GNUInstallDirs)
 
 find_package(Cabana)
 
+option(HACCabana_REQUIRE_ARBORX "Build HACCabana with Cabana-ArborX neighbor list support" ${Cabana_ENABLE_ARBORX})
+set(HACCabana_ENABLE_ARBORX ${HACCabana_REQUIRE_ARBORX})
+
 add_subdirectory(src)

cmake/HACCabanaConfig.cmakein

@@ -0,0 +1,9 @@
+include(CMakeFindDependencyMacro)
+list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_LIST_DIR}" )
+list(APPEND CMAKE_PREFIX_PATH @CMAKE_PREFIX_PATH@)
+find_dependency(Kokkos REQUIRED)
+find_dependency(Cabana REQUIRED)
+set(HACCabana_ENABLE_ARBORX @HACCabana_ENABLE_ARBORX@)
+if(HACCabana_ENABLE_ARBORX)
+  find_dependency(ArborX REQUIRED)
+endif()

src/CMakeLists.txt

@@ -1,3 +1,5 @@
+configure_file(HACCabana_Config.h.cmakein HACCabana_Config.h)
+
 file(GLOB HACCABANA_HEADERS *.h)
 file(GLOB HACCABANA_SOURCES *.cxx)
 

src/HACCabana_Config.h.cmakein

@@ -0,0 +1,6 @@
+#ifndef HACCABANA_CONFIG_HPP
+#define HACCABANA_CONFIG_HPP
+
+#cmakedefine HACCabana_ENABLE_ARBORX
+
+#endif

src/Neighbor.h

@@ -0,0 +1,85 @@
+#ifndef NEIGHBORS_H
+#define NEIGHBORS_H
+
+#include "HACCabana_Config.h"
+
+namespace HACCabana
+{
+  struct LinkedCellTag
+  {
+  };
+  struct ArborXTag
+  {
+  };
+
+  template <typename AoSoADevice, typename NeighborType>
+  struct Neighbors;
+
+  template <typename AoSoADevice>
+  struct Neighbors<AoSoADevice, LinkedCellTag>
+  {
+    using device_type = typename AoSoADevice::device_type;
+    using neighbor_type = Cabana::LinkedCellList<device_type>;
+    using tag_type = LinkedCellTag;
+    // Store to avoid reallocation during rebuild.
+    neighbor_type cell_list;
+
+    Neighbors(AoSoADevice aosoa_device,
+              const float, const float cm_size, const float min_pos,
+              const float max_pos)
+    {
+      auto position = Cabana::slice<HACCabana::Particles::Fields::Position>(
+          aosoa_device, "position");
+
+      // create the cell list on the GPU
+      // NOTE: fuzz particles (outside of overload) are not included
+      float dx = cm_size;
+      float x_min = min_pos;
+      float x_max = max_pos;
+
+      float grid_delta[3] = {dx, dx, dx};
+      float grid_min[3] = {x_min, x_min, x_min};
+      float grid_max[3] = {x_max, x_max, x_max};
+
+      cell_list = neighbor_type(position, grid_delta, grid_min, grid_max);
+    }
+
+    template <typename ParticlesType>
+    void update(ParticlesType aosoa_device)
+    {
+      cell_list.build(aosoa_device);
+      Cabana::permute(cell_list, aosoa_device);
+    }
+  };
+
+#ifdef HACCabana_ENABLE_ARBORX
+  template <typename AoSoADevice>
+  struct Neighbors<AoSoADevice, ArborXTag>
+  {
+    using device_type = typename AoSoADevice::device_type;
+    using neighbor_type =
+        Cabana::Experimental::CrsGraph<typename device_type::memory_space,
+                                       Cabana::FullNeighborTag>;
+    using tag_type = ArborXTag;
+    neighbor_type neighbor_list;
+
+    Neighbors(AoSoADevice aosoa_device, const float rmax2, const float,
+              const float, const float)
+    {
+      auto position = Cabana::slice<HACCabana::Particles::Fields::Position>(
+          aosoa_device, "position");
+      // Interface in Cabana to ArborX neighbor search.
+      neighbor_list = Cabana::Experimental::makeNeighborList<device_type>(
+          Cabana::FullNeighborTag{}, position, 0, position.size(), rmax2);
+    }
+  };
+#endif
+
+  template <typename NeighborType, typename AoSoADevice>
+  auto createNeighbors(AoSoADevice &aosoa_device, const float rmax2, const float cm_size, const float min_pos,
+                       const float max_pos)
+  {
+    return Neighbors<AoSoADevice, NeighborType>(aosoa_device, rmax2, cm_size, min_pos, max_pos);
+  }
+}
+#endif