Refactor:Replace the current fft with templates and polymorphism

examples/scf/pw_Si2/INPUT

@@ -9,4 +9,4 @@ scf_thr         1e-7
 scf_nmax        100
 device          cpu
 ks_solver       dav_subspace
-precision       double
+precision       single

source/Makefile.Objects

@@ -27,6 +27,7 @@ VPATH=./src_global:\
 ./module_base/module_mixing:\
 ./module_md:\
 ./module_basis/module_pw:\
+./module_basis/module_pw/module_fft:\
 ./module_esolver:\
 ./module_hsolver:\
 ./module_hsolver/kernels:\
@@ -168,7 +169,6 @@ OBJS_BASE=abfs-vector3_order.o\
     memory_op.o\
     device.o\
 
-
 OBJS_CELL=atom_pseudo.o\
     atom_spec.o\
     pseudo.o\
@@ -414,6 +414,9 @@ OBJS_PSI_INITIALIZER=psi_initializer.o\
                      psi_initializer_nao_random.o\
 
 OBJS_PW=fft.o\
+    fft_bundle.o\
+    fft_base.o\
+    fft_cpu.o\
     pw_basis.o\
     pw_basis_k.o\
     pw_basis_sup.o\

source/module_base/CMakeLists.txt

@@ -6,7 +6,6 @@ list (APPEND LIBM_SRC
   libm/sincos.cpp
 )
 endif()
-
 add_library(
     base
     OBJECT

source/module_basis/module_pw/CMakeLists.txt

@@ -1,3 +1,8 @@
+if (ENABLE_FLOAT_FFTW)
+  list (APPEND FFT_SRC
+    module_fft/fft_cpu_float.cpp
+  )
+endif()
 list(APPEND objects
     fft.cpp
     pw_basis.cpp
@@ -10,6 +15,10 @@ list(APPEND objects
     pw_init.cpp
     pw_transform.cpp
     pw_transform_k.cpp
+    module_fft/fft_base.cpp
+    module_fft/fft_bundle.cpp
+    module_fft/fft_cpu.cpp
+    ${FFT_SRC}
 )
 
 add_library(

source/module_basis/module_pw/fft.cpp

@@ -72,10 +72,11 @@ void FFT::initfft(int nx_in, int ny_in, int nz_in, int lixy_in, int rixy_in, int
     this->fftny = this->ny = ny_in;
     if (this->gamma_only)
     {
-        if (xprime)
+        if (xprime) {
             this->fftnx = int(nx / 2) + 1;
-        else
+        } else {
             this->fftny = int(ny / 2) + 1;
+}
     }
     this->nz = nz_in;
     this->ns = ns_in;
@@ -92,10 +93,10 @@ void FFT::initfft(int nx_in, int ny_in, int nz_in, int lixy_in, int rixy_in, int
     int maxgrids = (nsz > nrxx) ? nsz : nrxx;
     if (!this->mpifft)
     {
-        z_auxg = (std::complex<double>*)fftw_malloc(sizeof(fftw_complex) * maxgrids);
-        z_auxr = (std::complex<double>*)fftw_malloc(sizeof(fftw_complex) * maxgrids);
-        ModuleBase::Memory::record("FFT::grid", 2 * sizeof(fftw_complex) * maxgrids);
-        d_rspace = (double*)z_auxg;
+        // z_auxg = (std::complex<double>*)fftw_malloc(sizeof(fftw_complex) * maxgrids);
+        // z_auxr = (std::complex<double>*)fftw_malloc(sizeof(fftw_complex) * maxgrids);
+        // ModuleBase::Memory::record("FFT::grid", 2 * sizeof(fftw_complex) * maxgrids);
+        // d_rspace = (double*)z_auxg;
         // auxr_3d = static_cast<std::complex<double> *>(
         //     fftw_malloc(sizeof(fftw_complex) * (this->nx * this->ny * this->nz)));
 #if defined(__CUDA) || defined(__ROCM)
@@ -105,15 +106,15 @@ void FFT::initfft(int nx_in, int ny_in, int nz_in, int lixy_in, int rixy_in, int
             resmem_zd_op()(gpu_ctx, this->z_auxr_3d, this->nx * this->ny * this->nz);
         }
 #endif // defined(__CUDA) || defined(__ROCM)
-#if defined(__ENABLE_FLOAT_FFTW)
-        if (this->precision == "single")
-        {
-            c_auxg = (std::complex<float>*)fftw_malloc(sizeof(fftwf_complex) * maxgrids);
-            c_auxr = (std::complex<float>*)fftw_malloc(sizeof(fftwf_complex) * maxgrids);
-            ModuleBase::Memory::record("FFT::grid_s", 2 * sizeof(fftwf_complex) * maxgrids);
-            s_rspace = (float*)c_auxg;
-        }
-#endif // defined(__ENABLE_FLOAT_FFTW)
+// #if defined(__ENABLE_FLOAT_FFTW)
+//         if (this->precision == "single")
+//         {
+//             c_auxg = (std::complex<float>*)fftw_malloc(sizeof(fftwf_complex) * maxgrids);
+//             c_auxr = (std::complex<float>*)fftw_malloc(sizeof(fftwf_complex) * maxgrids);
+//             ModuleBase::Memory::record("FFT::grid_s", 2 * sizeof(fftwf_complex) * maxgrids);
+//             s_rspace = (float*)c_auxg;
+//         }
+// #endif // defined(__ENABLE_FLOAT_FFTW)
     }
     else
     {
@@ -353,62 +354,62 @@ void FFT::cleanFFT()
     if (planzfor)
     {
         fftw_destroy_plan(planzfor);
-        planzfor = NULL;
+        planzfor = nullptr;
     }
     if (planzbac)
     {
         fftw_destroy_plan(planzbac);
-        planzbac = NULL;
+        planzbac = nullptr;
     }
     if (planxfor1)
     {
         fftw_destroy_plan(planxfor1);
-        planxfor1 = NULL;
+        planxfor1 = nullptr;
     }
     if (planxbac1)
     {
         fftw_destroy_plan(planxbac1);
-        planxbac1 = NULL;
+        planxbac1 = nullptr;
     }
     if (planxfor2)
     {
         fftw_destroy_plan(planxfor2);
-        planxfor2 = NULL;
+        planxfor2 = nullptr;
     }
     if (planxbac2)
     {
         fftw_destroy_plan(planxbac2);
-        planxbac2 = NULL;
+        planxbac2 = nullptr;
     }
     if (planyfor)
     {
         fftw_destroy_plan(planyfor);
-        planyfor = NULL;
+        planyfor = nullptr;
     }
     if (planybac)
     {
         fftw_destroy_plan(planybac);
-        planybac = NULL;
+        planybac = nullptr;
     }
     if (planxr2c)
     {
         fftw_destroy_plan(planxr2c);
-        planxr2c = NULL;
+        planxr2c = nullptr;
     }
     if (planxc2r)
     {
         fftw_destroy_plan(planxc2r);
-        planxc2r = NULL;
+        planxc2r = nullptr;
     }
     if (planyr2c)
     {
         fftw_destroy_plan(planyr2c);
-        planyr2c = NULL;
+        planyr2c = nullptr;
     }
     if (planyc2r)
     {
         fftw_destroy_plan(planyc2r);
-        planyc2r = NULL;
+        planyc2r = nullptr;
     }
 
     // fftw_destroy_plan(this->plan3dforward);

source/module_basis/module_pw/module_fft/fft_base.cpp

@@ -0,0 +1,17 @@
+#include "fft_base.h"
+namespace ModulePW
+{
+template <typename FPTYPE>
+FFT_BASE<FPTYPE>::FFT_BASE()
+{
+}
+template <typename FPTYPE>
+FFT_BASE<FPTYPE>::~FFT_BASE()
+{
+}
+
+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_basis/module_pw/module_fft/fft_base.h

@@ -0,0 +1,62 @@
+#include <complex>
+#include <string>
+#include "fftw3.h"
+#ifndef FFT_BASE_H
+#define FFT_BASE_H
+namespace ModulePW
+{
+template <typename FPTYPE>
+class FFT_BASE
+{
+public:
+
+	FFT_BASE();
+	virtual  ~FFT_BASE(); 
+
+	// init parameters of fft
+	virtual __attribute__((weak))
+    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);
+
+	//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 ny=0;
+    int nx=0;
+    int nz=0;
+
+};
+}
+#endif // FFT_BASE_H