[Draft] Parallelization of non-bonded short-range forces with Cabana

CMakeLists.txt

@@ -130,6 +130,8 @@ FetchContent_Declare(
 # User input options
 # ##############################################################################
 
+option(ESPRESSO_BUILD_WITH_CABANA "Build with Cabana" ON)
+
 option(ESPRESSO_BUILD_WITH_PYTHON "Build with Python bindings" ON)
 option(ESPRESSO_BUILD_WITH_GSL "Build with GSL support" OFF)
 option(ESPRESSO_BUILD_WITH_FFTW "Build with FFTW support" ON)
@@ -482,6 +484,27 @@ endif()
 # Libraries
 #
 
+if(ESPRESSO_BUILD_WITH_CABANA)
+  # TODO Kokkos_DIR needs to be set here
+
+  find_package(Kokkos REQUIRED)
+
+  include(FetchContent)
+  if(EXISTS "${CMAKE_BINARY_DIR}/_deps/cabana-build")
+    FetchContent_Declare(cabana
+      GIT_REPOSITORY https://github.com/ECP-copa/Cabana.git
+      GIT_TAG 07573782590e223edc6e5f83ec0807be5941df14
+    )
+  else()# Fix cabana doxygen conflict until it is fixed upstream
+    FetchContent_Declare(cabana
+      GIT_REPOSITORY https://github.com/ECP-copa/Cabana.git
+      GIT_TAG 07573782590e223edc6e5f83ec0807be5941df14
+      PATCH_COMMAND patch -N --reject-file=- -i "${CMAKE_CURRENT_SOURCE_DIR}/cmake/patch_cabana_doxygen.diff" 
+    )
+  endif()
+  FetchContent_MakeAvailable(cabana)
+endif()
+
 if(ESPRESSO_BUILD_WITH_FFTW)
   find_package(FFTW3 REQUIRED)
 endif()

cmake/espresso_cmake_config.cmakein

@@ -23,6 +23,8 @@
 
 #cmakedefine ESPRESSO_BUILD_WITH_FPE
 
+#cmakedefine ESPRESSO_BUILD_WITH_CABANA
+
 #define PACKAGE_NAME "${PROJECT_NAME}"
 
 /**

cmake/patch_cabana_doxygen.diff

@@ -0,0 +1,17 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index c228d50a..00693256 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -216,12 +216,6 @@ if(Cabana_ENABLE_TESTING)
+   enable_testing()
+ endif()
+
+-# enable doxygen
+-find_package(Doxygen)
+-if(Doxygen_FOUND)
+-  doxygen_add_docs(doxygen core/src grid/src)
+-endif()
+-
+ ##---------------------------------------------------------------------------##
+ ## Libraries and Examples
+ ##---------------------------------------------------------------------------##

src/config/features.def

@@ -116,3 +116,5 @@ WALBERLA_FFT external
 VALGRIND external
 CALIPER external
 FPE external
+CABANA external
+

src/core/CMakeLists.txt

@@ -95,6 +95,10 @@ if(ESPRESSO_BUILD_WITH_FFTW)
   add_subdirectory(fft)
 endif()
 
+if(ESPRESSO_BUILD_WITH_CABANA)
+  target_link_libraries(espresso_core PRIVATE Cabana::Core)
+endif()
+
 add_subdirectory(accumulators)
 add_subdirectory(analysis)
 add_subdirectory(bond_breakage)

src/core/cabana_data.hpp

@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2010-2025 The ESPResSo project
+ *
+ * This file is part of ESPResSo.
+ *
+ * ESPResSo is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * ESPResSo is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#pragma once
+
+#include <iostream>
+
+#ifdef CABANA
+
+#include <Cabana_Core.hpp>
+#include "custom_verlet_list.hpp"
+#include <unordered_map>
+#endif
+
+#ifdef CABANA
+
+using data_types = Cabana::MemberTypes<double[3], double[3], int, int, int>;
+using memory_space = Kokkos::SharedSpace;
+using execution_space = Kokkos::DefaultExecutionSpace;
+
+using ListAlgorithm = Cabana::HalfNeighborTag;
+using ListType = Cabana::CustomVerletList<memory_space, ListAlgorithm, Cabana::VerletLayout2D>;
+
+class CabanaData {
+    ListType verlet_list;
+    std::unordered_map<int, int> id_to_index;
+
+public:
+    CabanaData() = default;
+    CabanaData(ListType verlet_list, std::unordered_map<int, int> id_to_index)
+      : verlet_list(verlet_list), id_to_index(id_to_index) {}
+
+    ListType get_verlet_list() const { return verlet_list; }
+    std::unordered_map<int, int> get_id_to_index() const { return id_to_index; }
+
+    ~CabanaData() {};
+
+
+};
+#endif 

src/core/cell_system/CellStructure.cpp

@@ -50,6 +50,38 @@
 #include <string>
 #include <utility>
 #include <vector>
+#include <iostream>
+
+#ifdef CABANA
+#include <Cabana_Core.hpp>
+#include "custom_verlet_list.hpp"
+#include "cabana_data.hpp"
+#endif
+
+#ifdef CABANA
+
+using data_types = Cabana::MemberTypes<double[3], double[3], int, int, int>;
+using memory_space = Kokkos::SharedSpace;
+using execution_space = Kokkos::DefaultExecutionSpace;
+
+using ListAlgorithm = Cabana::HalfNeighborTag;
+using ListType = Cabana::CustomVerletList<memory_space, ListAlgorithm, Cabana::VerletLayout2D>;
+
+
+CellStructure::~CellStructure() {
+  m_cabana_data.reset();
+};
+
+void CellStructure::set_cabana_data(std::unique_ptr<CabanaData> data) {
+  m_cabana_data = std::move(data);
+}
+
+CabanaData& CellStructure::get_cabana_data() {
+  return *m_cabana_data;
+}
+
+#endif
+
 
 CellStructure::CellStructure(BoxGeometry const &box)
     : m_decomposition{std::make_unique<AtomDecomposition>(box)} {}

src/core/cell_system/CellStructure.hpp

@@ -52,6 +52,10 @@
 #include <stdexcept>
 #include <utility>
 #include <vector>
+#include <any>
+
+// forward declaration to not have to import cabana
+class CabanaData;
 
 namespace Cells {
 enum Resort : unsigned {
@@ -603,6 +607,49 @@ struct CellStructure : public System::Leaf<CellStructure> {
     }
   }
 
+#ifdef CABANA
+private:
+    std::unique_ptr<CabanaData> m_cabana_data;
+
+public:
+  void set_cabana_data(std::unique_ptr<CabanaData> data);
+  CabanaData& get_cabana_data();
+
+  virtual ~CellStructure();
+
+  bool get_rebuild_verlet_list() const { return m_rebuild_verlet_list; }
+
+  template <class Kernel>
+  void cabana_link_cell(Kernel kernel) {
+    auto const local_cells_span = decomposition().local_cells();
+    auto const first = boost::make_indirect_iterator(local_cells_span.begin());
+    auto const last = boost::make_indirect_iterator(local_cells_span.end());
+
+    Algorithm::link_cell(first, last, [&kernel](Particle &p1, Particle &p2) {
+      kernel(p1, p2);
+    });
+  }
+
+  template <class Kernel, class VerletCriterion>
+  void cabana_verlet_list_loop(Kernel kernel,
+                               const VerletCriterion &verlet_criterion) {
+    if (m_rebuild_verlet_list) {
+      m_verlet_list.clear();
+
+      link_cell([&](Particle &p1, Particle &p2, Distance const &d) {
+        if (verlet_criterion(p1, p2, d)) {
+          m_verlet_list.emplace_back(&p1, &p2);
+        }
+      });
+      m_rebuild_verlet_list = false;
+    } 
+    for (auto const &pair : m_verlet_list) {
+      kernel(*pair.first, *pair.second);
+    } 
+  }
+#endif
+
+private:
   /** Non-bonded pair loop with verlet lists.
    *
    * @param pair_kernel Kernel to apply

src/core/communication.cpp

@@ -30,6 +30,10 @@
 #include <walberla_bridge/walberla_init.hpp>
 #endif
 
+#ifdef CABANA
+#include <Cabana_Core.hpp>
+#endif
+
 #include <utils/Vector.hpp>
 #include <utils/mpi/cart_comm.hpp>
 
@@ -84,10 +88,20 @@ void init(std::shared_ptr<boost::mpi::environment> mpi_env) {
 #ifdef CUDA
   cuda_on_program_start();
 #endif
-}
 
-void deinit() { Communication::m_callbacks.reset(); }
-} // namespace Communication
+#ifdef CABANA
+  Kokkos::initialize();
+#endif
+  }
+
+void deinit() {
+  Communication::m_callbacks.reset(); 
+
+#ifdef CABANA 
+  Kokkos::finalize();
+#endif
+  }
+}
 
 Communicator::Communicator()
     : comm{::comm_cart}, node_grid{}, this_node{::this_node}, size{-1} {}