p3m: Use OpenMP kernels

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

Author: jngrad

cmake/FindFFTW3.cmake

@@ -33,25 +33,28 @@ endif(FFTW3_INCLUDE_DIR)
 find_path(FFTW3_INCLUDE_DIR fftw3.h)
 find_library(FFTW3_LIBRARIES fftw3)
 find_library(FFTW3F_LIBRARIES fftw3f)
-find_path(FFTW3_MPI_INCLUDE_DIR fftw3-mpi.h)
-find_library(FFTW3_MPI_LIBRARIES fftw3_mpi)
-find_library(FFTW3F_MPI_LIBRARIES fftw3f_mpi)
+find_library(FFTW3_OMP_LIBRARIES fftw3_omp)
+find_library(FFTW3F_OMP_LIBRARIES fftw3f_omp)
 
 # handle the QUIETLY and REQUIRED arguments and set FFTW3_FOUND to TRUE if all
 # listed variables are TRUE
 include(FindPackageHandleStandardArgs)
 find_package_handle_standard_args(FFTW3 DEFAULT_MSG FFTW3_LIBRARIES FFTW3F_LIBRARIES
                                   FFTW3_INCLUDE_DIR)
 set(FPHSA_NAME_MISMATCHED 1)
-find_package_handle_standard_args(FFTW3_MPI DEFAULT_MSG FFTW3_MPI_LIBRARIES FFTW3F_MPI_LIBRARIES
-                                  FFTW3_MPI_INCLUDE_DIR)
+find_package_handle_standard_args(FFTW3_OMP DEFAULT_MSG FFTW3_OMP_LIBRARIES FFTW3F_OMP_LIBRARIES)
 unset(FPHSA_NAME_MISMATCHED)
 
-mark_as_advanced(FFTW3_LIBRARIES FFTW3F_LIBRARIES FFTW3_INCLUDE_DIR FFTW3_MPI_LIBRARIES FFTW3F_MPI_LIBRARIES FFTW3_MPI_INCLUDE_DIR)
+mark_as_advanced(FFTW3_LIBRARIES FFTW3F_LIBRARIES FFTW3_INCLUDE_DIR FFTW3_OMP_LIBRARIES FFTW3F_OMP_LIBRARIES)
 
 
 if(FFTW3_FOUND AND NOT TARGET FFTW3::FFTW3)
   add_library(FFTW3::FFTW3 INTERFACE IMPORTED)
   target_include_directories(FFTW3::FFTW3 INTERFACE "${FFTW3_INCLUDE_DIR}")
   target_link_libraries(FFTW3::FFTW3 INTERFACE "${FFTW3_LIBRARIES}" "${FFTW3F_LIBRARIES}")
 endif()
+if(FFTW3_FOUND AND NOT TARGET FFTW3::FFTW3_OMP)
+  add_library(FFTW3::FFTW3_OMP INTERFACE IMPORTED)
+  target_include_directories(FFTW3::FFTW3_OMP INTERFACE "${FFTW3_INCLUDE_DIR}")
+  target_link_libraries(FFTW3::FFTW3_OMP INTERFACE "${FFTW3_OMP_LIBRARIES}" "${FFTW3F_OMP_LIBRARIES}")
+endif()

maintainer/benchmarks/p3m.py

@@ -22,6 +22,7 @@
 import benchmarks
 import numpy as np
 import argparse
+import os
 
 parser = argparse.ArgumentParser(description="Benchmark P3M simulations. "
                                  "Save the results to a CSV file.")
@@ -83,7 +84,8 @@
 # System parameters
 #############################################################
 n_proc = system.cell_system.get_state()["n_nodes"]
-n_part = n_proc * args.particles_per_core
+n_threads = int(os.environ.get("OMP_NUM_THREADS", 1))
+n_part = n_proc * n_threads * args.particles_per_core
 # volume of N spheres with radius r: N * (4/3*pi*r^3)
 lj_sig = (lj_sigmas["cation"] + lj_sigmas["anion"]) / 2
 box_l = (n_part * 4. / 3. * np.pi * (lj_sig / 2.)**3
@@ -137,7 +139,7 @@
 # tuning and equilibration
 min_skin = 0.2
 max_skin = 1.6
-p3m_params = {"prefactor": args.prefactor, "accuracy": 1e-3}
+p3m_params = {"prefactor": args.prefactor, "accuracy": 1e-3, "timings": 100}
 p3m = p3m_class(**p3m_params)
 print("Quick equilibration")
 system.integrator.run(min(3 * measurement_steps, 1000))
@@ -156,6 +158,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,12 +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);
-    flat_array_t[index_out] = flat_array[index_in];
+  std::vector<T> flat_array_t;
+  flat_array_t.reserve(flat_array.size());
+  for_each_3d_lin<order_out>(mesh_start, shape, indices, index_out, [&]() {
+    auto const index_in = Utils::get_linear_index<order_in>(indices, shape);
+    assert(index_out == Utils::get_linear_index<order_out>(indices, shape));
+    flat_array_t.push_back(flat_array[index_in]);
   });
   return flat_array_t;
 }
@@ -451,10 +453,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 +469,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 +488,33 @@ 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 = 0;
+  for_each_3d_lin<Utils::MemoryOrder::COLUMN_MAJOR>(
+      mesh_start, mesh_stop, indices, local_index, [&]() {
+        assert(local_index ==
+               Utils::get_linear_index<Utils::MemoryOrder::COLUMN_MAJOR>(
+                   indices - mesh_start, mesh_stop));
+        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

@@ -18,6 +18,6 @@
 #
 
 if(ESPRESSO_BUILD_WITH_FFTW)
-  target_link_libraries(espresso_p3m PRIVATE FFTW3::FFTW3)
-  target_sources(espresso_p3m PRIVATE fft.cpp)
+  target_link_libraries(espresso_p3m PRIVATE FFTW3::FFTW3 FFTW3::FFTW3_OMP)
+  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/magnetostatics/CMakeLists.txt

@@ -22,5 +22,11 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/dipolar_direct_sum.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/dipolar_direct_sum_gpu.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/dlc.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/dp3m.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/scafacos_impl.cpp)
+
+if(ESPRESSO_BUILD_WITH_FFTW)
+  target_sources(espresso_p3m PRIVATE dp3m.cpp)
+  set_source_files_properties(
+    dp3m.cpp TARGET_DIRECTORY espresso_p3m
+    PROPERTIES INCLUDE_DIRECTORIES ${CMAKE_SOURCE_DIR}/src/core)
+endif()

src/core/npt.cpp

@@ -20,7 +20,9 @@
 
 #ifdef NPT
 
+#include "BoxGeometry.hpp"
 #include "PropagationMode.hpp"
+#include "cell_system/CellStructure.hpp"
 #include "communication.hpp"
 #include "config/config.hpp"
 #include "electrostatics/coulomb.hpp"

src/core/p3m/CMakeLists.txt

@@ -37,8 +37,15 @@ if(ESPRESSO_BUILD_WITH_FFTW)
   target_link_libraries(
     espresso_p3m
     PRIVATE
-      espresso::cpp_flags espresso::p3m::cpp_flags espresso::instrumentation
-      espresso::utils espresso::config Boost::mpi MPI::MPI_CXX Heffte::Heffte
+      espresso::cpp_flags
+      espresso::p3m::cpp_flags
+      espresso::instrumentation
+      espresso::utils
+      espresso::config
+      Boost::mpi
+      MPI::MPI_CXX
+      Heffte::Heffte
+      $<$<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>)
   set_target_properties(espresso_p3m PROPERTIES POSITION_INDEPENDENT_CODE ON)