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