p3m: Use OpenMP kernels

Co-authored-by: Rudolf Weeber <[email protected]>

Author: jngrad

maintainer/benchmarks/benchmarks.py

@@ -16,6 +16,7 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
+import os
 import sys
 import time
 import pathlib
@@ -108,7 +109,7 @@ def get_average_time(timings):
     return (avg, ci)
 
 
-def write_report(filepath, n_proc, timings, n_steps, label=''):
+def write_report(filepath, n_ranks, timings, n_steps, label=''):
     '''
     Append timing data to a CSV file. If it doesn't exist, it is created
     with a header.
@@ -117,7 +118,7 @@ def write_report(filepath, n_proc, timings, n_steps, label=''):
     ----------
     filepath: :obj:`str`
         Path to the CSV file.
-    n_proc: :obj:`int`
+    n_ranks: :obj:`int`
         Number of MPI ranks.
     timings: :obj:`ndarray` of :obj:`float`
         Timings.
@@ -127,11 +128,12 @@ def write_report(filepath, n_proc, timings, n_steps, label=''):
         Label to distinguish e.g. MD from MC or LB steps.
 
     '''
+    n_threads = int(os.environ.get("OMP_NUM_THREADS", 1))
     script = pathlib.Path(sys.argv[0]).name
     cmd = " ".join(x for x in sys.argv[1:] if not x.startswith("--output"))
     avg, ci = get_average_time(timings)
-    header = '"script","arguments","cores","mean","ci","nsteps","duration","label"\n'
-    report = f'"{script}","{cmd}",{n_proc},{avg:.3e},{ci:.3e},{n_steps},{np.sum(timings):.1f},"{label}"\n'  # nopep8
+    header = '"script","arguments","ranks","threads","mean","ci","nsteps","duration","label"\n'
+    report = f'"{script}","{cmd}",{n_ranks},{n_threads},{avg:.3e},{ci:.3e},{n_steps},{np.sum(timings):.1f},"{label}"\n'  # nopep8
     if pathlib.Path(filepath).is_file():
         header = ''
     with open(filepath, "a") as f:

maintainer/benchmarks/p3m.py

@@ -156,6 +156,9 @@
 print("Re-tune p3m")
 p3m = p3m_class(**p3m_params)
 system.electrostatics.solver = p3m
+print("Tune skin: {:.3f}".format(system.cell_system.tune_skin(
+    min_skin=min_skin, max_skin=max_skin, tol=0.05, int_steps=100,
+    adjust_max_skin=True)))
 
 
 if args.visualizer:

src/core/communication.cpp

@@ -25,6 +25,7 @@
 
 #include "cuda/init.hpp"
 #include "errorhandling.hpp"
+#include "fft/init.hpp"
 
 #ifdef WALBERLA
 #include <walberla_bridge/walberla_init.hpp>
@@ -90,6 +91,10 @@ void init(std::shared_ptr<boost::mpi::environment> mpi_env) {
   cuda_on_program_start();
 #endif
 
+#ifdef FFTW
+  fft_on_program_start();
+#endif
+
 #ifdef SHARED_MEMORY_PARALLELISM
   Kokkos::initialize();
 #endif

src/core/electrostatics/p3m.cpp

@@ -117,11 +117,14 @@ FloatType complex_norm2(std::complex<FloatType> const &z) {
 template <Utils::MemoryOrder order_in, Utils::MemoryOrder order_out, typename T>
 auto transpose(std::span<T> const &flat_array, Utils::Vector3i const &shape) {
   auto constexpr mesh_start = Utils::Vector3i::broadcast(0);
+  std::size_t index_out{};
   Utils::Vector3i indices{};
   std::vector<T> flat_array_t(flat_array.size());
-  for_each_3d(mesh_start, shape, indices, [&]() {
-    auto const index_in = Utils::get_linear_index(indices, shape, order_in);
-    auto const index_out = Utils::get_linear_index(indices, shape, order_out);
+  for_each_3d_lin<order_out>(mesh_start, shape, indices, index_out, [&]() {
+    auto const index_in = Utils::get_linear_index<order_in>(indices, shape);
+#ifdef ADDITIONAL_CHECKS
+    assert(index_out == Utils::get_linear_index<order_out>(indices, shape));
+#endif
     flat_array_t[index_out] = flat_array[index_in];
   });
   return flat_array_t;
@@ -451,10 +454,11 @@ void CoulombP3MImpl<FloatType, Architecture>::kernel_ks_charge_density() {
                             p3m.local_mesh.dim);
 
   // get real space charge density without ghost layers
-  auto charge_density_no_halos = extract_block(
-      p3m.rs_charge_density, p3m.local_mesh.dim, p3m.local_mesh.n_halo_ld,
-      p3m.local_mesh.dim - p3m.local_mesh.n_halo_ur,
-      Utils::MemoryOrder::ROW_MAJOR, Utils::MemoryOrder::COLUMN_MAJOR);
+  auto charge_density_no_halos =
+      extract_block<Utils::MemoryOrder::ROW_MAJOR,
+                    Utils::MemoryOrder::COLUMN_MAJOR>(
+          p3m.rs_charge_density, p3m.local_mesh.dim, p3m.local_mesh.n_halo_ld,
+          p3m.local_mesh.dim - p3m.local_mesh.n_halo_ur);
 
   // Set up the FFT using the Heffte library.
   // This is in global mesh coordinates without any ghost layers
@@ -466,14 +470,14 @@ void CoulombP3MImpl<FloatType, Architecture>::kernel_ks_charge_density() {
 
 template <typename FloatType, Arch Architecture>
 void CoulombP3MImpl<FloatType, Architecture>::kernel_rs_electric_field() {
-  auto constexpr mesh_start = Utils::Vector3i::broadcast(0);
-  auto const &mesh_stop = p3m.fft->ks_local_size();
+  auto const mesh_start = p3m.fft->ks_local_ld_index();
+  auto const mesh_stop = mesh_start + p3m.fft->ks_local_size();
   auto const &box_geo = *get_system().box_geo;
-  auto indices = Utils::Vector3i{};
+  Utils::Vector3i indices{};
 
   // hold electric field in k-space
   std::array<std::span<std::complex<FloatType>>, 3> ks_E_fields;
-  auto const fft_mesh_length = get_size_from_shape(mesh_stop);
+  auto const fft_mesh_length = get_size_from_shape(mesh_stop - mesh_start);
   for (auto d : {0u, 1u, 2u}) {
     auto const offset = d * fft_mesh_length;
     auto const begin = p3m.ks_E_fields_storage.begin() + offset;
@@ -485,30 +489,35 @@ void CoulombP3MImpl<FloatType, Architecture>::kernel_rs_electric_field() {
       Utils::Vector3<FloatType>((2. * std::numbers::pi) * box_geo.length_inv());
 
   // compute electric field, Eq. (3.49) @cite deserno00b
-  for_each_3d(mesh_start, mesh_stop, indices, [&]() {
-    auto const global_index = indices + p3m.fft->ks_local_ld_index();
-    auto const local_index = Utils::get_linear_index(
-        indices, mesh_stop, Utils::MemoryOrder::COLUMN_MAJOR);
-    auto const phi_hat = multiply_complex_by_real(
-        p3m.ks_charge_density[local_index], p3m.g_force[local_index]);
-
-    for (auto d : {0u, 1u, 2u}) {
-      // wave vector of the current mesh point
-      auto const k = FloatType(p3m.d_op[d][global_index[d]]) * wavevector[d];
-      // electric field in k-space
-      ks_E_fields[d][local_index] = multiply_complex_by_imaginary(phi_hat, k);
-    }
-  });
+  std::size_t local_index{};
+  for_each_3d_lin<Utils::MemoryOrder::COLUMN_MAJOR>(
+      mesh_start, mesh_stop, indices, local_index, [&]() {
+#ifdef ADDITIONAL_CHECKS
+        assert(local_index ==
+               Utils::get_linear_index<Utils::MemoryOrder::COLUMN_MAJOR>(
+                   indices - mesh_start, mesh_stop - mesh_start));
+#endif
+        auto const phi_hat = multiply_complex_by_real(
+            p3m.ks_charge_density[local_index], p3m.g_force[local_index]);
+
+        for (auto d : {0u, 1u, 2u}) {
+          // wave vector of the current mesh point
+          auto const k = FloatType(p3m.d_op[d][indices[d]]) * wavevector[d];
+          // electric field in k-space
+          ks_E_fields[d][local_index] =
+              multiply_complex_by_imaginary(phi_hat, k);
+        }
+      });
 
   // back-transform the k-space electric field to real space
-  auto const rs_mesh_size_no_halo =
-      get_size_from_shape(p3m.local_mesh.dim_no_halo);
-  p3m.fft->backward_batch(3, p3m.ks_E_fields_storage.data(),
-                          p3m.rs_E_fields_no_halo.data());
-
-  // add zeros around the E-field in real space to make room for ghost layers
   auto const size = p3m.local_mesh.ur_no_halo - p3m.local_mesh.ld_no_halo;
+  auto const rs_mesh_size_no_halo = Utils::product(size);
   for (auto d : {0u, 1u, 2u}) {
+    auto k_space = ks_E_fields[d].data();
+    auto real_space = p3m.rs_E_fields_no_halo.data() + d * rs_mesh_size_no_halo;
+    p3m.fft->backward(k_space, real_space);
+
+    // add zeros around the E-field in real space to make room for ghost layers
     auto const offset = d * rs_mesh_size_no_halo;
     auto const begin = p3m.rs_E_fields_no_halo.begin() + offset;
     auto f = std::span<std::complex<FloatType>>(begin, rs_mesh_size_no_halo);

src/core/fft/CMakeLists.txt

@@ -21,5 +21,5 @@ if(ESPRESSO_BUILD_WITH_FFTW)
   target_link_libraries(
     espresso_p3m PRIVATE $<$<BOOL:${fftw3_FOUND}>:espresso::fftw3>
                          $<$<BOOL:${fftw3_omp_FOUND}>:espresso::fftw3_omp>)
-  target_sources(espresso_p3m PRIVATE fft.cpp)
+  target_sources(espresso_p3m PRIVATE fft.cpp init.cpp)
 endif()

src/core/fft/fft.cpp

@@ -40,6 +40,10 @@
 
 #include <fftw3.h>
 
+#ifdef _OPENMP
+#include <omp.h>
+#endif
+
 #include <algorithm>
 #include <cmath>
 #include <cstddef>
@@ -677,27 +681,39 @@ int fft_data_struct<FloatType>::initialize_fft(
   /* === FFT Routines (Using FFTW / RFFTW package)=== */
   for (int i = 1; i < 4; i++) {
     if (init_tag) {
+#ifdef _OPENMP
+#pragma omp critical(fftw_destroy_plan_forward)
       forw[i].destroy_plan();
+#endif
     }
     forw[i].dir = FFTW_FORWARD;
+#ifdef _OPENMP
+#pragma omp critical(fftw_create_plan_forward)
     forw[i].plan_handle = fftw<FloatType>::plan_many_dft(
         1, &forw[i].new_mesh[2], forw[i].n_ffts, c_data, nullptr, 1,
         forw[i].new_mesh[2], c_data, nullptr, 1, forw[i].new_mesh[2],
         forw[i].dir, FFTW_PATIENT);
+#endif
     assert(forw[i].plan_handle);
   }
 
   /* === The BACK Direction === */
   /* this is needed because slightly different functions are used */
   for (int i = 1; i < 4; i++) {
     if (init_tag) {
+#ifdef _OPENMP
+#pragma omp critical(fftw_destroy_plan_backward)
       back[i].destroy_plan();
+#endif
     }
     back[i].dir = FFTW_BACKWARD;
+#ifdef _OPENMP
+#pragma omp critical(fftw_create_plan_backward)
     back[i].plan_handle = fftw<FloatType>::plan_many_dft(
         1, &forw[i].new_mesh[2], forw[i].n_ffts, c_data, nullptr, 1,
         forw[i].new_mesh[2], c_data, nullptr, 1, forw[i].new_mesh[2],
         back[i].dir, FFTW_PATIENT);
+#endif
     back[i].pack_function = pack_block_permute1;
     assert(back[i].plan_handle);
   }

src/core/fft/init.cpp

@@ -0,0 +1,46 @@
+/*
+ * 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/>.
+ */
+
+#include "init.hpp"
+
+#include "config/config.hpp"
+
+#include <cassert>
+
+#include <fftw3.h>
+
+#ifdef _OPENMP
+#include <omp.h>
+#endif
+
+#ifdef FFTW
+void fft_on_program_start() {
+#ifdef _OPENMP
+  int omp_num_threads = 1;
+#pragma omp parallel
+  {
+#pragma omp single
+    omp_num_threads = omp_get_num_threads();
+  }
+  [[maybe_unused]] auto const init_status_success = fftw_init_threads();
+  assert(init_status_success);
+  fftw_plan_with_nthreads(omp_num_threads);
+#endif // _OPENMP
+}
+#endif // FFTW

src/core/fft/init.hpp

@@ -0,0 +1,26 @@
+/*
+ * 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
+
+#include "config/config.hpp"
+
+#ifdef FFTW
+void fft_on_program_start();
+#endif

src/core/p3m/P3MFFT.hpp

@@ -24,9 +24,11 @@
 #include <boost/mpi/communicator.hpp>
 
 #include <heffte.h>
+#include <heffte_backends.h>
 
 #include <algorithm>
 #include <array>
+#include <initializer_list>
 #include <memory>
 
 template <typename T, std::size_t N>
@@ -45,6 +47,7 @@ template <typename FloatType> class P3MFFT {
   Utils::Vector3i m_global_mesh;
   std::shared_ptr<Box> in_box;
   std::shared_ptr<Box> out_box;
+  std::vector<std::complex<FloatType>> m_buffer;
   heffte::fft3d<backend_tag> fft3d;
 
 public:
@@ -69,8 +72,16 @@ template <typename FloatType> class P3MFFT {
     auto const global_box = heffte::box3d<>(
         {0, 0, 0}, to_array(m_global_mesh - Utils::Vector3i::broadcast(1)),
         to_array(m_memory_layout));
+    auto const n_procs = Utils::product(node_grid);
+    auto best_grid = node_grid;
+    for (auto i : {0u, 1u, 2u}) {
+      if (m_global_mesh[i] % 2 * n_procs == 0) {
+        best_grid = {n_procs, 1, 1};
+        break;
+      }
+    }
     auto all_boxes = heffte::split_world(
-        global_box, {node_grid[2], node_grid[1], node_grid[0]});
+        global_box, {best_grid[0], best_grid[1], best_grid[2]});
     out_box = std::make_shared<Box>(all_boxes[comm.rank()]);
     init_fft();
   }
@@ -95,8 +106,9 @@ template <typename FloatType> class P3MFFT {
     // 1-D pencils for sufficiently large problem, it is expected that the
     // pencil decomposition is better but for smaller problems, the slabs may
     // perform better (depending on hardware and backend)
-    options.use_pencils = false;
+    options.use_pencils = true;
     fft3d = heffte::fft3d<backend_tag>(*in_box, *out_box, comm, options);
+    m_buffer.resize(fft3d.size_workspace());
   }
 
   Utils::Vector3i ks_local_ld_index() const {
@@ -108,14 +120,11 @@ template <typename FloatType> class P3MFFT {
   Utils::Vector3i ks_local_size() const {
     return ks_local_ur_index() - ks_local_ld_index();
   }
-  template <typename T> auto forward(T &in) { return fft3d.forward(in); }
   template <typename In, typename Out> void forward(In in, Out out) {
-    fft3d.forward(in, out);
+    fft3d.forward(in, out, m_buffer.data());
   }
-  template <typename T> auto backward(T &in) { return fft3d.backward(in); }
-  template <typename T1, typename T2>
-  auto backward_batch(int n, T1 in, T2 out) {
-    return fft3d.backward(n, in, out);
+  template <typename T1, typename T2> auto backward(T1 &in, T2 &out) {
+    return fft3d.backward(in, out, m_buffer.data());
   }
   auto const &get_memory_layout() const { return m_memory_layout; }
 };