Shared memory parallelization for short range forces

CMakeLists.txt

@@ -580,6 +580,11 @@ if(ESPRESSO_BUILD_WITH_SHARED_MEMORY_PARALLELISM)
     set(CMAKE_SHARED_LIBRARY_PREFIX "lib")
     set(Kokkos_ENABLE_SERIAL ON CACHE BOOL "")
     set(Kokkos_ENABLE_OPENMP ON CACHE BOOL "")
+    set(Kokkos_ENABLE_IMPL_VIEW_LEGACY ON CACHE BOOL "")
+    set(Kokkos_ENABLE_COMPLEX_ALIGN ON CACHE BOOL "")
+    set(Kokkos_ENABLE_AGGRESSIVE_VECTORIZATION ON CACHE BOOL "")
+    set(Kokkos_ENABLE_HWLOC ON CACHE BOOL "")
+    set(Kokkos_ARCH_NATIVE ON CACHE BOOL "")
     FetchContent_MakeAvailable(kokkos)
     set(BUILD_SHARED_LIBS ${ESPRESSO_BUILD_SHARED_LIBS_DEFAULT})
     set(CMAKE_SHARED_LIBRARY_PREFIX "${ESPRESSO_SHARED_LIBRARY_PREFIX}")
@@ -604,7 +609,7 @@ if(ESPRESSO_BUILD_WITH_SHARED_MEMORY_PARALLELISM)
     FetchContent_Declare(
       cabana
       GIT_REPOSITORY https://github.com/ECP-copa/Cabana.git
-      GIT_TAG        ebfaa51 # 0.7.0 with patches
+      GIT_TAG        e76c1a1 # 0.7.0 with patches
       PATCH_COMMAND  patch -p0 < ${CMAKE_CURRENT_SOURCE_DIR}/cmake/cabana.patch
     )
     # cmake-format: on
@@ -628,6 +633,7 @@ if(ESPRESSO_BUILD_WITH_SHARED_MEMORY_PARALLELISM)
     list(APPEND SKIP_CLANG_TIDY_CHECKS "-modernize-use-bool-literals")
     list(APPEND SKIP_CLANG_TIDY_CHECKS "-modernize-use-equals-delete")
     list(APPEND SKIP_CLANG_TIDY_CHECKS "-modernize-use-equals-default")
+    list(APPEND SKIP_CLANG_TIDY_CHECKS "-modernize-use-nullptr")
     list(APPEND SKIP_CLANG_TIDY_CHECKS "-modernize-pass-by-value")
     list(APPEND SKIP_CLANG_TIDY_CHECKS "-modernize-loop-convert")
     list(APPEND SKIP_CLANG_TIDY_CHECKS "-modernize-return-braced-init-list")
@@ -636,6 +642,10 @@ if(ESPRESSO_BUILD_WITH_SHARED_MEMORY_PARALLELISM)
     list(APPEND SKIP_CLANG_TIDY_CHECKS "-readability-simplify-boolean-expr")
     list(APPEND SKIP_CLANG_TIDY_CHECKS
          "-readability-avoid-const-params-in-decls")
+    list(APPEND SKIP_CLANG_TIDY_CHECKS
+         "-clang-analyzer-optin.performance.Padding")
+    list(APPEND SKIP_CLANG_TIDY_CHECKS
+         "-clang-analyzer-security.insecureAPI.strcpy")
     espresso_override_clang_tidy_checks(
       ESPRESSO_CABANA_CXX_CLANG_TIDY "${SKIP_CLANG_TIDY_CHECKS}"
       "${SKIP_CLANG_TIDY_CHECKS_CXX}")

src/core/BoxGeometry.hpp

@@ -225,6 +225,40 @@ class BoxGeometry {
             get_mi_coord(a[2], b[2], 2)};
   }
 
+  /**
+   * @brief Get the minimum-image vector between two coordinates.
+   *
+   * @tparam T Floating point type.
+   *
+   * @param a0     x element of the terminal point.
+   * @param a1     y element of the terminal point.
+   * @param a2     z element of the terminal point.
+   * @param b0     x element of the initial point.
+   * @param b1     y element of the initial point.
+   * @param b2     z element of the initial point.
+   * @return Vector from @p b to @p a that minimizes the distance across
+   *         periodic images, i.e. <tt>a - b</tt>.
+   */
+  template <typename T>
+  Utils::Vector<T, 3> get_mi_vector(T const &a0, T const &a1, T const &a2,
+                                    T const &b0, T const &b1,
+                                    T const &b2) const {
+    if (type() == BoxType::LEES_EDWARDS) {
+      auto const shear_plane_normal = lees_edwards_bc().shear_plane_normal;
+      auto a_tmp = Utils::Vector<T, 3>{a0, a1, a2};
+      auto b_tmp = Utils::Vector<T, 3>{b0, b1, b2};
+      a_tmp[shear_plane_normal] = Algorithm::periodic_fold(
+          a_tmp[shear_plane_normal], m_length[shear_plane_normal]);
+      b_tmp[shear_plane_normal] = Algorithm::periodic_fold(
+          b_tmp[shear_plane_normal], m_length[shear_plane_normal]);
+      return lees_edwards_bc().distance(a_tmp - b_tmp, m_length, m_length_half,
+                                        m_length_inv, m_periodic);
+    }
+    assert(type() == BoxType::CUBOID);
+    return {get_mi_coord(a0, b0, 0), get_mi_coord(a1, b1, 1),
+            get_mi_coord(a2, b2, 2)};
+  }
+
   BoxType type() const { return m_type; }
   void set_type(BoxType type) { m_type = type; }
 

src/core/aosoa_pack.hpp

@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 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
+
+#ifdef SHARED_MEMORY_PARALLELISM
+
+#include "cell_system/CellStructure.hpp"
+
+#include <Cabana_Core.hpp>
+
+struct CellStructure::AoSoA_pack {
+  CellStructure::AoSoAType::member_slice_type<0> position;
+  CellStructure::AoSoAType::member_slice_type<1> charge;
+  CellStructure::AoSoAType::member_slice_type<2> id;
+  CellStructure::AoSoAType::member_slice_type<3> type;
+
+  AoSoA_pack() = default;
+
+  AoSoA_pack(CellStructure::AoSoAType &aosoa)
+      : position(Cabana::slice<0>(aosoa)), charge(Cabana::slice<1>(aosoa)),
+        id(Cabana::slice<2>(aosoa)), type(Cabana::slice<3>(aosoa)) {}
+};
+
+#endif // SHARED_MEMORY_PARALLELISM

src/core/cell_system/CellStructure.cpp

@@ -32,11 +32,13 @@
 #include "cell_system/CellStructureType.hpp"
 #include "communication.hpp"
 #include "lees_edwards/lees_edwards.hpp"
+#include "particle_enumeration.hpp"
 #include "particle_reduction.hpp"
 #include "system/System.hpp"
 
 #include <utils/Vector.hpp>
 #include <utils/contains.hpp>
+#include <utils/math/int_pow.hpp>
 #include <utils/math/sqr.hpp>
 
 #include <boost/mpi/collectives/all_reduce.hpp>
@@ -56,9 +58,139 @@
 #include <vector>
 
 #ifdef SHARED_MEMORY_PARALLELISM
+#include "aosoa_pack.hpp"
+#include "custom_verlet_list.hpp"
+#include <Cabana_Core.hpp>
+#include <Cabana_NeighborList.hpp>
 #include <Kokkos_Core.hpp>
 #endif
 
+CellStructure::~CellStructure() {
+#ifdef SHARED_MEMORY_PARALLELISM
+  if (m_local_force) {
+    m_local_force.reset();
+  }
+#ifdef ROTATION
+  if (m_local_torque) {
+    m_local_torque.reset();
+  }
+#endif
+#ifdef NPT
+  if (m_local_virial) {
+    m_local_virial.reset();
+  }
+#endif
+  if (m_aosoa) {
+    m_aosoa.reset();
+  }
+  if (m_particle_storage) {
+    m_particle_storage.reset();
+  }
+  if (m_verlet_list_cabana) {
+    m_verlet_list_cabana.reset();
+  }
+  // Kokkos handle can be freed after all Cabana containers have been freed
+  m_kokkos_handle.reset();
+#endif
+}
+
+#ifdef SHARED_MEMORY_PARALLELISM
+
+void CellStructure::set_kokkos_handle(
+    std::shared_ptr<Communication::KokkosHandle> handle) {
+  m_kokkos_handle = std::move(handle);
+}
+
+static auto estimate_max_counts(int max_prefactor, double pair_cutoff,
+                                std::size_t number_of_unique_particles) {
+  if (std::isinf(pair_cutoff)) {
+    return number_of_unique_particles;
+  }
+  if (pair_cutoff < 0.) {
+    pair_cutoff = 0.;
+  }
+  auto const volume = Utils::int_pow<3>(pair_cutoff);
+  auto max_counts = static_cast<std::size_t>(
+      std::ceil(static_cast<double>(max_prefactor) * volume));
+#ifdef COLLISION_DETECTION
+  std::size_t constexpr threshold_num = 64;
+#else
+  std::size_t constexpr threshold_num = 16;
+#endif
+  if (max_counts < threshold_num) {
+    max_counts = std::min(threshold_num, number_of_unique_particles);
+  }
+  return max_counts;
+}
+
+void CellStructure::rebuild_local_properties(std::size_t const num_threads,
+                                             double const pair_cutoff) {
+  assert(m_kokkos_handle);
+  auto const num_part = get_unique_particles().size();
+  m_local_force =
+      std::make_unique<ForceType>("local_force", num_part, num_threads);
+#ifdef ROTATION
+  m_local_torque =
+      std::make_unique<ForceType>("local_torque", num_part, num_threads);
+#endif
+#ifdef NPT
+  m_local_virial = std::make_unique<VirialType>("local_virial", num_threads);
+#endif
+  m_particle_storage = std::make_unique<AoSoAType>("particles", num_part);
+  m_particle_storage->resize(num_part);
+  // particle properties are defined in aosoa_pack.hpp
+  m_aosoa = std::make_unique<AoSoA_pack>(*m_particle_storage);
+
+  auto max_counts = estimate_max_counts(m_max_prefactor, pair_cutoff, num_part);
+  m_verlet_list_cabana = std::make_unique<ListType>(0ul, num_part, max_counts);
+}
+
+void CellStructure::reset_local_properties() {
+  Kokkos::deep_copy(get_local_force(), 0);
+#ifdef ROTATION
+  Kokkos::deep_copy(get_local_torque(), 0);
+#endif
+#ifdef NPT
+  Kokkos::deep_copy(get_local_virial(), 0);
+#endif
+}
+
+void CellStructure::set_index_map() {
+  auto &unique_particles = m_unique_particles;
+  unique_particles.clear();
+  unique_particles.resize(count_local_particles());
+  std::unordered_set<int> registered_index{};
+  using execution_space = Kokkos::DefaultExecutionSpace;
+  int n_threads = execution_space().concurrency();
+  std::vector<int> max_ids(n_threads);
+  enumerate_local_particles(
+      *this, [&unique_particles, &max_ids](std::size_t index, Particle &p) {
+        unique_particles[index] = &p;
+        const int thread_num = omp_get_thread_num();
+        max_ids[thread_num] = std::max(p.id(), max_ids[thread_num]);
+      });
+  int max_id = *(std::max_element(max_ids.begin(), max_ids.end()));
+  for (auto &p : ghost_particles()) {
+    auto const *local_particle = get_local_particle(p.id());
+    if (not local_particle) {
+      continue;
+    }
+    if (not local_particle->is_ghost()) {
+      continue;
+    }
+    if (registered_index.contains(p.id())) {
+      continue;
+    }
+    registered_index.insert(p.id());
+    unique_particles.emplace_back(&p);
+    max_id = std::max(p.id(), max_id);
+  }
+  registered_index.clear();
+  m_cached_max_local_particle_id = max_id;
+}
+
+#endif // SHARED_MEMORY_PARALLELISM
+
 CellStructure::CellStructure(BoxGeometry const &box)
     : m_decomposition{std::make_unique<AtomDecomposition>(box)} {}
 
@@ -236,6 +368,7 @@ void CellStructure::resort_particles(bool global_flag) {
 
   auto const &lebc = get_system().box_geo->lees_edwards_bc();
   m_rebuild_verlet_list = true;
+  m_rebuild_verlet_list_cabana = true;
   m_le_pos_offset_at_last_resort = lebc.pos_offset;
 
 #ifdef ADDITIONAL_CHECKS
@@ -285,6 +418,7 @@ void CellStructure::set_verlet_skin(double value) {
   assert(value >= 0.);
   m_verlet_skin = value;
   m_verlet_skin_set = true;
+  m_rebuild_verlet_list_cabana = true;
   get_system().on_verlet_skin_change();
 }