Encapsulate FFTs (#1055)

* wrap fft libraries

* implementation in cpp to avoid duplicate symbole

* delete some fftw code

* further cleaning

* typo

* pass fft plans by reference

* fix bug due to typo. Dammit, macros sometimes make it hard

* FFT wrapper also support cuFFT (not tested yet)

* eol

* further cleaning

* fix cuFFT, tested on Summit

* clean WarpX Complex checks

* pass directly the IntVect instead of the components (what were you thinking?)

* add some amrex prefix

* add a few const

* Should not need an FFT library to compile FDTD

* gather all FFT header files into one.

* eol

* minor changes in code comments.

* let dimension be chosen at runtime instead of at compile-time

* fix compilation on GPU

* add description to the wrapper namespace

1 files changed, 14 insertions, 204 deletions

diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
index 76299b7de..0f49e695b 100644
--- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
+++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp
@@ -9,26 +9,10 @@
 
 #include <map>
 
-
-
 #if WARPX_USE_PSATD
 
 using namespace amrex;
 
-#ifdef AMREX_USE_GPU
-#  ifdef AMREX_USE_FLOAT
-using cuPrecisionComplex = cuComplex;
-#  else
-using cuPrecisionComplex = cuDoubleComplex;
-#  endif
-#else
-#  ifdef AMREX_USE_FLOAT
-using fftw_precision_complex = fftwf_complex;
-#  else
-using fftw_precision_complex = fftw_complex;
-#  endif
-#endif
-
 /* \brief Initialize fields in spectral space, and FFT plans */
 SpectralFieldData::SpectralFieldData( const amrex::BoxArray& realspace_ba,
                                       const SpectralKSpace& k_space,
@@ -73,8 +57,8 @@ SpectralFieldData::SpectralFieldData( const amrex::BoxArray& realspace_ba,
 #endif
 
     // Allocate and initialize the FFT plans
-    forward_plan = FFTplans(spectralspace_ba, dm);
-    backward_plan = FFTplans(spectralspace_ba, dm);
+    forward_plan = AnyFFT::FFTplans(spectralspace_ba, dm);
+    backward_plan = AnyFFT::FFTplans(spectralspace_ba, dm);
     // Loop over boxes and allocate the corresponding plan
     // for each box owned by the local MPI proc
     for ( MFIter mfi(spectralspace_ba, dm); mfi.isValid(); ++mfi ){
@@ -82,98 +66,16 @@ SpectralFieldData::SpectralFieldData( const amrex::BoxArray& realspace_ba,
         // differ when using real-to-complex FFT. When initializing
         // the FFT plan, the valid dimensions are those of the real-space box.
         IntVect fft_size = realspace_ba[mfi].length();
-#ifdef AMREX_USE_GPU
-        // Create cuFFT plans
-        // Creating 3D plan for real to complex -- double precision
-        // Assuming CUDA is used for programming GPU
-        // Note that D2Z is inherently forward plan
-        // and  Z2D is inherently backward plan
-        cufftResult result;
-#  if (AMREX_SPACEDIM == 3)
-        result = cufftPlan3d( &forward_plan[mfi], fft_size[2], fft_size[1],fft_size[0],
-#    ifdef AMREX_USE_FLOAT
-                              CUFFT_R2C);
-#    else
-                              CUFFT_D2Z);
-#    endif
-        if ( result != CUFFT_SUCCESS ) {
-            amrex::Print() << " cufftplan3d forward failed! Error: " <<
-            cufftErrorToString(result) << "\n";
-        }
 
-        result = cufftPlan3d( &backward_plan[mfi], fft_size[2], fft_size[1],fft_size[0],
-#    ifdef AMREX_USE_FLOAT
-                              CUFFT_C2R);
-#    else
-                              CUFFT_Z2D);
-#    endif
-        if ( result != CUFFT_SUCCESS ) {
-           amrex::Print() << " cufftplan3d backward failed! Error: " <<
-            cufftErrorToString(result) << "\n";
-        }
-#  else
-        result = cufftPlan2d( &forward_plan[mfi], fft_size[1], fft_size[0],
-#    ifdef AMREX_USE_FLOAT
-                              CUFFT_R2C);
-#    else
-                              CUFFT_D2Z);
-#    endif
-        if ( result != CUFFT_SUCCESS ) {
-           amrex::Print() << " cufftplan2d forward failed! Error: " <<
-            cufftErrorToString(result) << "\n";
-        }
-
-        result = cufftPlan2d( &backward_plan[mfi], fft_size[1], fft_size[0],
-#    ifdef AMREX_USE_FLOAT
-                              CUFFT_C2R);
-#    else
-                              CUFFT_Z2D);
-#    endif
-        if ( result != CUFFT_SUCCESS ) {
-           amrex::Print() << " cufftplan2d backward failed! Error: " <<
-            cufftErrorToString(result) << "\n";
-        }
-#  endif
+        forward_plan[mfi] = AnyFFT::CreatePlan(
+            fft_size, tmpRealField[mfi].dataPtr(),
+            reinterpret_cast<AnyFFT::Complex*>( tmpSpectralField[mfi].dataPtr()),
+            AnyFFT::direction::R2C, AMREX_SPACEDIM);
 
-#else
-        // Create FFTW plans
-        forward_plan[mfi] =
-            // Swap dimensions: AMReX FAB are Fortran-order but FFTW is C-order
-#  if (AMREX_SPACEDIM == 3)
-#    ifdef AMREX_USE_FLOAT
-            fftwf_plan_dft_r2c_3d( fft_size[2], fft_size[1], fft_size[0],
-#    else
-            fftw_plan_dft_r2c_3d( fft_size[2], fft_size[1], fft_size[0],
-#    endif
-#  else
-#    ifdef AMREX_USE_FLOAT
-            fftwf_plan_dft_r2c_2d( fft_size[1], fft_size[0],
-#    else
-            fftw_plan_dft_r2c_2d( fft_size[1], fft_size[0],
-#    endif
-#  endif
-            tmpRealField[mfi].dataPtr(),
-            reinterpret_cast<fftw_precision_complex*>( tmpSpectralField[mfi].dataPtr() ),
-            FFTW_ESTIMATE );
-        backward_plan[mfi] =
-            // Swap dimensions: AMReX FAB are Fortran-order but FFTW is C-order
-#  if (AMREX_SPACEDIM == 3)
-#    ifdef AMREX_USE_FLOAT
-            fftwf_plan_dft_c2r_3d( fft_size[2], fft_size[1], fft_size[0],
-#    else
-            fftw_plan_dft_c2r_3d( fft_size[2], fft_size[1], fft_size[0],
-#    endif
-#  else
-#    ifdef AMREX_USE_FLOAT
-            fftwf_plan_dft_c2r_2d( fft_size[1], fft_size[0],
-#    else
-            fftw_plan_dft_c2r_2d( fft_size[1], fft_size[0],
-#    endif
-#  endif
-            reinterpret_cast<fftw_precision_complex*>( tmpSpectralField[mfi].dataPtr() ),
-            tmpRealField[mfi].dataPtr(),
-            FFTW_ESTIMATE );
-#endif
+        backward_plan[mfi] = AnyFFT::CreatePlan(
+            fft_size, tmpRealField[mfi].dataPtr(),
+            reinterpret_cast<AnyFFT::Complex*>( tmpSpectralField[mfi].dataPtr()),
+            AnyFFT::direction::C2R, AMREX_SPACEDIM);
     }
 }
 
@@ -182,20 +84,8 @@ SpectralFieldData::~SpectralFieldData()
 {
     if (tmpRealField.size() > 0){
         for ( MFIter mfi(tmpRealField); mfi.isValid(); ++mfi ){
-#ifdef AMREX_USE_GPU
-            // Destroy cuFFT plans
-            cufftDestroy( forward_plan[mfi] );
-            cufftDestroy( backward_plan[mfi] );
-#else
-            // Destroy FFTW plans
-#  ifdef AMREX_USE_FLOAT
-            fftwf_destroy_plan( forward_plan[mfi] );
-            fftwf_destroy_plan( backward_plan[mfi] );
-#  else
-            fftw_destroy_plan( forward_plan[mfi] );
-            fftw_destroy_plan( backward_plan[mfi] );
-#  endif
-#endif
+            AnyFFT::DestroyPlan(forward_plan[mfi]);
+            AnyFFT::DestroyPlan(backward_plan[mfi]);
         }
     }
 }
@@ -243,34 +133,7 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf,
         }
 
         // Perform Fourier transform from `tmpRealField` to `tmpSpectralField`
-#ifdef AMREX_USE_GPU
-        // Perform Fast Fourier Transform on GPU using cuFFT
-        // make sure that this is done on the same
-        // GPU stream as the above copy
-        cufftResult result;
-        cudaStream_t stream = amrex::Gpu::Device::cudaStream();
-        cufftSetStream ( forward_plan[mfi], stream);
-#  ifdef AMREX_USE_FLOAT
-        result = cufftExecR2C(
-#  else
-        result = cufftExecD2Z(
-#  endif
-            forward_plan[mfi],
-            tmpRealField[mfi].dataPtr(),
-            reinterpret_cast<cuPrecisionComplex*>(
-                tmpSpectralField[mfi].dataPtr()) );
-        if ( result != CUFFT_SUCCESS ) {
-           amrex::Print() <<
-           " forward transform using cufftExec failed ! Error: " <<
-           cufftErrorToString(result) << "\n";
-        }
-#else
-#  ifdef AMREX_USE_FLOAT
-        fftwf_execute( forward_plan[mfi] );
-#  else
-        fftw_execute( forward_plan[mfi] );
-#  endif
-#endif
+        AnyFFT::Execute(forward_plan[mfi]);
 
         // Copy the spectral-space field `tmpSpectralField` to the appropriate
         // index of the FabArray `fields` (specified by `field_index`)
@@ -358,34 +221,7 @@ SpectralFieldData::BackwardTransform( MultiFab& mf,
         }
 
         // Perform Fourier transform from `tmpSpectralField` to `tmpRealField`
-#ifdef AMREX_USE_GPU
-        // Perform Fast Fourier Transform on GPU using cuFFT.
-        // make sure that this is done on the same
-        // GPU stream as the above copy
-        cufftResult result;
-        cudaStream_t stream = amrex::Gpu::Device::cudaStream();
-        cufftSetStream ( backward_plan[mfi], stream);
-#  ifdef AMREX_USE_FLOAT
-        result = cufftExecC2R(
-#  else
-        result = cufftExecZ2D(
-#  endif
-            backward_plan[mfi],
-            reinterpret_cast<cuPrecisionComplex*>(
-            tmpSpectralField[mfi].dataPtr()),
-            tmpRealField[mfi].dataPtr() );
-        if ( result != CUFFT_SUCCESS ) {
-           amrex::Print() <<
-           " Backward transform using cufftexec failed! Error: " <<
-           cufftErrorToString(result) << "\n";
-        }
-#else
-#  ifdef AMREX_USE_FLOAT
-        fftwf_execute( backward_plan[mfi] );
-#  else
-        fftw_execute( backward_plan[mfi] );
-#  endif
-#endif
+        AnyFFT::Execute(backward_plan[mfi]);
 
         // Copy the temporary field `tmpRealField` to the real-space field `mf`
         // (only in the valid cells ; not in the guard cells)
@@ -430,30 +266,4 @@ SpectralFieldData::BackwardTransform( MultiFab& mf,
     }
 }
 
-#ifdef AMREX_USE_GPU
-std::string
-SpectralFieldData::cufftErrorToString (const cufftResult& err)
-{
-    const auto res2string = std::map<cufftResult, std::string>{
-        {CUFFT_SUCCESS, "CUFFT_SUCCESS"},
-        {CUFFT_INVALID_PLAN,"CUFFT_INVALID_PLAN"},
-        {CUFFT_ALLOC_FAILED,"CUFFT_ALLOC_FAILED"},
-        {CUFFT_INVALID_TYPE,"CUFFT_INVALID_TYPE"},
-        {CUFFT_INVALID_VALUE,"CUFFT_INVALID_VALUE"},
-        {CUFFT_INTERNAL_ERROR,"CUFFT_INTERNAL_ERROR"},
-        {CUFFT_EXEC_FAILED,"CUFFT_EXEC_FAILED"},
-        {CUFFT_SETUP_FAILED,"CUFFT_SETUP_FAILED"},
-        {CUFFT_INVALID_SIZE,"CUFFT_INVALID_SIZE"},
-        {CUFFT_UNALIGNED_DATA,"CUFFT_UNALIGNED_DATA"}};
-
-    const auto it = res2string.find(err);
-    if(it != res2string.end()){
-        return it->second;
-    }
-    else{
-        return std::to_string(err) +
-        " (unknown error code)";
-    }
-}
-#endif
 #endif // WARPX_USE_PSATD