Refactor:Replace the current fft with templates and polymorphism

source/Makefile.Objects

@@ -25,6 +25,7 @@ VPATH=./src_global:\
 ./module_base/module_container/ATen/ops:\
 ./module_base/module_device:\
 ./module_base/module_mixing:\
+./module_base/module_fft:\
 ./module_md:\
 ./module_basis/module_pw:\
 ./module_esolver:\
@@ -166,7 +167,9 @@ OBJS_BASE=abfs-vector3_order.o\
     broyden_mixing.o\
     memory_op.o\
     device.o\
-
+    fft_temp.o\
+    fft_base.o\
+    fft_cpu.o\
 
 OBJS_CELL=atom_pseudo.o\
     atom_spec.o\

source/module_base/CMakeLists.txt

@@ -6,7 +6,11 @@ list (APPEND LIBM_SRC
   libm/sincos.cpp
 )
 endif()
-
+if (ENABLE_FLOAT_FFTW)
+  list (APPEND FFT_SRC
+    module_fft/fft_cpu_float.cpp
+  )
+endif()
 add_library(
     base
     OBJECT
@@ -58,7 +62,11 @@ add_library(
     module_mixing/plain_mixing.cpp
     module_mixing/pulay_mixing.cpp
     module_mixing/broyden_mixing.cpp
+    module_fft/fft_base.cpp
+    module_fft/fft_temp.cpp
+    module_fft/fft_cpu.cpp
     ${LIBM_SRC}
+    ${FFT_SRC}
 )
 
 add_subdirectory(module_container)

source/module_base/module_fft/fft_base.cpp

@@ -0,0 +1,44 @@
+#include "fft_base.h"
+template <typename FPTYPE>
+FFT_BASE<FPTYPE>::FFT_BASE()
+{
+}
+template <typename FPTYPE>
+FFT_BASE<FPTYPE>::~FFT_BASE()
+{
+
+}
+template <typename FPTYPE>
+void FFT_BASE<FPTYPE>::initfft(int nx_in, int ny_in, int nz_in, int lixy_in, int rixy_in, int ns_in, int nplane_in, 
+				 int nproc_in, bool gamma_only_in, bool xprime_in, bool mpifft_in)
+{
+    this->gamma_only = gamma_only_in;
+    this->xprime = xprime_in;
+    this->fftnx = this->nx = nx_in;
+    this->fftny = this->ny = ny_in;
+    if (this->gamma_only)
+    {
+        if (xprime) {
+            this->fftnx = int(nx / 2) + 1;
+        } else {
+            this->fftny = int(ny / 2) + 1;
+}
+    }
+    this->nz = nz_in;
+    this->ns = ns_in;
+    this->lixy = lixy_in;
+    this->rixy = rixy_in;
+    this->nplane = nplane_in;
+    this->nproc = nproc_in;
+    this->mpifft = mpifft_in;
+    this->nxy = this->nx * this->ny;
+    this->fftnxy = this->fftnx * this->fftny;
+    const int nrxx = this->nxy * this->nplane;
+    const int nsz = this->nz * this->ns;
+    this->maxgrids = (nsz > nrxx) ? nsz : nrxx;
+}
+
+template FFT_BASE<float>::FFT_BASE();
+template FFT_BASE<double>::FFT_BASE();
+template FFT_BASE<float>::~FFT_BASE();
+template FFT_BASE<double>::~FFT_BASE();

source/module_base/module_fft/fft_base.h

@@ -0,0 +1,80 @@
+#include <complex>
+#include <string>
+#include "fftw3.h"
+#ifndef FFT_BASE_H
+#define FFT_BASE_H
+template <typename FPTYPE>
+class FFT_BASE
+{
+public:
+
+	FFT_BASE();
+	virtual  ~FFT_BASE(); 
+
+	// init parameters of fft
+	virtual void initfft(int nx_in, int ny_in, int nz_in, int lixy_in, int rixy_in, int ns_in, int nplane_in, 
+				 int nproc_in, bool gamma_only_in, bool xprime_in = true, bool mpifft_in = false);
+
+    virtual __attribute__((weak)) void initfftmode(int fft_mode_in);
+
+	//init fftw_plans
+	virtual void setupFFT()=0; 
+
+	//destroy fftw_plans
+	virtual void cleanFFT()=0;
+    //clear fftw_data
+    virtual void clear()=0;
+
+    // access the real space data
+    virtual __attribute__((weak)) FPTYPE* get_rspace_data() const;
+
+    virtual __attribute__((weak)) std::complex<FPTYPE>* get_auxr_data() const;
+
+    virtual __attribute__((weak)) std::complex<FPTYPE>* get_auxg_data() const;
+
+    virtual __attribute__((weak)) std::complex<FPTYPE>* get_auxr_3d_data() const;
+
+    //forward fft in x-y direction
+    virtual __attribute__((weak)) void fftxyfor(std::complex<FPTYPE>* in, std::complex<FPTYPE>* out) const;
+
+    virtual __attribute__((weak)) void fftxybac(std::complex<FPTYPE>* in, std::complex<FPTYPE>* out) const;
+
+    virtual __attribute__((weak)) void fftzfor(std::complex<FPTYPE>* in, std::complex<FPTYPE>* out) const;
+
+    virtual __attribute__((weak)) void fftzbac(std::complex<FPTYPE>* in, std::complex<FPTYPE>* out) const;
+
+    virtual __attribute__((weak)) void fftxyr2c(FPTYPE* in, std::complex<FPTYPE>* out) const;
+
+    virtual __attribute__((weak)) void fftxyc2r(std::complex<FPTYPE>* in, FPTYPE* out) const;
+
+    virtual __attribute__((weak)) void fft3D_forward(std::complex<FPTYPE>* in, std::complex<FPTYPE>* out) const;
+
+    virtual __attribute__((weak)) void fft3D_backward(std::complex<FPTYPE>* in, std::complex<FPTYPE>* out) const;
+
+protected:
+    int initflag = 0; // 0: not initialized; 1: initialized
+	int fftnx=0;
+    int fftny=0;
+	int fftnxy=0;
+	int ny=0;
+    int nx=0;
+    int nz=0;
+	int nxy=0;
+    int nplane=0; //number of x-y planes
+    bool gamma_only = false;
+    int lixy=0;
+    int rixy=0;// lixy: the left edge of the pw ball in the y direction; rixy: the right edge of the pw ball in the x or y direction
+    bool mpifft = false; // if use mpi fft, only used when define __FFTW3_MPI
+    int maxgrids = 0; // maxgrids = (nsz > nrxx) ? nsz : nrxx;
+    bool xprime = true; // true: when do recip2real, x-fft will be done last and when doing real2recip, x-fft will be done first; false: y-fft
+                         // For gamma_only, true: we use half x; false: we use half y
+	int ns=0; //number of sticks
+	int nproc=1; // number of proc.
+	int fft_mode = 0; ///< fftw mode 0: estimate, 1: measure, 2: patient, 3: exhaustive 
+
+public:
+    void set_device(std::string device_);
+    void set_precision(std::string precision_);
+
+};
+#endif // FFT_BASE_H