Improve P3M shared-memory parallelization

Use shared memory in P3M/DP3M charge/dipole/force/torque assignment
and real-space force calculation in the short-range loop.

Author: jngrad

doc/tutorials/charged_system/charged_system.ipynb

@@ -44,7 +44,7 @@
     "import espressomd.zn\n",
     "\n",
     "\n",
-    "espressomd.assert_features(['ELECTROSTATICS', 'P3M', 'WCA'])\n",
+    "espressomd.assert_features(['ELECTROSTATICS', 'P3M', 'WCA', 'EXTERNAL_FORCES'])\n",
     "\n",
     "import tqdm\n",
     "import numpy as np\n",

doc/tutorials/convert.py

@@ -188,7 +188,7 @@ def execute_notebook(nb, src, cell_separator, notebook_filepath):
     import importlib_wrapper as iw
     notebook_dirname = os.path.dirname(notebook_filepath)
     # disable OpenGL GUI
-    src_no_gui = iw.mock_es_visualization(src)
+    src_no_gui = iw.mock_es_visualization(src, force_mock=True)
     # update notebook with new code
     set_code_cells(nb, src_no_gui.split(cell_separator))
     # execute notebook

doc/tutorials/electrodes/electrodes_part1.ipynb

@@ -192,7 +192,7 @@
     "import espressomd.electrostatics\n",
     "import espressomd.electrostatic_extensions\n",
     "\n",
-    "espressomd.assert_features(['ELECTROSTATICS'])\n",
+    "espressomd.assert_features(['ELECTROSTATICS', 'EXTERNAL_FORCES'])\n",
     "plt.rcParams.update({'font.size': 18})"
    ]
   },
@@ -251,11 +251,7 @@
     "ICC_EPSILON_WALLS = 1e5     # epsilon outside the slit. Very large to mimic metal\n",
     "ICC_CONVERGENCE = 1e-3      # ICC numeric/performance parameters\n",
     "ICC_RELAXATION = 0.95\n",
-    "ICC_MAX_ITERATIONS = 1000\n",
-    "\n",
-    "# Lennard-Jones parameters\n",
-    "LJ_SIGMA = 1.0\n",
-    "LJ_EPSILON = 1.0                \n",
+    "ICC_MAX_ITERATIONS = 100\n",
     "\n",
     "# Particle parameters\n",
     "TYPES = {\"Cation\": 0, \"Anion\": 1  ,\"Electrodes\": 2}\n",
@@ -309,7 +305,7 @@
    "outputs": [],
    "source": [
     "# SOLUTION CELL\n",
-    "elc = espressomd.electrostatics.ELC(actor=p3m, gap_size=ELC_GAP, maxPWerror=MAX_PW_ERROR)"
+    "elc = espressomd.electrostatics.ELC(actor=p3m, gap_size=ELC_GAP, maxPWerror=MAX_PW_ERROR, check_neutrality=False)"
    ]
   },
   {

maintainer/parsing/importlib_wrapper.py

@@ -41,6 +41,7 @@ def configure_and_import(filepath,
                          script_suffix="",
                          move_to_script_dir=True,
                          mock_visualizer=True,
+                         force_visualizer_mock=True,
                          **parameters):
     """
     Copy a Python script to a new location and alter some lines of code:
@@ -69,6 +70,8 @@ def configure_and_import(filepath,
         if ``True``, substitute the visualizer with a ``Mock`` class in case
         of ``ImportError`` (use ``False`` if an ``ImportError`` is relevant
         to your test)
+    force_visualizer_mock : :obj:`bool`
+        if ``True``, always substitute the visualizer with a ``Mock`` class
     move_to_script_dir : :obj:`bool`
         if ``True``, move to the script's directory (useful when the script
         needs to load files hardcoded as relative paths, or when files are
@@ -101,7 +104,7 @@ def configure_and_import(filepath,
     code = disable_matplotlib_gui(code)
     # disable OpenGL GUI in case of ImportError using MagicMock()
     if mock_visualizer:
-        code = mock_es_visualization(code)
+        code = mock_es_visualization(code, force_visualizer_mock)
     # save changes to a new file
     output_filepath = filepath.parent / \
         f"{filepath.stem}_{script_suffix}_processed.py"
@@ -433,7 +436,7 @@ def visit_ImportFrom(self, node):
                         node.lineno, node.module, child.name, child.asname)
 
 
-def mock_es_visualization(code):
+def mock_es_visualization(code, force_mock=False):
     """
     Replace ``import espressomd.visualization`` by a ``MagicMock``
     when the visualization module is unavailable, by catching the
@@ -451,6 +454,12 @@ def mock_es_visualization(code):
     import unittest.mock
     import espressomd
     {1} = unittest.mock.MagicMock()
+""".lstrip()
+    if force_mock:
+        r_es_vis_mock = r"""
+import unittest.mock
+import espressomd
+{1} = unittest.mock.MagicMock()
 """.lstrip()
 
     visitor = GetEspressomdVisualizerImports()

src/config/include/config/config.hpp

@@ -93,4 +93,16 @@
 #define MAX_OBJECTS_IN_FLUID 10000
 #endif
 
+/**
+ * @brief Special cutoff value for an inactive interaction.
+ * Non-bonded potentials that have this cutoff are never evaluated.
+ */
+inline constexpr double INACTIVE_CUTOFF = -1.;
+
+/**
+ * @brief Special cutoff value for an inactive bond.
+ * Bonds that have this cutoff are never evaluated.
+ */
+inline constexpr double BONDED_INACTIVE_CUTOFF = -1.;
+
 #endif

src/core/analysis/statistics.cpp

@@ -149,10 +149,10 @@ Utils::Vector3d calc_linear_momentum(System::System const &system,
   if (include_particles) {
     momentum = reduce_over_local_particles<Utils::Vector3d>(
         *(system.cell_structure),
-        [](Particle const &p, Utils::Vector3d &res) {
-          res += p.mass() * p.v();
+        [](Utils::Vector3d &acc, Particle const &p) {
+          acc += p.mass() * p.v();
         },
-        [](Utils::Vector3d &a, Utils::Vector3d const &b) { a = a + b; });
+        [](Utils::Vector3d &acc, Utils::Vector3d const &v) { acc = acc + v; });
   }
   if (include_lbfluid and system.lb.is_solver_set()) {
     momentum += system.lb.get_momentum() * system.lb.get_lattice_speed();

src/core/aosoa_pack.hpp

@@ -19,6 +19,8 @@
 
 #pragma once
 
+#include <config/config.hpp>
+
 #ifdef SHARED_MEMORY_PARALLELISM
 
 #include "cell_system/CellStructure.hpp"

src/core/bonded_interactions/bonded_interaction_data.hpp

@@ -24,6 +24,8 @@
  *  For more information on how to add new interactions, see @ref bondedIA_new.
  */
 
+#include <config/config.hpp>
+
 #include "angle_common.hpp"
 #include "angle_cosine.hpp"
 #include "angle_cossquare.hpp"
@@ -56,11 +58,6 @@
 #include <variant>
 #include <vector>
 
-/* Special cutoff value for a disabled bond.
- * Bonds that have this cutoff are not visited during bond evaluation.
- */
-static constexpr double BONDED_INACTIVE_CUTOFF = -1.;
-
 /** Interaction type for unused bonded interaction slots */
 struct NoneBond {
   static constexpr int num = 0;

src/core/cell_system/CellStructure.cpp

@@ -29,8 +29,10 @@
 #include "BoxGeometry.hpp"
 #include "LocalBox.hpp"
 #include "Particle.hpp"
+#include "aosoa_pack.hpp"
 #include "cell_system/CellStructureType.hpp"
 #include "communication.hpp"
+#include "custom_verlet_list.hpp"
 #include "lees_edwards/lees_edwards.hpp"
 #include "particle_enumeration.hpp"
 #include "particle_reduction.hpp"
@@ -43,6 +45,13 @@
 
 #include <boost/mpi/collectives/all_reduce.hpp>
 
+#ifdef SHARED_MEMORY_PARALLELISM
+#include <Cabana_Core.hpp>
+#include <Cabana_NeighborList.hpp>
+#include <Kokkos_Core.hpp>
+#include <omp.h>
+#endif
+
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
@@ -57,45 +66,29 @@
 #include <variant>
 #include <vector>
 
+CellStructure::~CellStructure() {
 #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>
+  clear_local_properties();
+  // Kokkos handle can only be freed after all Cabana containers have been freed
+  m_kokkos_handle.reset();
 #endif
+}
 
-CellStructure::~CellStructure() {
 #ifdef SHARED_MEMORY_PARALLELISM
-  if (m_local_force) {
-    m_local_force.reset();
-  }
+void CellStructure::clear_local_properties() {
+  m_local_force.reset();
 #ifdef ROTATION
-  if (m_local_torque) {
-    m_local_torque.reset();
-  }
+  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();
+  m_local_virial.reset();
 #endif
+  m_aosoa.reset();
+  m_particle_storage.reset();
+  m_verlet_list_cabana.reset();
+  m_rebuild_verlet_list_cabana = true;
 }
 
-#ifdef SHARED_MEMORY_PARALLELISM
-
 void CellStructure::set_kokkos_handle(
     std::shared_ptr<Communication::KokkosHandle> handle) {
   m_kokkos_handle = std::move(handle);
@@ -123,9 +116,10 @@ static auto estimate_max_counts(int max_prefactor, double pair_cutoff,
   return max_counts;
 }
 
-void CellStructure::rebuild_local_properties(std::size_t const num_threads,
-                                             double const pair_cutoff) {
+void CellStructure::rebuild_local_properties(double const pair_cutoff) {
   assert(m_kokkos_handle);
+  using execution_space = Kokkos::DefaultExecutionSpace;
+  auto const num_threads = execution_space().concurrency();
   auto const num_part = get_unique_particles().size();
   m_local_force =
       std::make_unique<ForceType>("local_force", num_part, num_threads);
@@ -145,13 +139,17 @@ void CellStructure::rebuild_local_properties(std::size_t const num_threads,
   m_verlet_list_cabana = std::make_unique<ListType>(0ul, num_part, max_counts);
 }
 
+void CellStructure::reset_local_force() {
+  Kokkos::deep_copy(get_local_force(), 0.);
+}
+
 void CellStructure::reset_local_properties() {
-  Kokkos::deep_copy(get_local_force(), 0);
+  Kokkos::deep_copy(get_local_force(), 0.);
 #ifdef ROTATION
-  Kokkos::deep_copy(get_local_torque(), 0);
+  Kokkos::deep_copy(get_local_torque(), 0.);
 #endif
 #ifdef NPT
-  Kokkos::deep_copy(get_local_virial(), 0);
+  Kokkos::deep_copy(get_local_virial(), 0.);
 #endif
 }
 
@@ -490,12 +488,12 @@ bool CellStructure::check_resort_required(
     Utils::Vector3d const &additional_offset) const {
   auto const lim = Utils::sqr(m_verlet_skin / 2.) - additional_offset.norm2();
 
-  Reduction::AddPartialResultKernel<bool> add_partial = [lim](Particle const &p,
-                                                              bool &result) {
-    if ((p.pos() - p.pos_at_last_verlet_update()).norm2() > lim) {
-      result = true;
-    }
-  };
+  Reduction::AddPartialResultKernel<bool> add_partial =
+      [lim](bool &result, Particle const &p) {
+        if ((p.pos() - p.pos_at_last_verlet_update()).norm2() > lim) {
+          result = true;
+        }
+      };
 
   Reduction::ReductionOp<bool> reduce_op = [](bool &acc, bool const &val) {
     acc |= val;

src/core/cell_system/CellStructure.hpp

@@ -729,8 +729,9 @@ class CellStructure : public System::Leaf<CellStructure> {
   auto get_max_id() const { return m_max_id; }
 
   void set_kokkos_handle(std::shared_ptr<Communication::KokkosHandle> handle);
-  void rebuild_local_properties(std::size_t num_threads, double pair_cutoff);
+  void rebuild_local_properties(double pair_cutoff);
   void reset_local_properties();
+  void reset_local_force();
 
   auto &get_local_force() { return *m_local_force; }
 #ifdef ROTATION
@@ -742,24 +743,24 @@ class CellStructure : public System::Leaf<CellStructure> {
   auto &get_aosoa() { return *m_aosoa; }
   auto const &get_unique_particles() const { return m_unique_particles; }
   auto const &get_verlet_list_cabana() const { return *m_verlet_list_cabana; }
+  void clear_local_properties();
 
   [[nodiscard]] auto is_verlet_list_cabana_rebuild_needed() const {
     return m_rebuild_verlet_list_cabana or (not use_verlet_list);
   }
 
   /**
    * @brief Reset local properties of the Verlet list.
-   * @param n_threads Number of threads.
    * @param cutoff    Pair interaction cutoff.
    * @return True if a rebuild is needed.
    */
-  [[nodiscard]] auto prepare_verlet_list_cabana(int n_threads, double cutoff) {
+  [[nodiscard]] auto prepare_verlet_list_cabana(double cutoff) {
     auto const rebuild = is_verlet_list_cabana_rebuild_needed();
     if (rebuild) {
       // If we have to rebuild, we need to count the particles
       set_index_map(); // parallelized index_map
       // Create essential variables for MD
-      rebuild_local_properties(n_threads, cutoff);
+      rebuild_local_properties(cutoff);
     } else {
       // If we do not rebuild we can use the saved map
       reset_local_properties();