FFTW: Initialize Threads

Co-authored-by: Severin Diederichs <severin.diederichs@desy.de>
ECP-WarpX · Jun 29, 2021 · f1e879d · f1e879d
1 parent d508049
commit f1e879d
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 1 deletion.
diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
@@ -145,6 +145,8 @@ SpectralFieldData::ForwardTransform (const int lev,
 #endif
 
     // Loop over boxes
+    // Note: we do NOT OpenMP parallelize here, since we use OpenMP threads for
+    //       the FFTs on each box!
     for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){
         if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers)
         {
@@ -247,6 +249,8 @@ SpectralFieldData::BackwardTransform( const int lev,
 #endif
 
     // Loop over boxes
+    // Note: we do NOT OpenMP parallelize here, since we use OpenMP threads for
+    //       the iFFTs on each box!
     for ( MFIter mfi(mf); mfi.isValid(); ++mfi ){
         if (cost && WarpX::load_balance_costs_update_algo == LoadBalanceCostsUpdateAlgo::Timers)
         {

diff --git a/Source/FieldSolver/SpectralSolver/WrapFFTW.cpp b/Source/FieldSolver/SpectralSolver/WrapFFTW.cpp
@@ -1,4 +1,4 @@
-/* Copyright 2019-2020
+/* Copyright 2019-2021
  *
  * This file is part of WarpX.
  *
@@ -32,6 +32,16 @@ namespace AnyFFT
     {
         FFTplan fft_plan;
 
+#if defined(AMREX_USE_OMP) && defined(WarpX_FFTW_OMP)
+#   ifdef AMREX_USE_FLOAT
+        fftwf_init_threads();
+        fftwf_plan_with_nthreads(omp_get_max_threads());
+#   else
+        fftw_init_threads();
+        fftw_plan_with_nthreads(omp_get_max_threads());
+#   endif
+#endif
+
         // Initialize fft_plan.m_plan with the vendor fft plan.
         // Swap dimensions: AMReX FAB are Fortran-order but FFTW is C-order
         if (dir == direction::R2C){