Kokkos implementation of for_each_local_particle()

Author: RudolfWeeber

CMakeLists.txt

@@ -503,6 +503,8 @@ if(ESPRESSO_BUILD_WITH_FFTW)
 endif()
 
 if(ESPRESSO_BUILD_WITH_SHARED_MEMORY_PARALLELISM)
+  find_package(OpenMP REQUIRED)
+
   if(NOT EXISTS ${FETCHCONTENT_BASE_DIR}/kokkos-src)
     find_package(Kokkos 4.3 QUIET)
   endif()

src/core/CMakeLists.txt

@@ -90,7 +90,11 @@ install(TARGETS espresso_core
 target_link_libraries(
   espresso_core
   PRIVATE
-    espresso::config espresso::utils::mpi espresso::shapes espresso::cpp_flags
+    espresso::config
+    espresso::utils::mpi
+    espresso::shapes
+    espresso::cpp_flags
+    $<$<BOOL:${ESPRESSO_BUILD_WITH_SHARED_MEMORY_PARALLELISM}>:OpenMP::OpenMP_CXX>
     $<$<BOOL:${ESPRESSO_BUILD_WITH_SHARED_MEMORY_PARALLELISM}>:Kokkos::kokkos>
     $<$<BOOL:${ESPRESSO_BUILD_WITH_SHARED_MEMORY_PARALLELISM}>:Cabana::Core>
   PUBLIC espresso::utils MPI::MPI_CXX Random123 espresso::particle_observables

src/core/cell_system/CellStructure.cpp

@@ -51,6 +51,10 @@
 #include <utility>
 #include <vector>
 
+#ifdef SHARED_MEMORY_PARALLELISM
+#include <Kokkos_Core.hpp>
+#endif
+
 CellStructure::CellStructure(BoxGeometry const &box)
     : m_decomposition{std::make_unique<AtomDecomposition>(box)} {}
 
@@ -327,3 +331,22 @@ void CellStructure::update_ghosts_and_resort_particle(unsigned data_parts) {
     ghosts_update(data_parts & ~resort_only_parts);
   }
 }
+
+#ifdef SHARED_MEMORY_PARALLELISM
+void CellStructure::parallel_for_each_particle_impl(
+    std::span<Cell *const> cells, ParticleUnaryOp &f) const {
+  if (cells.size() > 1) {
+    Kokkos::parallel_for( // loop over cells
+        "for_each_local_particle", cells.size(), [&](int cell_idx) {
+          for (auto &p : cells[cell_idx]->particles())
+            f(p);
+        });
+  } else if (cells.size() == 1) {
+    auto const &cell = cells[0];
+    Kokkos::parallel_for( // loop over particles
+        "for_each_local_particle", cell->particles().size(),
+        [&](int part_idx) { f(*(cell->particles().begin() + part_idx)); });
+  }
+}
+
+#endif

src/core/cell_system/CellStructure.hpp

@@ -45,6 +45,7 @@
 #include <algorithm>
 #include <cassert>
 #include <concepts>
+#include <functional>
 #include <iterator>
 #include <memory>
 #include <optional>
@@ -54,10 +55,7 @@
 #include <utility>
 #include <vector>
 
-template <typename Callable>
-concept ParticleCallback = requires(Callable c, Particle &p) {
-  { c(p) } -> std::same_as<void>;
-};
+using ParticleUnaryOp = std::function<void(Particle &)>;
 
 namespace Cells {
 enum Resort : unsigned {
@@ -280,14 +278,26 @@ struct CellStructure : public System::Leaf<CellStructure> {
   ParticleRange ghost_particles() const {
     return Cells::particles(decomposition().ghost_cells());
   }
+  /** @brief whether to use parallel version of for_each_local_particle */
+  bool use_parallel_for_each_local_particle() const {
+#ifdef SHARED_MEMORY_PARALLELISM
+    return true;
+#else
+    return false;
+#endif
+  }
 
   /**
    * @brief Run a kernel on all local particles.
    * The kernel is assumed to be thread-safe.
    */
-  template <typename Kernel>
-    requires ParticleCallback<Kernel>
-  void for_each_local_particle(Kernel f) const {
+  void for_each_local_particle(ParticleUnaryOp &&f) const {
+#ifdef SHARED_MEMORY_PARALLELISM
+    if (use_parallel_for_each_local_particle()) {
+      parallel_for_each_particle_impl(decomposition().local_cells(), f);
+      return;
+    }
+#endif
     for (auto &p : local_particles()) {
       f(p);
     }
@@ -297,14 +307,16 @@ struct CellStructure : public System::Leaf<CellStructure> {
    * @brief Run a kernel on all ghost particles.
    * The kernel is assumed to be thread-safe.
    */
-  template <typename Kernel>
-    requires ParticleCallback<Kernel>
-  void for_each_ghost_particle(Kernel f) const {
+  void for_each_ghost_particle(ParticleUnaryOp &&f) const {
     for (auto &p : ghost_particles()) {
       f(p);
     }
   }
 
+#ifdef SHARED_MEMORY_PARALLELISM
+  void parallel_for_each_particle_impl(std::span<Cell *const> cells,
+                                       ParticleUnaryOp &f) const;
+#endif
 private:
   /** Cell system dependent function to find the right cell for a
    *  particle.