add change on pw_basis_k

Author: A-006

source/source_basis/module_pw/pw_transform_k.cpp

@@ -366,7 +366,7 @@ void PW_Basis_K::convolution(const base_device::DEVICE_CPU* ctx,
     auto* auxg = this->fft_bundle.get_auxg_data<double>();
     auto* auxr=this->fft_bundle.get_auxr_data<double>();
 
-     memset(auxg, 0, this->nst * this->nz * 2 * 8);
+    memset(auxg, 0, this->nst * this->nz * 2 * 8);
     const int startig = ik * this->npwk_max;
     const int npwk = this->npwk[ik];
 
@@ -385,28 +385,28 @@ void PW_Basis_K::convolution(const base_device::DEVICE_CPU* ctx,
     this->gathers_scatterp(auxg, auxr);
 
     this->fft_bundle.fftxybac(auxr, auxr);
+
+    #ifdef _OPENMP
+    #pragma omp parallel for simd schedule(static) aligned(auxr, input1: 64)
+    #endif
     for (int ir = 0; ir < size; ir++)
     {
         auxr[ir] *= input1[ir];
     }
-
     // 3d fft
     this->fft_bundle.fftxyfor(auxr, auxr);
 
     this->gatherp_scatters(auxr, auxg);
 
     this->fft_bundle.fftzfor(auxg, auxg);
     // copy the result from the auxr to the out ,while consider the add
-    if (add)
-    {
-        double tmpfac = factor / double(this->nxyz);
+    double tmpfac = factor / double(this->nxyz);
 #ifdef _OPENMP
 #pragma omp parallel for schedule(static, 4096 / sizeof(double))
 #endif
-        for (int igl = 0; igl < npwk; ++igl)
-        {
-            output[igl] += tmpfac * auxg[this->igl2isz_k[igl + startig]];
-        }
+    for (int igl = 0; igl < npwk; ++igl)
+    {
+        output[igl] += tmpfac * auxg[this->igl2isz_k[igl + startig]];
     }
     ModuleBase::timer::tick(this->classname, "convolution");
 }

source/source_io/cal_ldos.cpp

@@ -140,7 +140,7 @@ void stm_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
 
             for (int ib = 0; ib < nbands; ib++)
             {
-                pelec->basis->recip2real(&psi(ib, 0), wfcr.data(), ik);
+                pelec->basis->recip_to_real<std::complex<double>,base_device::DEVICE_CPU>(&psi(ib, 0), wfcr.data(), ik);
 
                 const double eigenval = (pelec->ekb(ik, ib) - efermi) * ModuleBase::Ry_to_eV;
                 double weight = en > 0 ? pelec->klist->wk[ik] - pelec->wg(ik, ib) : pelec->wg(ik, ib);
@@ -210,7 +210,7 @@ void ldos_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
 
         for (int ib = 0; ib < nbands; ib++)
         {
-            pelec->basis->recip2real(&psi(ib, 0), wfcr.data(), ik);
+            pelec->basis->recip_to_real<std::complex<double>,base_device::DEVICE_CPU>(&psi(ib, 0), wfcr.data(), ik);
             const double weight = pelec->klist->wk[ik] / ucell.omega;
 
             for (int ir = 0; ir < pelec->basis->nrxx; ir++)

source/source_io/cal_mlkedf_descriptors.cpp

@@ -472,7 +472,7 @@ void Cal_MLKEDF_Descriptors::getF_KS(
                     wfcr[ig] = psi->operator()(ibnd, ig) * std::complex<double>(0.0, fact);
                 }
 
-                pw_psi->recip2real(wfcr, wfcr, ik);
+                pw_psi->recip_to_real<std::complex<double>,base_device::DEVICE_CPU>(wfcr, wfcr, ik);
 
                 for (int ir = 0; ir < this->nx; ++ir)
                 {

source/source_io/get_wf_lcao.cpp

@@ -179,7 +179,7 @@ void Get_wf_lcao::begin(const UnitCell& ucell,
                 // Calculate real-space wave functions
                 psi_g.fix_k(is);
                 std::vector<std::complex<double>> wfc_r(pw_wfc->nrxx);
-                pw_wfc->recip2real(&psi_g(ib, 0), wfc_r.data(), is);
+                pw_wfc->recip_to_real<std::complex<double>,base_device::DEVICE_CPU>(&psi_g(ib, 0), wfc_r.data(), is);
 
                 // Extract real and imaginary parts
                 std::vector<double> wfc_real(pw_wfc->nrxx);
@@ -399,7 +399,7 @@ void Get_wf_lcao::begin(const UnitCell& ucell,
 
                 // Calculate real-space wave functions
                 std::vector<std::complex<double>> wfc_r(pw_wfc->nrxx);
-                pw_wfc->recip2real(&psi_g(ib, 0), wfc_r.data(), ik);
+                pw_wfc->recip_to_real<std::complex<double>,base_device::DEVICE_CPU>(&psi_g(ib, 0), wfc_r.data(), ik);
 
                 // Extract real and imaginary parts
                 std::vector<double> wfc_real(pw_wfc->nrxx);
@@ -551,7 +551,7 @@ void Get_wf_lcao::set_pw_wfc(const ModulePW::PW_Basis_K* pw_wfc,
     }
 
     // call FFT
-    pw_wfc->real2recip(Porter.data(), &wfc_g(ib, 0), ik);
+    pw_wfc->real_to_recip<std::complex<double>,base_device::DEVICE_CPU>(Porter.data(), &wfc_g(ib, 0), ik);
 }
 
 #ifdef __MPI

source/source_io/unk_overlap_pw.cpp

@@ -93,7 +93,7 @@ std::complex<double> unkOverlap_pw::unkdotp_G0(const ModulePW::PW_Basis* rhopw,
     }
 
     // (3) calculate the overlap in ik_L and ik_R
-    wfcpw->real2recip(psi_r, psi_r, ik_R);
+    wfcpw->real_to_recip<std::complex<double>,base_device::DEVICE_CPU>(psi_r, psi_r, ik_R);
 
     for (int ig = 0; ig < evc->get_ngk(ik_R); ig++)
     {
@@ -197,8 +197,8 @@ std::complex<double> unkOverlap_pw::unkdotp_soc_G0(const ModulePW::PW_Basis* rho
 
     // (2) fft and get value
     rhopw->recip2real(phase, phase);
-    wfcpw->recip2real(&evc[0](ik_L, iband_L, 0), psi_up, ik_L);
-    wfcpw->recip2real(&evc[0](ik_L, iband_L, npwx), psi_down, ik_L);
+    wfcpw->recip_to_real<std::complex<double>,base_device::DEVICE_CPU>(&evc[0](ik_L, iband_L, 0), psi_up, ik_L);
+    wfcpw->recip_to_real<std::complex<double>,base_device::DEVICE_CPU>(&evc[0](ik_L, iband_L, npwx), psi_down, ik_L);
 
     for (int ir = 0; ir < wfcpw->nrxx; ir++)
     {
@@ -207,8 +207,8 @@ std::complex<double> unkOverlap_pw::unkdotp_soc_G0(const ModulePW::PW_Basis* rho
     }
 
     // (3) calculate the overlap in ik_L and ik_R
-    wfcpw->real2recip(psi_up, psi_up, ik_L);
-    wfcpw->real2recip(psi_down, psi_down, ik_L);
+    wfcpw->real_to_recip<std::complex<double>,base_device::DEVICE_CPU>(psi_up, psi_up, ik_L);
+    wfcpw->real_to_recip<std::complex<double>,base_device::DEVICE_CPU>(psi_down, psi_down, ik_L);
 
     for (int i = 0; i < PARAM.globalv.npol; i++)
     {