6 files changed, 93 insertions, 52 deletions
diff --git a/Source/FieldSolver/Make.package b/Source/FieldSolver/Make.package
index 7eea52ee7..25e21c146 100644
--- a/Source/FieldSolver/Make.package
+++ b/Source/FieldSolver/Make.package
@@ -2,9 +2,10 @@ CEXE_headers += WarpX_K.H
 CEXE_headers += WarpX_FDTD.H
 CEXE_sources += WarpXPushFieldsEM.cpp
 ifeq ($(USE_PSATD),TRUE)
-  CEXE_sources += WarpXFFT.cpp
-  F90EXE_sources += WarpX_fft.F90
   include $(WARPX_HOME)/Source/FieldSolver/SpectralSolver/Make.package
+  ifeq ($(USE_CUDA),FALSE)
+    include $(WARPX_HOME)/Source/FieldSolver/PicsarHybridSpectralSolver/Make.package
+  endif
 endif
 ifeq ($(USE_OPENBC_POISSON),TRUE)
   F90EXE_sources += openbc_poisson_solver.F90
diff --git a/Source/FieldSolver/PicsarHybridSpectralSolver/Make.package b/Source/FieldSolver/PicsarHybridSpectralSolver/Make.package
new file mode 100644
index 000000000..e185525e6
--- /dev/null
+++ b/Source/FieldSolver/PicsarHybridSpectralSolver/Make.package
@@ -0,0 +1,6 @@
+CEXE_headers += PicsarHybridFFTData.H
+CEXE_sources += PicsarHybridSpectralSolver.cpp
+F90EXE_sources += picsar_hybrid_spectral.F90
+
+INCLUDE_LOCATIONS += $(WARPX_HOME)/Source/FieldSolver/PicsarHybridSpectralSolver
+VPATH_LOCATIONS   += $(WARPX_HOME)/Source/FieldSolver/PicsarHybridSpectralSolver
+\ No newline at end of file
diff --git a/Source/FieldSolver/PicsarHybridSpectralSolver/PicsarHybridFFTData.H b/Source/FieldSolver/PicsarHybridSpectralSolver/PicsarHybridFFTData.H
new file mode 100644
index 000000000..4e97ab675
--- /dev/null
+++ b/Source/FieldSolver/PicsarHybridSpectralSolver/PicsarHybridFFTData.H
@@ -0,0 +1,32 @@
+#ifndef WARPX_PICSAR_HYBRID_FFTDATA_H_
+#define WARPX_PICSAR_HYBRID_FFTDATA_H_
+
+// FFTData is a stuct containing a 22 pointers to auxiliary arrays
+// 1-11: padded arrays in real space (required by FFTW); 12-22: arrays in spectral space
+struct FFTData
+{
+    static constexpr int N = 22;
+    void* data[N] = { nullptr };
+
+    ~FFTData () {
+        for (int i = 0; i < N; ++i) { // The memory is allocated with fftw_alloc.
+            fftw_free(data[i]);
+            data[i] = nullptr;
+        }
+    }
+
+    FFTData () = default;
+
+    FFTData (FFTData && rhs) noexcept {
+        for (int i = 0; i < N; ++i) {
+            data[i] = rhs.data[i];
+            rhs.data[i] = nullptr;
+        }
+    }
+
+    FFTData (FFTData const&) = delete;
+    void operator= (FFTData const&) = delete;
+    void operator= (FFTData&&) = delete;
+};
+
+#endif // WARPX_PICSAR_HYBRID_FFTDATA_H_
+\ No newline at end of file
diff --git a/Source/FieldSolver/WarpXFFT.cpp b/Source/FieldSolver/PicsarHybridSpectralSolver/PicsarHybridSpectralSolver.cpp
index f9cd7570a..26c93086a 100644
--- a/Source/FieldSolver/WarpXFFT.cpp
+++ b/Source/FieldSolver/PicsarHybridSpectralSolver/PicsarHybridSpectralSolver.cpp
@@ -5,10 +5,10 @@
 
 using namespace amrex;
 
-constexpr int WarpX::FFTData::N;
+constexpr int FFTData::N;
 
 namespace {
-static std::unique_ptr<WarpX::FFTData> nullfftdata; // This for process with nz_fft=0
+static std::unique_ptr<FFTData> nullfftdata; // This for process with nz_fft=0
 
 /** \brief Returns an "owner mask" which 1 for all cells, except
  *  for the duplicated (physical) cells of a nodal grid.
@@ -122,6 +122,8 @@ WarpX::AllocLevelDataFFT (int lev)
     static_assert(std::is_standard_layout<FFTData>::value, "FFTData must have standard layout");
     static_assert(sizeof(FFTData) == sizeof(void*)*FFTData::N, "sizeof FFTData is wrong");
 
+
+
     InitFFTComm(lev);
 
     BoxArray ba_fp_fft;
@@ -366,52 +368,8 @@ WarpX::FreeFFT (int lev)
 }
 
 void
-WarpX::PushPSATD (amrex::Real a_dt)
-{
-    for (int lev = 0; lev <= finest_level; ++lev) {
-        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(dt[lev] == a_dt, "dt must be consistent");
-        if (fft_hybrid_mpi_decomposition){
-            PushPSATD_hybridFFT(lev, a_dt);
-        } else {
-            PushPSATD_localFFT(lev, a_dt);
-        }
-    }
-}
-
-void WarpX::PushPSATD_localFFT (int lev, amrex::Real /* dt */)
-{
-    auto& solver = *spectral_solver_fp[lev];
-
-    // Perform forward Fourier transform
-    solver.ForwardTransform(*Efield_fp[lev][0], SpectralFieldIndex::Ex);
-    solver.ForwardTransform(*Efield_fp[lev][1], SpectralFieldIndex::Ey);
-    solver.ForwardTransform(*Efield_fp[lev][2], SpectralFieldIndex::Ez);
-    solver.ForwardTransform(*Bfield_fp[lev][0], SpectralFieldIndex::Bx);
-    solver.ForwardTransform(*Bfield_fp[lev][1], SpectralFieldIndex::By);
-    solver.ForwardTransform(*Bfield_fp[lev][2], SpectralFieldIndex::Bz);
-    solver.ForwardTransform(*current_fp[lev][0], SpectralFieldIndex::Jx);
-    solver.ForwardTransform(*current_fp[lev][1], SpectralFieldIndex::Jy);
-    solver.ForwardTransform(*current_fp[lev][2], SpectralFieldIndex::Jz);
-    solver.ForwardTransform(*rho_fp[lev], SpectralFieldIndex::rho_old, 0);
-    solver.ForwardTransform(*rho_fp[lev], SpectralFieldIndex::rho_new, 1);
-
-    // Advance fields in spectral space
-    solver.pushSpectralFields();
-
-    // Perform backward Fourier Transform
-    solver.BackwardTransform(*Efield_fp[lev][0], SpectralFieldIndex::Ex);
-    solver.BackwardTransform(*Efield_fp[lev][1], SpectralFieldIndex::Ey);
-    solver.BackwardTransform(*Efield_fp[lev][2], SpectralFieldIndex::Ez);
-    solver.BackwardTransform(*Bfield_fp[lev][0], SpectralFieldIndex::Bx);
-    solver.BackwardTransform(*Bfield_fp[lev][1], SpectralFieldIndex::By);
-    solver.BackwardTransform(*Bfield_fp[lev][2], SpectralFieldIndex::Bz);
-}
-
-void
 WarpX::PushPSATD_hybridFFT (int lev, amrex::Real /* dt */)
 {
-#ifndef AMREX_USE_CUDA // Running on CPU ; use PICSAR code for the hybrid FFT
-
     BL_PROFILE_VAR_NS("WarpXFFT::CopyDualGrid", blp_copy);
     BL_PROFILE_VAR_NS("PICSAR::FftPushEB", blp_push_eb);
 
@@ -486,8 +444,4 @@ WarpX::PushPSATD_hybridFFT (int lev, amrex::Real /* dt */)
     {
         amrex::Abort("WarpX::PushPSATD: TODO");
     }
-#else // AMREX_USE_CUDA is defined ; running on GPU
-        amrex::Abort("The option `psatd.fft_hybrid_mpi_decomposition` does not work on GPU.");
-#endif
-
 }
diff --git a/Source/FieldSolver/WarpX_fft.F90 b/Source/FieldSolver/PicsarHybridSpectralSolver/picsar_hybrid_spectral.F90
index 35357a63f..35357a63f 100644
--- a/Source/FieldSolver/WarpX_fft.F90
+++ b/Source/FieldSolver/PicsarHybridSpectralSolver/picsar_hybrid_spectral.F90
diff --git a/Source/FieldSolver/WarpXPushFieldsEM.cpp b/Source/FieldSolver/WarpXPushFieldsEM.cpp
index c53e13f8f..bea008598 100644
--- a/Source/FieldSolver/WarpXPushFieldsEM.cpp
+++ b/Source/FieldSolver/WarpXPushFieldsEM.cpp
@@ -16,6 +16,54 @@
 
 using namespace amrex;
 
+#ifdef WARPX_USE_PSATD
+
+void
+WarpX::PushPSATD (amrex::Real a_dt)
+{
+    for (int lev = 0; lev <= finest_level; ++lev) {
+        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(dt[lev] == a_dt, "dt must be consistent");
+        if (fft_hybrid_mpi_decomposition){
+#ifndef AMREX_USE_CUDA // Only available on CPU
+            PushPSATD_hybridFFT(lev, a_dt);
+#endif
+        } else {
+            PushPSATD_localFFT(lev, a_dt);
+        }
+    }
+}
+
+void WarpX::PushPSATD_localFFT (int lev, amrex::Real /* dt */)
+{
+    auto& solver = *spectral_solver_fp[lev];
+
+    // Perform forward Fourier transform
+    solver.ForwardTransform(*Efield_fp[lev][0], SpectralFieldIndex::Ex);
+    solver.ForwardTransform(*Efield_fp[lev][1], SpectralFieldIndex::Ey);
+    solver.ForwardTransform(*Efield_fp[lev][2], SpectralFieldIndex::Ez);
+    solver.ForwardTransform(*Bfield_fp[lev][0], SpectralFieldIndex::Bx);
+    solver.ForwardTransform(*Bfield_fp[lev][1], SpectralFieldIndex::By);
+    solver.ForwardTransform(*Bfield_fp[lev][2], SpectralFieldIndex::Bz);
+    solver.ForwardTransform(*current_fp[lev][0], SpectralFieldIndex::Jx);
+    solver.ForwardTransform(*current_fp[lev][1], SpectralFieldIndex::Jy);
+    solver.ForwardTransform(*current_fp[lev][2], SpectralFieldIndex::Jz);
+    solver.ForwardTransform(*rho_fp[lev], SpectralFieldIndex::rho_old, 0);
+    solver.ForwardTransform(*rho_fp[lev], SpectralFieldIndex::rho_new, 1);
+
+    // Advance fields in spectral space
+    solver.pushSpectralFields();
+
+    // Perform backward Fourier Transform
+    solver.BackwardTransform(*Efield_fp[lev][0], SpectralFieldIndex::Ex);
+    solver.BackwardTransform(*Efield_fp[lev][1], SpectralFieldIndex::Ey);
+    solver.BackwardTransform(*Efield_fp[lev][2], SpectralFieldIndex::Ez);
+    solver.BackwardTransform(*Bfield_fp[lev][0], SpectralFieldIndex::Bx);
+    solver.BackwardTransform(*Bfield_fp[lev][1], SpectralFieldIndex::By);
+    solver.BackwardTransform(*Bfield_fp[lev][2], SpectralFieldIndex::Bz);
+}
+
+#endif
+
 void
 WarpX::EvolveB (Real a_dt)
 {