From bdbb70b0b9ef15f0a4f38a074c078d0cd2e9c102 Mon Sep 17 00:00:00 2001 From: ubuntu <3158793232@qq.com> Date: Wed, 19 Mar 2025 09:40:14 +0800 Subject: [PATCH 1/8] update recip_to_real in the rhpw --- source/module_basis/module_pw/pw_basis.h | 36 +- .../module_basis/module_pw/pw_transform.cpp | 361 ++++++++++++++---- source/module_elecstate/elecstate_pw.cpp | 3 +- 3 files changed, 329 insertions(+), 71 deletions(-) diff --git a/source/module_basis/module_pw/pw_basis.h b/source/module_basis/module_pw/pw_basis.h index 47827000eb7..4f1ad9a1ee8 100644 --- a/source/module_basis/module_pw/pw_basis.h +++ b/source/module_basis/module_pw/pw_basis.h @@ -267,6 +267,35 @@ class PW_Basis const bool add = false, const FPTYPE factor = 1.0) const; // in:(nz, ns) ; out(nplane,nx*ny) + template + void real_to_recip(const std::complex* in, + std::complex* out, + const Device* ctx=get_default_device_ctx(), + const int ik=0, + const bool add = false, + const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny) ; out(nz, ns) + template + void recip_to_real(const std::complex* in, + std::complex* out, + const Device* ctx=get_default_device_ctx(), + const int ik=0, + const bool add = false, + const FPTYPE factor = 1.0) const; // in:(nz, ns) ; out(nplane,nx*ny) + + template + void real_to_recip(FPTYPE* in, + std::complex* out, + const Device* ctx=get_default_device_ctx(), + const int ik=0, + const bool add = false, + const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny) ; out(nz, ns) + template + void recip_to_real(const std::complex* in, + FPTYPE* out, + const Device* ctx=get_default_device_ctx() , + const int ik=0, + const bool add = false, + const FPTYPE factor = 1.0) const; // in:(nz, ns) ; out(nplane,nx*ny) protected: //gather planes and scatter sticks of all processors template @@ -282,15 +311,16 @@ class PW_Basis using resmem_int_op = base_device::memory::resize_memory_op; using delmem_int_op = base_device::memory::delete_memory_op; - using syncmem_int_h2d_op - = base_device::memory::synchronize_memory_op; - + using syncmem_int_h2d_op = base_device::memory::synchronize_memory_op; + // using default_device_cpu = base_device::DEVICE_CPU; + void set_device(std::string device_); void set_precision(std::string precision_); protected: std::string device = "cpu"; std::string precision = "double"; + static const base_device::DEVICE_CPU* get_default_device_ctx(); }; } diff --git a/source/module_basis/module_pw/pw_transform.cpp b/source/module_basis/module_pw/pw_transform.cpp index 8e458b25616..38d4bfc49df 100644 --- a/source/module_basis/module_pw/pw_transform.cpp +++ b/source/module_basis/module_pw/pw_transform.cpp @@ -1,12 +1,19 @@ -#include "module_fft/fft_bundle.h" -#include -#include "pw_basis.h" -#include #include "module_base/global_function.h" #include "module_base/timer.h" +#include "module_basis/module_pw/kernels/pw_op.h" +#include "module_fft/fft_bundle.h" +#include "pw_basis.h" #include "pw_gatherscatter.h" -namespace ModulePW { +#include +#include + +namespace ModulePW +{ + const base_device::DEVICE_CPU* PW_Basis::get_default_device_ctx() { + static const base_device::DEVICE_CPU* default_device_cpu; + return default_device_cpu; +} /** * @brief transform real space to reciprocal space * @details c(g)=\int dr*f(r)*exp(-ig*r) @@ -24,25 +31,25 @@ void PW_Basis::real2recip(const std::complex* in, assert(this->gamma_only == false); #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ir = 0 ; ir < this->nrxx ; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { this->fft_bundle.get_auxr_data()[ir] = in[ir]; } - this->fft_bundle.fftxyfor(fft_bundle.get_auxr_data(),fft_bundle.get_auxr_data()); + this->fft_bundle.fftxyfor(fft_bundle.get_auxr_data(), fft_bundle.get_auxr_data()); this->gatherp_scatters(this->fft_bundle.get_auxr_data(), this->fft_bundle.get_auxg_data()); - - this->fft_bundle.fftzfor(fft_bundle.get_auxg_data(),fft_bundle.get_auxg_data()); - if(add) + this->fft_bundle.fftzfor(fft_bundle.get_auxg_data(), fft_bundle.get_auxg_data()); + + if (add) { FPTYPE tmpfac = factor / FPTYPE(this->nxyz); #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ig = 0 ; ig < this->npw ; ++ig) + for (int ig = 0; ig < this->npw; ++ig) { out[ig] += tmpfac * this->fft_bundle.get_auxg_data()[this->ig2isz[ig]]; } @@ -51,9 +58,9 @@ void PW_Basis::real2recip(const std::complex* in, { FPTYPE tmpfac = 1.0 / FPTYPE(this->nxyz); #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ig = 0 ; ig < this->npw ; ++ig) + for (int ig = 0; ig < this->npw; ++ig) { out[ig] = tmpfac * this->fft_bundle.get_auxg_data()[this->ig2isz[ig]]; } @@ -72,44 +79,44 @@ template void PW_Basis::real2recip(const FPTYPE* in, std::complex* out, const bool add, const FPTYPE factor) const { ModuleBase::timer::tick(this->classname, "real2recip"); - if(this->gamma_only) + if (this->gamma_only) { const int npy = this->ny * this->nplane; #ifdef _OPENMP -#pragma omp parallel for collapse(2) schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for collapse(2) schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ix = 0 ; ix < this->nx ; ++ix) + for (int ix = 0; ix < this->nx; ++ix) { - for(int ipy = 0 ; ipy < npy ; ++ipy) + for (int ipy = 0; ipy < npy; ++ipy) { - this->fft_bundle.get_rspace_data()[ix*npy + ipy] = in[ix*npy + ipy]; + this->fft_bundle.get_rspace_data()[ix * npy + ipy] = in[ix * npy + ipy]; } } - this->fft_bundle.fftxyr2c(fft_bundle.get_rspace_data(),fft_bundle.get_auxr_data()); + this->fft_bundle.fftxyr2c(fft_bundle.get_rspace_data(), fft_bundle.get_auxr_data()); } else { #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ir = 0 ; ir < this->nrxx ; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { - this->fft_bundle.get_auxr_data()[ir] = std::complex(in[ir],0); + this->fft_bundle.get_auxr_data()[ir] = std::complex(in[ir], 0); } - this->fft_bundle.fftxyfor(fft_bundle.get_auxr_data(),fft_bundle.get_auxr_data()); + this->fft_bundle.fftxyfor(fft_bundle.get_auxr_data(), fft_bundle.get_auxr_data()); } this->gatherp_scatters(this->fft_bundle.get_auxr_data(), this->fft_bundle.get_auxg_data()); - - this->fft_bundle.fftzfor(fft_bundle.get_auxg_data(),fft_bundle.get_auxg_data()); - if(add) + this->fft_bundle.fftzfor(fft_bundle.get_auxg_data(), fft_bundle.get_auxg_data()); + + if (add) { FPTYPE tmpfac = factor / FPTYPE(this->nxyz); #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ig = 0 ; ig < this->npw ; ++ig) + for (int ig = 0; ig < this->npw; ++ig) { out[ig] += tmpfac * this->fft_bundle.get_auxg_data()[this->ig2isz[ig]]; } @@ -118,9 +125,9 @@ void PW_Basis::real2recip(const FPTYPE* in, std::complex* out, const boo { FPTYPE tmpfac = 1.0 / FPTYPE(this->nxyz); #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ig = 0 ; ig < this->npw ; ++ig) + for (int ig = 0; ig < this->npw; ++ig) { out[ig] = tmpfac * this->fft_bundle.get_auxg_data()[this->ig2isz[ig]]; } @@ -144,32 +151,32 @@ void PW_Basis::recip2real(const std::complex* in, ModuleBase::timer::tick(this->classname, "recip2real"); assert(this->gamma_only == false); #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int i = 0 ; i < this->nst * this->nz ; ++i) + for (int i = 0; i < this->nst * this->nz; ++i) { fft_bundle.get_auxg_data()[i] = std::complex(0, 0); } #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ig = 0 ; ig < this->npw ; ++ig) + for (int ig = 0; ig < this->npw; ++ig) { this->fft_bundle.get_auxg_data()[this->ig2isz[ig]] = in[ig]; } this->fft_bundle.fftzbac(fft_bundle.get_auxg_data(), fft_bundle.get_auxg_data()); - this->gathers_scatterp(this->fft_bundle.get_auxg_data(),this->fft_bundle.get_auxr_data()); + this->gathers_scatterp(this->fft_bundle.get_auxg_data(), this->fft_bundle.get_auxr_data()); + + this->fft_bundle.fftxybac(fft_bundle.get_auxr_data(), fft_bundle.get_auxr_data()); - this->fft_bundle.fftxybac(fft_bundle.get_auxr_data(),fft_bundle.get_auxr_data()); - - if(add) + if (add) { #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ir = 0 ; ir < this->nrxx ; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { out[ir] += factor * this->fft_bundle.get_auxr_data()[ir]; } @@ -177,9 +184,9 @@ void PW_Basis::recip2real(const std::complex* in, else { #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ir = 0 ; ir < this->nrxx ; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { out[ir] = this->fft_bundle.get_auxr_data()[ir]; } @@ -199,17 +206,17 @@ void PW_Basis::recip2real(const std::complex* in, FPTYPE* out, const boo { ModuleBase::timer::tick(this->classname, "recip2real"); #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int i = 0 ; i < this->nst * this->nz ; ++i) + for (int i = 0; i < this->nst * this->nz; ++i) { fft_bundle.get_auxg_data()[i] = std::complex(0, 0); } #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ig = 0 ; ig < this->npw ; ++ig) + for (int ig = 0; ig < this->npw; ++ig) { this->fft_bundle.get_auxg_data()[this->ig2isz[ig]] = in[ig]; } @@ -217,49 +224,49 @@ void PW_Basis::recip2real(const std::complex* in, FPTYPE* out, const boo this->gathers_scatterp(this->fft_bundle.get_auxg_data(), this->fft_bundle.get_auxr_data()); - if(this->gamma_only) + if (this->gamma_only) { - this->fft_bundle.fftxyc2r(fft_bundle.get_auxr_data(),fft_bundle.get_rspace_data()); + this->fft_bundle.fftxyc2r(fft_bundle.get_auxr_data(), fft_bundle.get_rspace_data()); // r2c in place const int npy = this->ny * this->nplane; - if(add) + if (add) { #ifdef _OPENMP -#pragma omp parallel for collapse(2) schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for collapse(2) schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ix = 0 ; ix < this->nx ; ++ix) + for (int ix = 0; ix < this->nx; ++ix) { - for(int ipy = 0 ; ipy < npy ; ++ipy) + for (int ipy = 0; ipy < npy; ++ipy) { - out[ix*npy + ipy] += factor * this->fft_bundle.get_rspace_data()[ix*npy + ipy]; + out[ix * npy + ipy] += factor * this->fft_bundle.get_rspace_data()[ix * npy + ipy]; } } } else { #ifdef _OPENMP -#pragma omp parallel for collapse(2) schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for collapse(2) schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ix = 0 ; ix < this->nx ; ++ix) + for (int ix = 0; ix < this->nx; ++ix) { - for(int ipy = 0 ; ipy < npy ; ++ipy) + for (int ipy = 0; ipy < npy; ++ipy) { - out[ix*npy + ipy] = this->fft_bundle.get_rspace_data()[ix*npy + ipy]; + out[ix * npy + ipy] = this->fft_bundle.get_rspace_data()[ix * npy + ipy]; } } } } else { - this->fft_bundle.fftxybac(fft_bundle.get_auxr_data(),fft_bundle.get_auxr_data()); - if(add) + this->fft_bundle.fftxybac(fft_bundle.get_auxr_data(), fft_bundle.get_auxr_data()); + if (add) { #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ir = 0 ; ir < this->nrxx ; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { out[ir] += factor * this->fft_bundle.get_auxr_data()[ir].real(); } @@ -267,9 +274,9 @@ void PW_Basis::recip2real(const std::complex* in, FPTYPE* out, const boo else { #ifdef _OPENMP -#pragma omp parallel for schedule(static, 4096/sizeof(FPTYPE)) +#pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE)) #endif - for(int ir = 0 ; ir < this->nrxx ; ++ir) + for (int ir = 0; ir < this->nrxx; ++ir) { out[ir] = this->fft_bundle.get_auxr_data()[ir].real(); } @@ -278,6 +285,228 @@ void PW_Basis::recip2real(const std::complex* in, FPTYPE* out, const boo ModuleBase::timer::tick(this->classname, "recip2real"); } +template <> +void PW_Basis::real_to_recip(const std::complex* in, + std::complex* out, + const base_device::DEVICE_CPU* ctx, + const int ik, + const bool add, + const float factor) const +{ + this->real2recip(in, out, add, factor); +} +template <> +void PW_Basis::real_to_recip(const std::complex* in, + std::complex* out, + const base_device::DEVICE_CPU* ctx, + const int ik, + const bool add, + const double factor) const +{ + this->real2recip(in, out, add, factor); +} + +template <> +void PW_Basis::recip_to_real(const std::complex* in, + std::complex* out, + const base_device::DEVICE_CPU* ctx, + const int ik, + const bool add, + const float factor) const +{ + this->recip2real(in, out, add, factor); +} +template <> +void PW_Basis::recip_to_real(const std::complex* in, + std::complex* out, + const base_device::DEVICE_CPU* ctx, + const int ik, + const bool add, + const double factor) const +{ + this->recip2real(in, out, add, factor); +} + +template <> +void PW_Basis::real_to_recip(float* in, + std::complex* out, + const base_device::DEVICE_CPU* ctx, + const int ik, + const bool add, + const float factor) const +{ + this->real2recip(in, out, add, factor); +} +template <> +void PW_Basis::real_to_recip(double* in, + std::complex* out, + const base_device::DEVICE_CPU* ctx, + const int ik, + const bool add, + const double factor) const +{ + this->real2recip(in, out, add, factor); +} + +template <> +void PW_Basis::recip_to_real(const std::complex* in, + float* out, + const base_device::DEVICE_CPU* ctx, + const int ik, + const bool add, + const float factor) const +{ + this->recip2real(in, out, add, factor); +} + +template <> +void PW_Basis::recip_to_real(const std::complex* in, + double* out, + const base_device::DEVICE_CPU* ctx, + const int ik, + const bool add, + const double factor) const +{ + this->recip2real(in, out, add, factor); +} + +#if (defined(__CUDA) || defined(__ROCM)) +template <> +void PW_Basis::real_to_recip(const std::complex* in, + std::complex* out, + const base_device::DEVICE_GPU* ctx, + const int ik, + const bool add, + const float factor) const +{ + ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); + assert(this->gamma_only == false); + assert(this->poolnproc == 1); + + base_device::memory::synchronize_memory_op, base_device::DEVICE_GPU, base_device::DEVICE_GPU>()( + this->fft_bundle.get_auxr_3d_data(), + in, + this->nrxx); + + this->fft_bundle.fft3D_forward(ctx, + this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.get_auxr_3d_data()); + + set_real_to_recip_output_op()(ctx, + npw, + this->nxyz, + add, + factor, + this->ig2isz, + this->fft_bundle.get_auxr_3d_data(), + out); + ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); +} +template <> +void PW_Basis::real_to_recip(const std::complex* in, + std::complex* out, + const base_device::DEVICE_GPU* ctx, + const int ik, + const bool add, + const double factor) const +{ + ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); + assert(this->gamma_only == false); + assert(this->poolnproc == 1); + + base_device::memory::synchronize_memory_op, + base_device::DEVICE_GPU, + base_device::DEVICE_GPU>()(this->fft_bundle.get_auxr_3d_data(), + in, + this->nrxx); + + this->fft_bundle.fft3D_forward(ctx, + this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.get_auxr_3d_data()); + + set_real_to_recip_output_op()(ctx, + npw, + this->nxyz, + add, + factor, + this->ig2isz, + this->fft_bundle.get_auxr_3d_data(), + out); + ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); +} + +template <> +void PW_Basis::recip_to_real(const std::complex* in, + std::complex* out, + const base_device::DEVICE_GPU* ctx, + const int ik, + const bool add, + const float factor) const +{ + ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); + assert(this->gamma_only == false); + assert(this->poolnproc == 1); + // ModuleBase::GlobalFunc::ZEROS(fft_bundle.get_auxr_3d_data(), this->nxyz); + base_device::memory::set_memory_op, base_device::DEVICE_GPU>()( + this->fft_bundle.get_auxr_3d_data(), + 0, + this->nxyz); + + set_3d_fft_box_op()(ctx, + npw, + this->ig2isz, + in, + this->fft_bundle.get_auxr_3d_data()); + this->fft_bundle.fft3D_backward(ctx, + this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.get_auxr_3d_data()); + + set_recip_to_real_output_op()(ctx, + this->nrxx, + add, + factor, + this->fft_bundle.get_auxr_3d_data(), + out); + + ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); +} +template <> +void PW_Basis::recip_to_real(const std::complex* in, + std::complex* out, + const base_device::DEVICE_GPU* ctx, + const int ik, + const bool add, + const double factor) const +{ + ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); + assert(this->gamma_only == false); + assert(this->poolnproc == 1); + // ModuleBase::GlobalFunc::ZEROS(fft_bundle.get_auxr_3d_data(), this->nxyz); + base_device::memory::set_memory_op, base_device::DEVICE_GPU>()( + this->fft_bundle.get_auxr_3d_data(), + 0, + this->nxyz); + + set_3d_fft_box_op()(ctx, + npw, + this->ig2isz, + in, + this->fft_bundle.get_auxr_3d_data()); + this->fft_bundle.fft3D_backward(ctx, + this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.get_auxr_3d_data()); + + set_recip_to_real_output_op()(ctx, + this->nrxx, + add, + factor, + this->fft_bundle.get_auxr_3d_data(), + out); + + ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); +} +#endif + template void PW_Basis::real2recip(const float* in, std::complex* out, const bool add, @@ -311,4 +540,4 @@ template void PW_Basis::recip2real(const std::complex* in, std::complex* out, const bool add, const double factor) const; -} \ No newline at end of file +} // namespace ModulePW \ No newline at end of file diff --git a/source/module_elecstate/elecstate_pw.cpp b/source/module_elecstate/elecstate_pw.cpp index 739e497515c..e09a1a3dfa7 100644 --- a/source/module_elecstate/elecstate_pw.cpp +++ b/source/module_elecstate/elecstate_pw.cpp @@ -443,13 +443,12 @@ void ElecStatePW::add_usrho(const psi::Psi& psi) { this->addusdens_g(becsum, rhog); } - // transform back to real space using dense grids if (PARAM.globalv.double_grid || PARAM.globalv.use_uspp) { for (int is = 0; is < PARAM.inp.nspin; is++) { - this->charge->rhopw->recip2real(this->rhog[is], this->rho[is]); + this->charge->rhopw->recip_to_real(this->rhog[is], this->rho[is]); } } } From 926f608371e5bb6b9ae93a577e97300440261e15 Mon Sep 17 00:00:00 2001 From: ubuntu <3158793232@qq.com> Date: Fri, 21 Mar 2025 16:07:36 +0800 Subject: [PATCH 2/8] add the operator in pw_basis --- source/module_basis/module_pw/CMakeLists.txt | 1 + .../module_pw/kernels/cuda/pw_op.cu | 80 ++++++ .../module_basis/module_pw/kernels/pw_op.cpp | 23 ++ source/module_basis/module_pw/kernels/pw_op.h | 33 +++ .../module_pw/kernels/rocm/pw_op.hip.cu | 80 ++++++ source/module_basis/module_pw/pw_basis.h | 168 ++++++++++--- .../module_basis/module_pw/pw_transform.cpp | 231 +----------------- .../module_pw/pw_transform_gpu.cpp | 173 +++++++++++++ source/module_elecstate/elecstate_pw.cpp | 2 +- 9 files changed, 535 insertions(+), 256 deletions(-) create mode 100644 source/module_basis/module_pw/pw_transform_gpu.cpp diff --git a/source/module_basis/module_pw/CMakeLists.txt b/source/module_basis/module_pw/CMakeLists.txt index e365e12b5e5..50b7ec34306 100644 --- a/source/module_basis/module_pw/CMakeLists.txt +++ b/source/module_basis/module_pw/CMakeLists.txt @@ -30,6 +30,7 @@ list(APPEND objects pw_distributer.cpp pw_init.cpp pw_transform.cpp + pw_transform_gpu.cpp pw_transform_k.cpp module_fft/fft_bundle.cpp module_fft/fft_cpu.cpp diff --git a/source/module_basis/module_pw/kernels/cuda/pw_op.cu b/source/module_basis/module_pw/kernels/cuda/pw_op.cu index a9128db3187..d9418ca486d 100644 --- a/source/module_basis/module_pw/kernels/cuda/pw_op.cu +++ b/source/module_basis/module_pw/kernels/cuda/pw_op.cu @@ -41,6 +41,24 @@ __global__ void set_recip_to_real_output( } } +template +__global__ void set_recip_to_real_output( + const int nrxx, + const bool add, + const FPTYPE factor, + const thrust::complex* in, + FPTYPE* out) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if(idx >= nrxx) {return;} + if(add) { + out[idx] += factor * in[idx].real(); + } + else { + out[idx] = in[idx].real(); + } +} + template __global__ void set_real_to_recip_output( const int npwk, @@ -61,6 +79,26 @@ __global__ void set_real_to_recip_output( } } +template +__global__ void set_real_to_recip_output( + const int npwk, + const int nxyz, + const bool add, + const FPTYPE factor, + const int* box_index, + const thrust::complex* in, + FPTYPE* out) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if(idx >= npwk) {return;} + if(add) { + out[idx] += factor / nxyz * in[box_index[idx]].real(); + } + else { + out[idx] = in[box_index[idx]].real() / nxyz; + } +} + template void set_3d_fft_box_op::operator()(const base_device::DEVICE_GPU* /*dev*/, const int npwk, @@ -97,6 +135,25 @@ void set_recip_to_real_output_op::operator()(co cudaCheckOnDebug(); } +template +void set_recip_to_real_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, + const int nrxx, + const bool add, + const FPTYPE factor, + const std::complex* in, + FPTYPE* out) +{ + const int block = (nrxx + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + set_recip_to_real_output<<>>( + nrxx, + add, + factor, + reinterpret_cast*>(in), + reinterpret_cast(out)); + + cudaCheckOnDebug(); +} + template void set_real_to_recip_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, const int npwk, @@ -120,6 +177,29 @@ void set_real_to_recip_output_op::operator()(co cudaCheckOnDebug(); } +template +void set_real_to_recip_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, + const int npwk, + const int nxyz, + const bool add, + const FPTYPE factor, + const int* box_index, + const std::complex* in, + FPTYPE* out) +{ + const int block = (npwk + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + set_real_to_recip_output<<>>( + npwk, + nxyz, + add, + factor, + box_index, + reinterpret_cast*>(in), + reinterpret_cast(out)); + + cudaCheckOnDebug(); +} + template struct set_3d_fft_box_op; template struct set_recip_to_real_output_op; template struct set_real_to_recip_output_op; diff --git a/source/module_basis/module_pw/kernels/pw_op.cpp b/source/module_basis/module_pw/kernels/pw_op.cpp index b5fb4533540..e39e8a6a7c7 100644 --- a/source/module_basis/module_pw/kernels/pw_op.cpp +++ b/source/module_basis/module_pw/kernels/pw_op.cpp @@ -39,6 +39,29 @@ struct set_recip_to_real_output_op } } } + + void operator()(const base_device::DEVICE_CPU* /*dev*/, + const int nrxx, + const bool add, + const FPTYPE factor, + const std::complex* in, + FPTYPE* out) + { + if (add) + { + for (int ir = 0; ir < nrxx; ++ir) + { + out[ir] += factor * in[ir].real(); + } + } + else + { + for (int ir = 0; ir < nrxx; ++ir) + { + out[ir] = in[ir].real(); + } + } + } }; template diff --git a/source/module_basis/module_pw/kernels/pw_op.h b/source/module_basis/module_pw/kernels/pw_op.h index 8415ad96778..b5471906061 100644 --- a/source/module_basis/module_pw/kernels/pw_op.h +++ b/source/module_basis/module_pw/kernels/pw_op.h @@ -45,6 +45,14 @@ struct set_recip_to_real_output_op { const FPTYPE factor, const std::complex* in, std::complex* out); + + void operator() ( + const Device* dev, + const int nrxx, + const bool add, + const FPTYPE factor, + const std::complex* in, + FPTYPE* out); }; template @@ -70,6 +78,16 @@ struct set_real_to_recip_output_op { const int* box_index, const std::complex* in, std::complex* out); + + void operator() ( + const Device* dev, + const int npw_k, + const int nxyz, + const bool add, + const FPTYPE factor, + const int* box_index, + const std::complex* in, + FPTYPE* out); }; #if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM @@ -93,6 +111,13 @@ struct set_recip_to_real_output_op const FPTYPE factor, const std::complex* in, std::complex* out); + + void operator()(const base_device::DEVICE_GPU* dev, + const int nrxx, + const bool add, + const FPTYPE factor, + const std::complex* in, + FPTYPE* out); }; template @@ -106,6 +131,14 @@ struct set_real_to_recip_output_op const int* box_index, const std::complex* in, std::complex* out); + void operator()(const base_device::DEVICE_GPU* dev, + const int npw_k, + const int nxyz, + const bool add, + const FPTYPE factor, + const int* box_index, + const std::complex* in, + FPTYPE* out); }; #endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM diff --git a/source/module_basis/module_pw/kernels/rocm/pw_op.hip.cu b/source/module_basis/module_pw/kernels/rocm/pw_op.hip.cu index a3f5fe2c2bd..44be7d0671f 100644 --- a/source/module_basis/module_pw/kernels/rocm/pw_op.hip.cu +++ b/source/module_basis/module_pw/kernels/rocm/pw_op.hip.cu @@ -42,6 +42,24 @@ __global__ void set_recip_to_real_output( } } +template +__global__ void set_recip_to_real_output( + const int nrxx, + const bool add, + const FPTYPE factor, + const thrust::complex* in, + FPTYPE* out) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if(idx >= nrxx) {return;} + if(add) { + out[idx] += factor * in[idx].real(); + } + else { + out[idx] = in[idx].real(); + } +} + template __global__ void set_real_to_recip_output( const int npwk, @@ -62,6 +80,26 @@ __global__ void set_real_to_recip_output( } } +template +__global__ void set_real_to_recip_output( + const int npwk, + const int nxyz, + const bool add, + const FPTYPE factor, + const int* box_index, + const thrust::complex* in, + FPTYPE* out) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if(idx >= npwk) {return;} + if(add) { + out[idx] += factor / nxyz * in[box_index[idx]].real(); + } + else { + out[idx] = in[box_index[idx]].real() / nxyz; + } +} + template void set_3d_fft_box_op::operator()(const base_device::DEVICE_GPU* /*dev*/, const int npwk, @@ -98,6 +136,25 @@ void set_recip_to_real_output_op::operator()(co hipCheckOnDebug(); } +template +void set_recip_to_real_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, + const int nrxx, + const bool add, + const FPTYPE factor, + const std::complex* in, + FPTYPE* out) +{ + const int block = (nrxx + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + hipLaunchKernelGGL(HIP_KERNEL_NAME(set_recip_to_real_output), dim3(block), dim3(THREADS_PER_BLOCK), 0, 0, + nrxx, + add, + factor, + reinterpret_cast*>(in), + reinterpret_cast(out)); + + hipCheckOnDebug(); +} + template void set_real_to_recip_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, const int npwk, @@ -121,6 +178,29 @@ void set_real_to_recip_output_op::operator()(co hipCheckOnDebug(); } +template +void set_real_to_recip_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, + const int npwk, + const int nxyz, + const bool add, + const FPTYPE factor, + const int* box_index, + const std::complex* in, + FPTYPE* out) +{ + const int block = (npwk + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; + hipLaunchKernelGGL(HIP_KERNEL_NAME(set_real_to_recip_output), dim3(block), dim3(THREADS_PER_BLOCK), 0, 0, + npwk, + nxyz, + add, + factor, + box_index, + reinterpret_cast*>(in), + reinterpret_cast(out)); + + hipCheckOnDebug(); +} + template struct set_3d_fft_box_op; template struct set_recip_to_real_output_op; template struct set_real_to_recip_output_op; diff --git a/source/module_basis/module_pw/pw_basis.h b/source/module_basis/module_pw/pw_basis.h index 4f1ad9a1ee8..2baf937af83 100644 --- a/source/module_basis/module_pw/pw_basis.h +++ b/source/module_basis/module_pw/pw_basis.h @@ -1,6 +1,7 @@ #ifndef PWBASIS_H #define PWBASIS_H +#include "module_base/macros.h" #include "module_base/module_device/memory_op.h" #include "module_base/matrix.h" #include "module_base/matrix3.h" @@ -245,7 +246,7 @@ class PW_Basis // FFT ft; FFT_Bundle fft_bundle; //The position of pointer in and out can be equal(in-place transform) or different(out-of-place transform). - + template void real2recip(const FPTYPE* in, std::complex* out, @@ -266,36 +267,148 @@ class PW_Basis std::complex* out, const bool add = false, const FPTYPE factor = 1.0) const; // in:(nz, ns) ; out(nplane,nx*ny) + + template + void real2recip_gpu(const FPTYPE* in, + std::complex* out, + const bool add = false, + const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny) ; out(nz, ns) + template + void real2recip_gpu(const std::complex* in, + std::complex* out, + const bool add = false, + const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny) ; out(nz, ns) + template + void recip2real_gpu(const std::complex* in, + FPTYPE* out, + const bool add = false, + const FPTYPE factor = 1.0) const; // in:(nz, ns) ; out(nplane,nx*ny) + template + void recip2real_gpu(const std::complex* in, + std::complex* out, + const bool add = false, + const FPTYPE factor = 1.0) const; // in:(nz, ns) ; out(nplane,nx*ny) - template - void real_to_recip(const std::complex* in, - std::complex* out, - const Device* ctx=get_default_device_ctx(), - const int ik=0, - const bool add = false, - const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny) ; out(nz, ns) - template - void recip_to_real(const std::complex* in, - std::complex* out, - const Device* ctx=get_default_device_ctx(), - const int ik=0, + /** + * @brief Converts data from reciprocal space to real space on Cpu + * + * This function handles the conversion of data from reciprocal space (Fourier space) to real space. + * It supports complex types as input. + * The output can be either the same fundamental type or the underlying real type of a complex type. + * + * @tparam FPTYPE The type of the input data, which can only be a compelx type (e.g., std::complex, std::complex) + * @tparam Device The device type, must be base_device::DEVICE_CPU. + * @tparam std::enable_if::type>::value, int>::type + * SFINAE constraint to ensure that FPTYPE is a complex type. + * @tparam std::enable_if::value, int>::type + * SFINAE constraint to ensure that Device is base_device::DEVICE_CPU. + * + * @param in Pointer to the input data array in reciprocal space. + * @param out Pointer to the output data array in real space. If FPTYPE is a complex type, + * this should point to an array of the underlying real type. + * @param add Boolean flag indicating whether to add the result to the existing values in the output array. + * @param factor A scaling factor applied to the output values. + */ + template ::type>::value + && (std::is_same::type>::value + || std::is_same::value) + && std::is_same::value, + int>::type + = 0> + void recip_to_real(TK* in, TR* out, const bool add = false, const typename GetTypeReal::type factor = 1.0) const + { + this->recip2real(in, out, add, factor); + }; + + /** + * @brief Converts data from reciprocal space (Fourier space) to real space. + * + * This function handles the conversion of data from reciprocal space (typically after a Fourier transform) + * to real space, supporting scenarios where the input is of a complex type and the output is of the underlying + * fundamental type. + * + * @tparam FPTYPE The underlying fundamental type of the input data (e.g., float, double). + * Note that the actual type passed should be std::complex. + * @tparam Device The device type, which can be any supported device type (e.g., base_device::DEVICE_CPU). + * + * @param in Pointer to the input data array in reciprocal space, of type std::complex*. + * @param out Pointer to the output data array in real space, of type FPTYPE*. + * @param add Optional boolean flag, default value false, indicating whether to add the result to the existing + * values in the output array. + * @param factor Optional scaling factor, default value 1.0, applied to the output values. + */ + template ::type>::value + && (std::is_same::type>::value + || std::is_same::value) + && std::is_same::value, + int>::type + = 0> + void recip_to_real(TK* in, + TR* out, const bool add = false, - const FPTYPE factor = 1.0) const; // in:(nz, ns) ; out(nplane,nx*ny) - - template - void real_to_recip(FPTYPE* in, - std::complex* out, - const Device* ctx=get_default_device_ctx(), - const int ik=0, + const typename GetTypeReal::type factor = 1.0) const + { + this->recip2real_gpu(in, out, add, factor); + }; + + // template ::type>::value, int>::type = 0> + // void recip_to_real(FPTYPE* in, + // FPTYPE* out, + // const bool add = false, + // const typename GetTypeReal::type factor = 1.0) const; + /** + * @brief Converts data from real space to reciprocal space (Fourier space). + * + * This function handles the conversion of data from real space to reciprocal space (typically after performing a + * Fourier transform), supporting scenarios where the input is of a fundamental type (e.g., float, double) and the + * output is of a complex type. + * + * @tparam FPTYPE The underlying fundamental type of the input data (e.g., float, double). + * SFINAE constraint ensures that FPTYPE is a fundamental type, not a complex type. + * @tparam Device The device type, which must be base_device::DEVICE_CPU. + * @tparam std::enable_if::type>::value, int>::type + * SFINAE constraint to ensure that FPTYPE is a fundamental type. + * @tparam std::enable_if::value, int>::type + * SFINAE constraint to ensure that Device is base_device::DEVICE_CPU. + * + * @param in Pointer to the input data array in real space, of type FPTYPE*. + * @param out Pointer to the output data array in reciprocal space, of type std::complex*. + * @param ik Optional parameter, default value 0, representing some index or identifier. + * @param add Optional boolean flag, default value false, indicating whether to add the result to the existing + * values in the output array. + * @param factor Optional scaling factor, default value 1.0, applied to the output values. + */ + template ::type>::value + && (std::is_same::type>::value || std::is_same::value) + && std::is_same::value ,int>::type = 0> + void real_to_recip(TR* in, + TK* out, const bool add = false, - const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny) ; out(nz, ns) - template - void recip_to_real(const std::complex* in, - FPTYPE* out, - const Device* ctx=get_default_device_ctx() , - const int ik=0, + const typename GetTypeReal::type factor = 1.0) const + { + this->real2recip(in, out, add, factor); + } + + template ::type>::value + && (std::is_same::type>::value || std::is_same::value) + && !std::is_same::value ,int>::type = 0> + void real_to_recip(TR* in, + TK* out, const bool add = false, - const FPTYPE factor = 1.0) const; // in:(nz, ns) ; out(nplane,nx*ny) + const typename GetTypeReal::type factor = 1.0) const; + protected: //gather planes and scatter sticks of all processors template @@ -325,6 +438,5 @@ class PW_Basis } #endif // PWBASIS_H - #include "pw_basis_sup.h" #include "pw_basis_big.h" //temporary it will be removed diff --git a/source/module_basis/module_pw/pw_transform.cpp b/source/module_basis/module_pw/pw_transform.cpp index 38d4bfc49df..4f34221775c 100644 --- a/source/module_basis/module_pw/pw_transform.cpp +++ b/source/module_basis/module_pw/pw_transform.cpp @@ -10,10 +10,10 @@ namespace ModulePW { - const base_device::DEVICE_CPU* PW_Basis::get_default_device_ctx() { - static const base_device::DEVICE_CPU* default_device_cpu; - return default_device_cpu; -} +// const base_device::DEVICE_CPU* PW_Basis::get_default_device_ctx() { +// static const base_device::DEVICE_CPU* default_device_cpu; +// return default_device_cpu; +// } /** * @brief transform real space to reciprocal space * @details c(g)=\int dr*f(r)*exp(-ig*r) @@ -284,229 +284,6 @@ void PW_Basis::recip2real(const std::complex* in, FPTYPE* out, const boo } ModuleBase::timer::tick(this->classname, "recip2real"); } - -template <> -void PW_Basis::real_to_recip(const std::complex* in, - std::complex* out, - const base_device::DEVICE_CPU* ctx, - const int ik, - const bool add, - const float factor) const -{ - this->real2recip(in, out, add, factor); -} -template <> -void PW_Basis::real_to_recip(const std::complex* in, - std::complex* out, - const base_device::DEVICE_CPU* ctx, - const int ik, - const bool add, - const double factor) const -{ - this->real2recip(in, out, add, factor); -} - -template <> -void PW_Basis::recip_to_real(const std::complex* in, - std::complex* out, - const base_device::DEVICE_CPU* ctx, - const int ik, - const bool add, - const float factor) const -{ - this->recip2real(in, out, add, factor); -} -template <> -void PW_Basis::recip_to_real(const std::complex* in, - std::complex* out, - const base_device::DEVICE_CPU* ctx, - const int ik, - const bool add, - const double factor) const -{ - this->recip2real(in, out, add, factor); -} - -template <> -void PW_Basis::real_to_recip(float* in, - std::complex* out, - const base_device::DEVICE_CPU* ctx, - const int ik, - const bool add, - const float factor) const -{ - this->real2recip(in, out, add, factor); -} -template <> -void PW_Basis::real_to_recip(double* in, - std::complex* out, - const base_device::DEVICE_CPU* ctx, - const int ik, - const bool add, - const double factor) const -{ - this->real2recip(in, out, add, factor); -} - -template <> -void PW_Basis::recip_to_real(const std::complex* in, - float* out, - const base_device::DEVICE_CPU* ctx, - const int ik, - const bool add, - const float factor) const -{ - this->recip2real(in, out, add, factor); -} - -template <> -void PW_Basis::recip_to_real(const std::complex* in, - double* out, - const base_device::DEVICE_CPU* ctx, - const int ik, - const bool add, - const double factor) const -{ - this->recip2real(in, out, add, factor); -} - -#if (defined(__CUDA) || defined(__ROCM)) -template <> -void PW_Basis::real_to_recip(const std::complex* in, - std::complex* out, - const base_device::DEVICE_GPU* ctx, - const int ik, - const bool add, - const float factor) const -{ - ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); - assert(this->gamma_only == false); - assert(this->poolnproc == 1); - - base_device::memory::synchronize_memory_op, base_device::DEVICE_GPU, base_device::DEVICE_GPU>()( - this->fft_bundle.get_auxr_3d_data(), - in, - this->nrxx); - - this->fft_bundle.fft3D_forward(ctx, - this->fft_bundle.get_auxr_3d_data(), - this->fft_bundle.get_auxr_3d_data()); - - set_real_to_recip_output_op()(ctx, - npw, - this->nxyz, - add, - factor, - this->ig2isz, - this->fft_bundle.get_auxr_3d_data(), - out); - ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); -} -template <> -void PW_Basis::real_to_recip(const std::complex* in, - std::complex* out, - const base_device::DEVICE_GPU* ctx, - const int ik, - const bool add, - const double factor) const -{ - ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); - assert(this->gamma_only == false); - assert(this->poolnproc == 1); - - base_device::memory::synchronize_memory_op, - base_device::DEVICE_GPU, - base_device::DEVICE_GPU>()(this->fft_bundle.get_auxr_3d_data(), - in, - this->nrxx); - - this->fft_bundle.fft3D_forward(ctx, - this->fft_bundle.get_auxr_3d_data(), - this->fft_bundle.get_auxr_3d_data()); - - set_real_to_recip_output_op()(ctx, - npw, - this->nxyz, - add, - factor, - this->ig2isz, - this->fft_bundle.get_auxr_3d_data(), - out); - ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); -} - -template <> -void PW_Basis::recip_to_real(const std::complex* in, - std::complex* out, - const base_device::DEVICE_GPU* ctx, - const int ik, - const bool add, - const float factor) const -{ - ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); - assert(this->gamma_only == false); - assert(this->poolnproc == 1); - // ModuleBase::GlobalFunc::ZEROS(fft_bundle.get_auxr_3d_data(), this->nxyz); - base_device::memory::set_memory_op, base_device::DEVICE_GPU>()( - this->fft_bundle.get_auxr_3d_data(), - 0, - this->nxyz); - - set_3d_fft_box_op()(ctx, - npw, - this->ig2isz, - in, - this->fft_bundle.get_auxr_3d_data()); - this->fft_bundle.fft3D_backward(ctx, - this->fft_bundle.get_auxr_3d_data(), - this->fft_bundle.get_auxr_3d_data()); - - set_recip_to_real_output_op()(ctx, - this->nrxx, - add, - factor, - this->fft_bundle.get_auxr_3d_data(), - out); - - ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); -} -template <> -void PW_Basis::recip_to_real(const std::complex* in, - std::complex* out, - const base_device::DEVICE_GPU* ctx, - const int ik, - const bool add, - const double factor) const -{ - ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); - assert(this->gamma_only == false); - assert(this->poolnproc == 1); - // ModuleBase::GlobalFunc::ZEROS(fft_bundle.get_auxr_3d_data(), this->nxyz); - base_device::memory::set_memory_op, base_device::DEVICE_GPU>()( - this->fft_bundle.get_auxr_3d_data(), - 0, - this->nxyz); - - set_3d_fft_box_op()(ctx, - npw, - this->ig2isz, - in, - this->fft_bundle.get_auxr_3d_data()); - this->fft_bundle.fft3D_backward(ctx, - this->fft_bundle.get_auxr_3d_data(), - this->fft_bundle.get_auxr_3d_data()); - - set_recip_to_real_output_op()(ctx, - this->nrxx, - add, - factor, - this->fft_bundle.get_auxr_3d_data(), - out); - - ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); -} -#endif - template void PW_Basis::real2recip(const float* in, std::complex* out, const bool add, diff --git a/source/module_basis/module_pw/pw_transform_gpu.cpp b/source/module_basis/module_pw/pw_transform_gpu.cpp new file mode 100644 index 00000000000..c17e72ca50e --- /dev/null +++ b/source/module_basis/module_pw/pw_transform_gpu.cpp @@ -0,0 +1,173 @@ +#include "pw_basis.h" +#include "module_base/timer.h" +#include "module_basis/module_pw/kernels/pw_op.h" +namespace ModulePW +{ +#if (defined(__CUDA) || defined(__ROCM)) +template +void PW_Basis::real2recip_gpu(const FPTYPE* in, + std::complex* out, + const bool add, + const FPTYPE factor) const +{ + ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); + assert(this->gamma_only == false); + assert(this->poolnproc == 1); + base_device::DEVICE_GPU* ctx; + // base_device::memory::synchronize_memory_op, + // base_device::DEVICE_GPU, + // base_device::DEVICE_GPU>()(this->fft_bundle.get_auxr_3d_data(), + // in, + // this->nrxx); + + this->fft_bundle.fft3D_forward(ctx, + this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.get_auxr_3d_data()); + + set_real_to_recip_output_op()(ctx, + npw, + this->nxyz, + add, + factor, + this->ig2isz, + this->fft_bundle.get_auxr_3d_data(), + out); + ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); +} +template +void PW_Basis::real2recip_gpu(const std::complex* in, + std::complex* out, + const bool add, + const FPTYPE factor) const +{ + ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); + assert(this->gamma_only == false); + assert(this->poolnproc == 1); + base_device::DEVICE_GPU* ctx; + base_device::memory::synchronize_memory_op, + base_device::DEVICE_GPU, + base_device::DEVICE_GPU>()(this->fft_bundle.get_auxr_3d_data(), + in, + this->nrxx); + + this->fft_bundle.fft3D_forward(ctx, + this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.get_auxr_3d_data()); + + set_real_to_recip_output_op()(ctx, + npw, + this->nxyz, + add, + factor, + this->ig2isz, + this->fft_bundle.get_auxr_3d_data(), + out); + ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); +} + +template +void PW_Basis::recip2real_gpu(const std::complex* in, + FPTYPE* out, + const bool add, + const FPTYPE factor) const +{ + ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); + assert(this->gamma_only == false); + assert(this->poolnproc == 1); + base_device::DEVICE_GPU* ctx; + // ModuleBase::GlobalFunc::ZEROS(fft_bundle.get_auxr_3d_data(), this->nxyz); + base_device::memory::set_memory_op, base_device::DEVICE_GPU>()( + this->fft_bundle.get_auxr_3d_data(), + 0, + this->nxyz); + + set_3d_fft_box_op()(ctx, + npw, + this->ig2isz, + in, + this->fft_bundle.get_auxr_3d_data()); + this->fft_bundle.fft3D_backward(ctx, + this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.get_auxr_3d_data()); + + set_recip_to_real_output_op()(ctx, + this->nrxx, + add, + factor, + this->fft_bundle.get_auxr_3d_data(), + out); + + ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); +} +template + void PW_Basis::recip2real_gpu(const std::complex *in, + std::complex *out, + const bool add, + const FPTYPE factor) const +{ + base_device::DEVICE_GPU* ctx; + ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); + assert(this->gamma_only == false); + assert(this->poolnproc == 1); + // ModuleBase::GlobalFunc::ZEROS(fft_bundle.get_auxr_3d_data(), this->nxyz); + base_device::memory::set_memory_op, base_device::DEVICE_GPU>()( + this->fft_bundle.get_auxr_3d_data(), + 0, + this->nxyz); + + set_3d_fft_box_op()(ctx, + npw, + this->ig2isz, + in, + this->fft_bundle.get_auxr_3d_data()); + this->fft_bundle.fft3D_backward(ctx, + this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.get_auxr_3d_data()); + + set_recip_to_real_output_op()(ctx, + this->nrxx, + add, + factor, + this->fft_bundle.get_auxr_3d_data(), + out); + + ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); +} +template void PW_Basis::real2recip_gpu(const double* in, + std::complex* out, + const bool add, + const double factor) const; +template void PW_Basis::real2recip_gpu(const float* in, + std::complex* out, + const bool add, + const float factor) const; + +template void PW_Basis::real2recip_gpu(const std::complex* in, + std::complex* out, + const bool add, + const double factor) const; +template void PW_Basis::real2recip_gpu(const std::complex* in, + std::complex* out, + const bool add, + const float factor) const; + +template void PW_Basis::recip2real_gpu(const std::complex* in, + double* out, + const bool add, + const double factor) const; +template void PW_Basis::recip2real_gpu(const std::complex* in, + float* out, + const bool add, + const float factor) const; + +template void PW_Basis::recip2real_gpu(const std::complex* in, + std::complex* out, + const bool add, + const double factor) const; +template void PW_Basis::recip2real_gpu(const std::complex* in, + std::complex* out, + const bool add, + const float factor) const; + +#endif +} // namespace ModulePW \ No newline at end of file diff --git a/source/module_elecstate/elecstate_pw.cpp b/source/module_elecstate/elecstate_pw.cpp index e09a1a3dfa7..ef7a168752e 100644 --- a/source/module_elecstate/elecstate_pw.cpp +++ b/source/module_elecstate/elecstate_pw.cpp @@ -448,7 +448,7 @@ void ElecStatePW::add_usrho(const psi::Psi& psi) { for (int is = 0; is < PARAM.inp.nspin; is++) { - this->charge->rhopw->recip_to_real(this->rhog[is], this->rho[is]); + this->charge->rhopw->recip2real(this->rhog[is], this->rho[is]); } } } From b96bb777cf6ee8712a73b6126e70809fe9fd05f9 Mon Sep 17 00:00:00 2001 From: ubuntu <3158793232@qq.com> Date: Fri, 21 Mar 2025 16:22:14 +0800 Subject: [PATCH 3/8] remove ctx in the file --- .../module_pw/kernels/cuda/pw_op.cu | 15 +++++-------- .../module_basis/module_pw/kernels/pw_op.cpp | 12 ++++------ source/module_basis/module_pw/kernels/pw_op.h | 20 +++++------------ .../module_pw/kernels/rocm/pw_op.hip.cu | 15 +++++-------- .../module_pw/pw_transform_gpu.cpp | 22 +++++++------------ .../module_basis/module_pw/pw_transform_k.cpp | 18 +++++---------- .../module_pw/pw_transform_k_dsp.cpp | 10 ++++----- 7 files changed, 37 insertions(+), 75 deletions(-) diff --git a/source/module_basis/module_pw/kernels/cuda/pw_op.cu b/source/module_basis/module_pw/kernels/cuda/pw_op.cu index d9418ca486d..1e0fce284e1 100644 --- a/source/module_basis/module_pw/kernels/cuda/pw_op.cu +++ b/source/module_basis/module_pw/kernels/cuda/pw_op.cu @@ -100,8 +100,7 @@ __global__ void set_real_to_recip_output( } template -void set_3d_fft_box_op::operator()(const base_device::DEVICE_GPU* /*dev*/, - const int npwk, +void set_3d_fft_box_op::operator()(const int npwk, const int* box_index, const std::complex* in, std::complex* out) @@ -117,8 +116,7 @@ void set_3d_fft_box_op::operator()(const base_d } template -void set_recip_to_real_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, - const int nrxx, +void set_recip_to_real_output_op::operator()(const int nrxx, const bool add, const FPTYPE factor, const std::complex* in, @@ -136,8 +134,7 @@ void set_recip_to_real_output_op::operator()(co } template -void set_recip_to_real_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, - const int nrxx, +void set_recip_to_real_output_op::operator()(const int nrxx, const bool add, const FPTYPE factor, const std::complex* in, @@ -155,8 +152,7 @@ void set_recip_to_real_output_op::operator()(co } template -void set_real_to_recip_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, - const int npwk, +void set_real_to_recip_output_op::operator()(const int npwk, const int nxyz, const bool add, const FPTYPE factor, @@ -178,8 +174,7 @@ void set_real_to_recip_output_op::operator()(co } template -void set_real_to_recip_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, - const int npwk, +void set_real_to_recip_output_op::operator()(const int npwk, const int nxyz, const bool add, const FPTYPE factor, diff --git a/source/module_basis/module_pw/kernels/pw_op.cpp b/source/module_basis/module_pw/kernels/pw_op.cpp index e39e8a6a7c7..a1bafae69ba 100644 --- a/source/module_basis/module_pw/kernels/pw_op.cpp +++ b/source/module_basis/module_pw/kernels/pw_op.cpp @@ -5,8 +5,7 @@ namespace ModulePW { template struct set_3d_fft_box_op { - void operator()(const base_device::DEVICE_CPU* /*dev*/, - const int npwk, + void operator()(const int npwk, const int* box_index, const std::complex* in, std::complex* out) @@ -21,8 +20,7 @@ struct set_3d_fft_box_op template struct set_recip_to_real_output_op { - void operator()(const base_device::DEVICE_CPU* /*dev*/, - const int nrxx, + void operator()(const int nrxx, const bool add, const FPTYPE factor, const std::complex* in, @@ -40,8 +38,7 @@ struct set_recip_to_real_output_op } } - void operator()(const base_device::DEVICE_CPU* /*dev*/, - const int nrxx, + void operator()(const int nrxx, const bool add, const FPTYPE factor, const std::complex* in, @@ -67,8 +64,7 @@ struct set_recip_to_real_output_op template struct set_real_to_recip_output_op { - void operator()(const base_device::DEVICE_CPU* /*dev*/, - const int npw_k, + void operator()(const int npw_k, const int nxyz, const bool add, const FPTYPE factor, diff --git a/source/module_basis/module_pw/kernels/pw_op.h b/source/module_basis/module_pw/kernels/pw_op.h index b5471906061..6221c8b68cf 100644 --- a/source/module_basis/module_pw/kernels/pw_op.h +++ b/source/module_basis/module_pw/kernels/pw_op.h @@ -19,7 +19,6 @@ struct set_3d_fft_box_op { /// Output Parameters /// @param out - output psi within the 3D box(in recip space) void operator() ( - const Device* dev, const int npwk, const int* box_index, const std::complex* in, @@ -39,7 +38,6 @@ struct set_recip_to_real_output_op { /// Output Parameters /// @param out - output psi within the 3D box(in real space) void operator() ( - const Device* dev, const int nrxx, const bool add, const FPTYPE factor, @@ -47,7 +45,6 @@ struct set_recip_to_real_output_op { std::complex* out); void operator() ( - const Device* dev, const int nrxx, const bool add, const FPTYPE factor, @@ -70,7 +67,6 @@ struct set_real_to_recip_output_op { /// Output Parameters /// @param out - output psi within the 3D box(in recip space) void operator() ( - const Device* dev, const int npw_k, const int nxyz, const bool add, @@ -80,7 +76,6 @@ struct set_real_to_recip_output_op { std::complex* out); void operator() ( - const Device* dev, const int npw_k, const int nxyz, const bool add, @@ -95,8 +90,7 @@ struct set_real_to_recip_output_op { template struct set_3d_fft_box_op { - void operator()(const base_device::DEVICE_GPU* dev, - const int npwk, + void operator()(const int npwk, const int* box_index, const std::complex* in, std::complex* out); @@ -105,15 +99,13 @@ struct set_3d_fft_box_op template struct set_recip_to_real_output_op { - void operator()(const base_device::DEVICE_GPU* dev, - const int nrxx, + void operator()(const int nrxx, const bool add, const FPTYPE factor, const std::complex* in, std::complex* out); - void operator()(const base_device::DEVICE_GPU* dev, - const int nrxx, + void operator()(const int nrxx, const bool add, const FPTYPE factor, const std::complex* in, @@ -123,16 +115,14 @@ struct set_recip_to_real_output_op template struct set_real_to_recip_output_op { - void operator()(const base_device::DEVICE_GPU* dev, - const int npw_k, + void operator()(const int npw_k, const int nxyz, const bool add, const FPTYPE factor, const int* box_index, const std::complex* in, std::complex* out); - void operator()(const base_device::DEVICE_GPU* dev, - const int npw_k, + void operator()(const int npw_k, const int nxyz, const bool add, const FPTYPE factor, diff --git a/source/module_basis/module_pw/kernels/rocm/pw_op.hip.cu b/source/module_basis/module_pw/kernels/rocm/pw_op.hip.cu index 44be7d0671f..51b50c6b8b4 100644 --- a/source/module_basis/module_pw/kernels/rocm/pw_op.hip.cu +++ b/source/module_basis/module_pw/kernels/rocm/pw_op.hip.cu @@ -101,8 +101,7 @@ __global__ void set_real_to_recip_output( } template -void set_3d_fft_box_op::operator()(const base_device::DEVICE_GPU* /*dev*/, - const int npwk, +void set_3d_fft_box_op::operator()(const int npwk, const int* box_index, const std::complex* in, std::complex* out) @@ -118,8 +117,7 @@ void set_3d_fft_box_op::operator()(const base_d } template -void set_recip_to_real_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, - const int nrxx, +void set_recip_to_real_output_op::operator()(const int nrxx, const bool add, const FPTYPE factor, const std::complex* in, @@ -137,8 +135,7 @@ void set_recip_to_real_output_op::operator()(co } template -void set_recip_to_real_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, - const int nrxx, +void set_recip_to_real_output_op::operator()(const int nrxx, const bool add, const FPTYPE factor, const std::complex* in, @@ -156,8 +153,7 @@ void set_recip_to_real_output_op::operator()(co } template -void set_real_to_recip_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, - const int npwk, +void set_real_to_recip_output_op::operator()(const int npwk, const int nxyz, const bool add, const FPTYPE factor, @@ -179,8 +175,7 @@ void set_real_to_recip_output_op::operator()(co } template -void set_real_to_recip_output_op::operator()(const base_device::DEVICE_GPU* /*dev*/, - const int npwk, +void set_real_to_recip_output_op::operator()(const int npwk, const int nxyz, const bool add, const FPTYPE factor, diff --git a/source/module_basis/module_pw/pw_transform_gpu.cpp b/source/module_basis/module_pw/pw_transform_gpu.cpp index c17e72ca50e..db89d2fbaf4 100644 --- a/source/module_basis/module_pw/pw_transform_gpu.cpp +++ b/source/module_basis/module_pw/pw_transform_gpu.cpp @@ -3,7 +3,7 @@ #include "module_basis/module_pw/kernels/pw_op.h" namespace ModulePW { -#if (defined(__CUDA) || defined(__ROCM)) +// #if (defined(__CUDA) || defined(__ROCM)) template void PW_Basis::real2recip_gpu(const FPTYPE* in, std::complex* out, @@ -24,8 +24,7 @@ void PW_Basis::real2recip_gpu(const FPTYPE* in, this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); - set_real_to_recip_output_op()(ctx, - npw, + set_real_to_recip_output_op()(npw, this->nxyz, add, factor, @@ -54,8 +53,7 @@ void PW_Basis::real2recip_gpu(const std::complex* in, this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); - set_real_to_recip_output_op()(ctx, - npw, + set_real_to_recip_output_op()(npw, this->nxyz, add, factor, @@ -81,8 +79,7 @@ void PW_Basis::recip2real_gpu(const std::complex* in, 0, this->nxyz); - set_3d_fft_box_op()(ctx, - npw, + set_3d_fft_box_op()(npw, this->ig2isz, in, this->fft_bundle.get_auxr_3d_data()); @@ -90,8 +87,7 @@ void PW_Basis::recip2real_gpu(const std::complex* in, this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); - set_recip_to_real_output_op()(ctx, - this->nrxx, + set_recip_to_real_output_op()(this->nrxx, add, factor, this->fft_bundle.get_auxr_3d_data(), @@ -115,8 +111,7 @@ template 0, this->nxyz); - set_3d_fft_box_op()(ctx, - npw, + set_3d_fft_box_op()(npw, this->ig2isz, in, this->fft_bundle.get_auxr_3d_data()); @@ -124,8 +119,7 @@ template this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); - set_recip_to_real_output_op()(ctx, - this->nrxx, + set_recip_to_real_output_op()(this->nrxx, add, factor, this->fft_bundle.get_auxr_3d_data(), @@ -169,5 +163,5 @@ template void PW_Basis::recip2real_gpu(const std::complex* in, const bool add, const float factor) const; -#endif +// #endif } // namespace ModulePW \ No newline at end of file diff --git a/source/module_basis/module_pw/pw_transform_k.cpp b/source/module_basis/module_pw/pw_transform_k.cpp index 3d75f07f6f1..4e3a9ecc9d1 100644 --- a/source/module_basis/module_pw/pw_transform_k.cpp +++ b/source/module_basis/module_pw/pw_transform_k.cpp @@ -361,8 +361,7 @@ void PW_Basis_K::real_to_recip(const base_device::DEVICE_GPU* ctx, const int startig = ik * this->npwk_max; const int npw_k = this->npwk[ik]; - set_real_to_recip_output_op()(ctx, - npw_k, + set_real_to_recip_output_op()(npw_k, this->nxyz, add, factor, @@ -393,8 +392,7 @@ void PW_Basis_K::real_to_recip(const base_device::DEVICE_GPU* ctx, const int startig = ik * this->npwk_max; const int npw_k = this->npwk[ik]; - set_real_to_recip_output_op()(ctx, - npw_k, + set_real_to_recip_output_op()(npw_k, this->nxyz, add, factor, @@ -424,15 +422,13 @@ void PW_Basis_K::recip_to_real(const base_device::DEVICE_GPU* ctx, const int startig = ik * this->npwk_max; const int npw_k = this->npwk[ik]; - set_3d_fft_box_op()(ctx, - npw_k, + set_3d_fft_box_op()(npw_k, this->ig2ixyz_k + startig, in, this->fft_bundle.get_auxr_3d_data()); this->fft_bundle.fft3D_backward(ctx, this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); - set_recip_to_real_output_op()(ctx, - this->nrxx, + set_recip_to_real_output_op()(this->nrxx, add, factor, this->fft_bundle.get_auxr_3d_data(), @@ -460,15 +456,13 @@ void PW_Basis_K::recip_to_real(const base_device::DEVICE_GPU* ctx, const int startig = ik * this->npwk_max; const int npw_k = this->npwk[ik]; - set_3d_fft_box_op()(ctx, - npw_k, + set_3d_fft_box_op()(npw_k, this->ig2ixyz_k + startig, in, this->fft_bundle.get_auxr_3d_data()); this->fft_bundle.fft3D_backward(ctx, this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); - set_recip_to_real_output_op()(ctx, - this->nrxx, + set_recip_to_real_output_op()(this->nrxx, add, factor, this->fft_bundle.get_auxr_3d_data(), diff --git a/source/module_basis/module_pw/pw_transform_k_dsp.cpp b/source/module_basis/module_pw/pw_transform_k_dsp.cpp index b292e25f0a0..1632a20c5e8 100644 --- a/source/module_basis/module_pw/pw_transform_k_dsp.cpp +++ b/source/module_basis/module_pw/pw_transform_k_dsp.cpp @@ -32,8 +32,7 @@ void PW_Basis_K::real2recip_dsp(const std::complex* in, auxr); this->fft_bundle.resource_handler(0); // copy the result from the auxr to the out ,while consider the add - set_real_to_recip_output_op()(ctx, - npw_k, + set_real_to_recip_output_op()(npw_k, this->nxyz, add, factor, @@ -58,7 +57,7 @@ void PW_Basis_K::recip2real_dsp(const std::complex* in, const int startig = ik * this->npwk_max; const int npw_k = this->npwk[ik]; // copy the mapping form the type of stick to the 3dfft - set_3d_fft_box_op()(ctx, npw_k, this->ig2ixyz_k_cpu.data() + startig, in, auxr); + set_3d_fft_box_op()(npw_k, this->ig2ixyz_k_cpu.data() + startig, in, auxr); // use 3d fft backward this->fft_bundle.resource_handler(1); this->fft_bundle.fft3D_backward(gpux, auxr, auxr); @@ -107,7 +106,7 @@ void PW_Basis_K::convolution(const base_device::DEVICE_CPU* ctx, const int npw_k = this->npwk[ik]; // copy the mapping form the type of stick to the 3dfft - set_3d_fft_box_op()(ctx, npw_k, this->ig2ixyz_k_cpu.data() + startig, input, auxr); + set_3d_fft_box_op()(npw_k, this->ig2ixyz_k_cpu.data() + startig, input, auxr); // use 3d fft backward this->fft_bundle.fft3D_backward(gpux, auxr, auxr); @@ -120,8 +119,7 @@ void PW_Basis_K::convolution(const base_device::DEVICE_CPU* ctx, // 3d fft this->fft_bundle.fft3D_forward(gpux, auxr, auxr); // copy the result from the auxr to the out ,while consider the add - set_real_to_recip_output_op()(ctx, - npw_k, + set_real_to_recip_output_op()(npw_k, this->nxyz, add, factor, From a105a3cd92f40778404fc5928daf2121d778625d Mon Sep 17 00:00:00 2001 From: ubuntu <3158793232@qq.com> Date: Fri, 21 Mar 2025 16:42:04 +0800 Subject: [PATCH 4/8] remove ctx in bundle --- .../module_pw/kernels/test/pw_op_test.cpp | 12 ++++++------ .../module_basis/module_pw/module_fft/fft_bundle.cpp | 12 ++++-------- .../module_basis/module_pw/module_fft/fft_bundle.h | 8 ++++---- source/module_basis/module_pw/pw_transform_gpu.cpp | 12 ++++-------- source/module_basis/module_pw/pw_transform_k.cpp | 8 ++++---- source/module_basis/module_pw/pw_transform_k_dsp.cpp | 9 ++++----- 6 files changed, 26 insertions(+), 35 deletions(-) diff --git a/source/module_basis/module_pw/kernels/test/pw_op_test.cpp b/source/module_basis/module_pw/kernels/test/pw_op_test.cpp index 6adac4613fe..e9bc07bd9ea 100644 --- a/source/module_basis/module_pw/kernels/test/pw_op_test.cpp +++ b/source/module_basis/module_pw/kernels/test/pw_op_test.cpp @@ -72,7 +72,7 @@ class TestModulePWPWMultiDevice : public ::testing::Test TEST_F(TestModulePWPWMultiDevice, set_3d_fft_box_op_cpu) { std::vector> res(out_1.size(), std::complex{0, 0}); - set_3d_fft_box_cpu_op()(cpu_ctx, this->npwk, box_index.data(), in_1.data(), res.data()); + set_3d_fft_box_cpu_op()(this->npwk, box_index.data(), in_1.data(), res.data()); for (int ii = 0; ii < this->nxyz; ii++) { EXPECT_LT(std::abs(res[ii] - out_1[ii]), 1e-12); } @@ -81,7 +81,7 @@ TEST_F(TestModulePWPWMultiDevice, set_3d_fft_box_op_cpu) TEST_F(TestModulePWPWMultiDevice, set_recip_to_real_output_op_cpu) { std::vector> res(out_2.size(), std::complex{0, 0}); - set_recip_to_real_output_cpu_op()(cpu_ctx, this->nxyz, this->add, this->factor, in_2.data(), res.data()); + set_recip_to_real_output_cpu_op()(this->nxyz, this->add, this->factor, in_2.data(), res.data()); for (int ii = 0; ii < this->nxyz; ii++) { EXPECT_LT(std::abs(res[ii] - out_2[ii]), 1e-12); } @@ -90,7 +90,7 @@ TEST_F(TestModulePWPWMultiDevice, set_recip_to_real_output_op_cpu) TEST_F(TestModulePWPWMultiDevice, set_real_to_recip_output_op_cpu) { std::vector> res = out_3_init; - set_real_to_recip_output_cpu_op()(cpu_ctx, this->npwk, this->nxyz, true, this->factor, box_index.data(), in_3.data(), res.data()); + set_real_to_recip_output_cpu_op()(this->npwk, this->nxyz, true, this->factor, box_index.data(), in_3.data(), res.data()); for (int ii = 0; ii < out_3.size(); ii++) { EXPECT_LT(std::abs(res[ii] - out_3[ii]), 5e-6); } @@ -109,7 +109,7 @@ TEST_F(TestModulePWPWMultiDevice, set_3d_fft_box_op_gpu) synchronize_memory_complex_h2d_op()(d_res, res.data(), res.size()); synchronize_memory_complex_h2d_op()(d_in_1, in_1.data(), in_1.size()); - set_3d_fft_box_gpu_op()(gpu_ctx, this->npwk, d_box_index, d_in_1, d_res); + set_3d_fft_box_gpu_op()(this->npwk, d_box_index, d_in_1, d_res); synchronize_memory_complex_d2h_op()(res.data(), d_res, res.size()); @@ -130,7 +130,7 @@ TEST_F(TestModulePWPWMultiDevice, set_recip_to_real_output_op_gpu) synchronize_memory_complex_h2d_op()(d_res, res.data(), res.size()); synchronize_memory_complex_h2d_op()(d_in_2, in_2.data(), in_2.size()); - set_recip_to_real_output_gpu_op()(gpu_ctx, this->nxyz, this->add, this->factor, d_in_2, d_res); + set_recip_to_real_output_gpu_op()(this->nxyz, this->add, this->factor, d_in_2, d_res); synchronize_memory_complex_d2h_op()(res.data(), d_res, res.size()); @@ -153,7 +153,7 @@ TEST_F(TestModulePWPWMultiDevice, set_real_to_recip_output_op_gpu) synchronize_memory_complex_h2d_op()(d_res, res.data(), res.size()); synchronize_memory_complex_h2d_op()(d_in_3, in_3.data(), in_3.size()); - set_real_to_recip_output_gpu_op()(gpu_ctx, this->npwk, this->nxyz, true, this->factor, d_box_index, d_in_3, d_res); + set_real_to_recip_output_gpu_op()(this->npwk, this->nxyz, true, this->factor, d_box_index, d_in_3, d_res); synchronize_memory_complex_d2h_op()(res.data(), d_res, res.size()); diff --git a/source/module_basis/module_pw/module_fft/fft_bundle.cpp b/source/module_basis/module_pw/module_fft/fft_bundle.cpp index 7289e8ab029..270f9854980 100644 --- a/source/module_basis/module_pw/module_fft/fft_bundle.cpp +++ b/source/module_basis/module_pw/module_fft/fft_bundle.cpp @@ -227,30 +227,26 @@ void FFT_Bundle::fftxyc2r(std::complex* in, double* out) const } template <> -void FFT_Bundle::fft3D_forward(const base_device::DEVICE_GPU* ctx, - std::complex* in, +void FFT_Bundle::fft3D_forward(std::complex* in, std::complex* out) const { fft_float->fft3D_forward(in, out); } template <> -void FFT_Bundle::fft3D_forward(const base_device::DEVICE_GPU* ctx, - std::complex* in, +void FFT_Bundle::fft3D_forward(std::complex* in, std::complex* out) const { fft_double->fft3D_forward(in, out); } template <> -void FFT_Bundle::fft3D_backward(const base_device::DEVICE_GPU* ctx, - std::complex* in, +void FFT_Bundle::fft3D_backward(std::complex* in, std::complex* out) const { fft_float->fft3D_backward(in, out); } template <> -void FFT_Bundle::fft3D_backward(const base_device::DEVICE_GPU* ctx, - std::complex* in, +void FFT_Bundle::fft3D_backward(std::complex* in, std::complex* out) const { fft_double->fft3D_backward(in, out); diff --git a/source/module_basis/module_pw/module_fft/fft_bundle.h b/source/module_basis/module_pw/module_fft/fft_bundle.h index 1982a79a0c7..5b8b527cc40 100644 --- a/source/module_basis/module_pw/module_fft/fft_bundle.h +++ b/source/module_basis/module_pw/module_fft/fft_bundle.h @@ -188,10 +188,10 @@ class FFT_Bundle template void fftxyc2r(std::complex* in, FPTYPE* out) const; - template - void fft3D_forward(const Device* ctx, std::complex* in, std::complex* out) const; - template - void fft3D_backward(const Device* ctx, std::complex* in, std::complex* out) const; + template + void fft3D_forward(std::complex* in, std::complex* out) const; + template + void fft3D_backward(std::complex* in, std::complex* out) const; private: int fft_mode = 0; diff --git a/source/module_basis/module_pw/pw_transform_gpu.cpp b/source/module_basis/module_pw/pw_transform_gpu.cpp index db89d2fbaf4..b5970a48a66 100644 --- a/source/module_basis/module_pw/pw_transform_gpu.cpp +++ b/source/module_basis/module_pw/pw_transform_gpu.cpp @@ -20,8 +20,7 @@ void PW_Basis::real2recip_gpu(const FPTYPE* in, // in, // this->nrxx); - this->fft_bundle.fft3D_forward(ctx, - this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.fft3D_forward(this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); set_real_to_recip_output_op()(npw, @@ -49,8 +48,7 @@ void PW_Basis::real2recip_gpu(const std::complex* in, in, this->nrxx); - this->fft_bundle.fft3D_forward(ctx, - this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.fft3D_forward(this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); set_real_to_recip_output_op()(npw, @@ -83,8 +81,7 @@ void PW_Basis::recip2real_gpu(const std::complex* in, this->ig2isz, in, this->fft_bundle.get_auxr_3d_data()); - this->fft_bundle.fft3D_backward(ctx, - this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.fft3D_backward(this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); set_recip_to_real_output_op()(this->nrxx, @@ -115,8 +112,7 @@ template this->ig2isz, in, this->fft_bundle.get_auxr_3d_data()); - this->fft_bundle.fft3D_backward(ctx, - this->fft_bundle.get_auxr_3d_data(), + this->fft_bundle.fft3D_backward(this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); set_recip_to_real_output_op()(this->nrxx, diff --git a/source/module_basis/module_pw/pw_transform_k.cpp b/source/module_basis/module_pw/pw_transform_k.cpp index 4e3a9ecc9d1..776b493d35c 100644 --- a/source/module_basis/module_pw/pw_transform_k.cpp +++ b/source/module_basis/module_pw/pw_transform_k.cpp @@ -357,7 +357,7 @@ void PW_Basis_K::real_to_recip(const base_device::DEVICE_GPU* ctx, in, this->nrxx); - this->fft_bundle.fft3D_forward(ctx, this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); + this->fft_bundle.fft3D_forward(this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); const int startig = ik * this->npwk_max; const int npw_k = this->npwk[ik]; @@ -388,7 +388,7 @@ void PW_Basis_K::real_to_recip(const base_device::DEVICE_GPU* ctx, in, this->nrxx); - this->fft_bundle.fft3D_forward(ctx, this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); + this->fft_bundle.fft3D_forward(this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); const int startig = ik * this->npwk_max; const int npw_k = this->npwk[ik]; @@ -426,7 +426,7 @@ void PW_Basis_K::recip_to_real(const base_device::DEVICE_GPU* ctx, this->ig2ixyz_k + startig, in, this->fft_bundle.get_auxr_3d_data()); - this->fft_bundle.fft3D_backward(ctx, this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); + this->fft_bundle.fft3D_backward(this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); set_recip_to_real_output_op()(this->nrxx, add, @@ -460,7 +460,7 @@ void PW_Basis_K::recip_to_real(const base_device::DEVICE_GPU* ctx, this->ig2ixyz_k + startig, in, this->fft_bundle.get_auxr_3d_data()); - this->fft_bundle.fft3D_backward(ctx, this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); + this->fft_bundle.fft3D_backward(this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); set_recip_to_real_output_op()(this->nrxx, add, diff --git a/source/module_basis/module_pw/pw_transform_k_dsp.cpp b/source/module_basis/module_pw/pw_transform_k_dsp.cpp index 1632a20c5e8..56ba26eb2af 100644 --- a/source/module_basis/module_pw/pw_transform_k_dsp.cpp +++ b/source/module_basis/module_pw/pw_transform_k_dsp.cpp @@ -27,8 +27,7 @@ void PW_Basis_K::real2recip_dsp(const std::complex* in, // 3d fft this->fft_bundle.resource_handler(1); - this->fft_bundle.fft3D_forward(gpux, - auxr, + this->fft_bundle.fft3D_forward(auxr, auxr); this->fft_bundle.resource_handler(0); // copy the result from the auxr to the out ,while consider the add @@ -60,7 +59,7 @@ void PW_Basis_K::recip2real_dsp(const std::complex* in, set_3d_fft_box_op()(npw_k, this->ig2ixyz_k_cpu.data() + startig, in, auxr); // use 3d fft backward this->fft_bundle.resource_handler(1); - this->fft_bundle.fft3D_backward(gpux, auxr, auxr); + this->fft_bundle.fft3D_backward(auxr, auxr); this->fft_bundle.resource_handler(0); if (add) { @@ -109,7 +108,7 @@ void PW_Basis_K::convolution(const base_device::DEVICE_CPU* ctx, set_3d_fft_box_op()(npw_k, this->ig2ixyz_k_cpu.data() + startig, input, auxr); // use 3d fft backward - this->fft_bundle.fft3D_backward(gpux, auxr, auxr); + this->fft_bundle.fft3D_backward(auxr, auxr); for (int ir = 0; ir < size; ir++) { @@ -117,7 +116,7 @@ void PW_Basis_K::convolution(const base_device::DEVICE_CPU* ctx, } // 3d fft - this->fft_bundle.fft3D_forward(gpux, auxr, auxr); + this->fft_bundle.fft3D_forward(auxr, auxr); // copy the result from the auxr to the out ,while consider the add set_real_to_recip_output_op()(npw_k, this->nxyz, From 977b713f5fb103056e152d7995120d3882051ed1 Mon Sep 17 00:00:00 2001 From: ubuntu <3158793232@qq.com> Date: Fri, 21 Mar 2025 16:49:35 +0800 Subject: [PATCH 5/8] update compile bug --- source/module_basis/module_pw/pw_transform_gpu.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/module_basis/module_pw/pw_transform_gpu.cpp b/source/module_basis/module_pw/pw_transform_gpu.cpp index b5970a48a66..1898ff23ebb 100644 --- a/source/module_basis/module_pw/pw_transform_gpu.cpp +++ b/source/module_basis/module_pw/pw_transform_gpu.cpp @@ -3,7 +3,7 @@ #include "module_basis/module_pw/kernels/pw_op.h" namespace ModulePW { -// #if (defined(__CUDA) || defined(__ROCM)) +#if (defined(__CUDA) || defined(__ROCM)) template void PW_Basis::real2recip_gpu(const FPTYPE* in, std::complex* out, @@ -159,5 +159,5 @@ template void PW_Basis::recip2real_gpu(const std::complex* in, const bool add, const float factor) const; -// #endif +#endif } // namespace ModulePW \ No newline at end of file From ee48127a32eaae8c139373f69ca8582c88c6e0d5 Mon Sep 17 00:00:00 2001 From: ubuntu <3158793232@qq.com> Date: Sun, 23 Mar 2025 19:53:19 +0800 Subject: [PATCH 6/8] add T,device --- source/module_basis/module_pw/pw_basis_k.h | 72 ++++++++++++++- .../module_basis/module_pw/pw_transform_k.cpp | 89 +++++++++++++++++++ .../module_pw/test/test-other.cpp | 8 +- source/module_elecstate/elecstate_pw.cpp | 8 +- .../module_elecstate/elecstate_pw_cal_tau.cpp | 4 +- .../module_xc/test/xc3_mock.h | 18 ++-- .../module_xc/xc_functional_gradcorr.cpp | 3 +- .../hamilt_ofdft/ml_data_descriptor.cpp | 4 +- .../hamilt_pwdft/operator_pw/meta_pw.cpp | 4 +- .../hamilt_pwdft/operator_pw/op_exx_pw.cpp | 10 +-- .../hamilt_pwdft/operator_pw/veff_pw.cpp | 12 +-- .../hamilt_stodft/sto_iter.cpp | 2 +- source/module_io/get_pchg_pw.h | 4 +- 13 files changed, 195 insertions(+), 43 deletions(-) diff --git a/source/module_basis/module_pw/pw_basis_k.h b/source/module_basis/module_pw/pw_basis_k.h index ae5076bba9e..ef578852ae3 100644 --- a/source/module_basis/module_pw/pw_basis_k.h +++ b/source/module_basis/module_pw/pw_basis_k.h @@ -161,7 +161,7 @@ class PW_Basis_K : public PW_Basis #endif - template + template void real_to_recip(const Device* ctx, const std::complex* in, std::complex* out, @@ -176,6 +176,76 @@ class PW_Basis_K : public PW_Basis const bool add = false, const FPTYPE factor = 1.0) const; // in:(nz, ns) ; out(nplane,nx*ny) + + template ::value, int>::type = 0> + void real_to_recip(const TK* in, + TK* out, + const int ik, + const bool add = false, + const typename GetTypeReal::type factor = 1.0) const + { + #if defined(__DSP) + this->recip2real_dsp(in, out, ik, add, factor); + #else + this->real2recip(in,out,ik,add,factor); + #endif + } + template ::value, int>::type = 0> + void recip_to_real(const TK* in, + TK* out, + const int ik, + const bool add = false, + const typename GetTypeReal::type factor = 1.0) const + { + + #if defined(__DSP) + this->recip2real_dsp(in,out,ik,add,factor); + #else + this->recip2real(in,out,ik,add,factor); + #endif + } + template + void real2recip_gpu(const std::complex* in, + std::complex* out, + const int ik, + const bool add = false, + const FPTYPE factor = 1.0) const; // in:(nplane,nx*ny) ; out(nz, ns) + + template + void recip2real_gpu(const std::complex* in, + std::complex* out, + const int ik, + const bool add = false, + const FPTYPE factor = 1.0) const; // in:(nz, ns) ; out(nplane,nx*ny) + + template ::value, int>::type = 0> + void real_to_recip(const FPTYPE* in, + FPTYPE* out, + const int ik, + const bool add = false, + const typename GetTypeReal::type factor = 1.0) const + { + this->real2recip_gpu(in, out, ik, add, factor); + } + + template ::value, int>::type = 0> + void recip_to_real(const TK* in, + TK* out, + const int ik, + const bool add = false, + const typename GetTypeReal::type factor = 1.0) const + { + this->recip2real_gpu(in, out, ik, add, factor); + } + public: //operator: //get (G+K)^2: diff --git a/source/module_basis/module_pw/pw_transform_k.cpp b/source/module_basis/module_pw/pw_transform_k.cpp index 776b493d35c..a709b60429b 100644 --- a/source/module_basis/module_pw/pw_transform_k.cpp +++ b/source/module_basis/module_pw/pw_transform_k.cpp @@ -470,6 +470,95 @@ void PW_Basis_K::recip_to_real(const base_device::DEVICE_GPU* ctx, ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); } + +template +void PW_Basis_K::real2recip_gpu(const std::complex* in, + std::complex* out, + const int ik, + const bool add, + const FPTYPE factor) const +{ + ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); + assert(this->gamma_only == false); + assert(this->poolnproc == 1); + + base_device::memory::synchronize_memory_op, + base_device::DEVICE_GPU, + base_device::DEVICE_GPU>()(this->fft_bundle.get_auxr_3d_data(), + in, + this->nrxx); + + this->fft_bundle.fft3D_forward(this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); + + const int startig = ik * this->npwk_max; + const int npw_k = this->npwk[ik]; + set_real_to_recip_output_op()(npw_k, + this->nxyz, + add, + factor, + this->ig2ixyz_k + startig, + this->fft_bundle.get_auxr_3d_data(), + out); + ModuleBase::timer::tick(this->classname, "real_to_recip gpu"); +} +template +void PW_Basis_K::recip2real_gpu(const std::complex* in, + std::complex* out, + const int ik, + const bool add, + const FPTYPE factor) const +{ + ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); + assert(this->gamma_only == false); + assert(this->poolnproc == 1); + // ModuleBase::GlobalFunc::ZEROS(fft_bundle.get_auxr_3d_data(), this->nxyz); + base_device::memory::set_memory_op, base_device::DEVICE_GPU>()( + this->fft_bundle.get_auxr_3d_data(), + 0, + this->nxyz); + + const int startig = ik * this->npwk_max; + const int npw_k = this->npwk[ik]; + + set_3d_fft_box_op()(npw_k, + this->ig2ixyz_k + startig, + in, + this->fft_bundle.get_auxr_3d_data()); + this->fft_bundle.fft3D_backward(this->fft_bundle.get_auxr_3d_data(), this->fft_bundle.get_auxr_3d_data()); + + set_recip_to_real_output_op()(this->nrxx, + add, + factor, + this->fft_bundle.get_auxr_3d_data(), + out); + + ModuleBase::timer::tick(this->classname, "recip_to_real gpu"); +} + +template void PW_Basis_K::real2recip_gpu(const std::complex*, + std::complex*, + const int, + const bool, + const float) const; + +template void PW_Basis_K::real2recip_gpu(const std::complex*, + std::complex*, + const int, + const bool, + const double) const; + +template void PW_Basis_K::recip2real_gpu(const std::complex*, + std::complex*, + const int, + const bool, + const float) const; + +template void PW_Basis_K::recip2real_gpu(const std::complex*, + std::complex*, + const int, + const bool, + const double) const; + #endif template void PW_Basis_K::real2recip(const float* in, diff --git a/source/module_basis/module_pw/test/test-other.cpp b/source/module_basis/module_pw/test/test-other.cpp index 308f4c6c689..c5eb0093b27 100644 --- a/source/module_basis/module_pw/test/test-other.cpp +++ b/source/module_basis/module_pw/test/test-other.cpp @@ -76,26 +76,26 @@ TEST_F(PWTEST,test_other) } #endif - pwktest.recip_to_real(ctx, rhog1, rhor1, ik); + pwktest.recip_to_real,base_device::DEVICE_CPU>(rhog1, rhor1, ik); pwktest.recip2real(rhog2, rhor2, ik); for(int ir = 0 ; ir < nrxx; ++ir) { EXPECT_NEAR(std::abs(rhor1[ir]),std::abs(rhor2[ir]),1e-8); } - pwktest.real_to_recip(ctx, rhor1, rhog1, ik); + pwktest.real_to_recip,base_device::DEVICE_CPU>(rhor1, rhog1, ik); pwktest.real2recip(rhor2, rhog2, ik); for(int ig = 0 ; ig < npwk; ++ig) { EXPECT_NEAR(std::abs(rhog1[ig]),std::abs(rhog2[ig]),1e-8); } #ifdef __ENABLE_FLOAT_FFTW - pwktest.recip_to_real(ctx, rhofg1, rhofr1, ik); + pwktest.recip_to_real,base_device::DEVICE_CPU>(rhofg1, rhofr1, ik); pwktest.recip2real(rhofg2, rhofr2, ik); for(int ir = 0 ; ir < nrxx; ++ir) { EXPECT_NEAR(std::abs(rhofr1[ir]),std::abs(rhofr2[ir]),1e-6); } - pwktest.real_to_recip(ctx, rhofr1, rhofg1, ik); + pwktest.real_to_recip,base_device::DEVICE_CPU>(rhofr1, rhofg1, ik); pwktest.real2recip(rhofr2, rhofg2, ik); for(int ig = 0 ; ig < npwk; ++ig) { diff --git a/source/module_elecstate/elecstate_pw.cpp b/source/module_elecstate/elecstate_pw.cpp index ef7a168752e..76a4ce6c42b 100644 --- a/source/module_elecstate/elecstate_pw.cpp +++ b/source/module_elecstate/elecstate_pw.cpp @@ -202,9 +202,9 @@ void ElecStatePW::rhoBandK(const psi::Psi& psi) /// be care of when smearing_sigma is large, wg would less than 0 /// - this->basis->recip_to_real(this->ctx, &psi(ibnd,0), this->wfcr, ik); + this->basis->recip_to_real( &psi(ibnd,0), this->wfcr, ik); - this->basis->recip_to_real(this->ctx, &psi(ibnd,npwx), this->wfcr_another_spin, ik); + this->basis->recip_to_real( &psi(ibnd,npwx), this->wfcr_another_spin, ik); const auto w1 = static_cast(this->wg(ik, ibnd) / ucell->omega); @@ -230,7 +230,7 @@ void ElecStatePW::rhoBandK(const psi::Psi& psi) /// only occupied band should be calculated. /// - this->basis->recip_to_real(this->ctx, &psi(ibnd,0), this->wfcr, ik); + this->basis->recip_to_real(&psi(ibnd,0), this->wfcr, ik); const auto w1 = static_cast(this->wg(ik, ibnd) / ucell->omega); @@ -258,7 +258,7 @@ void ElecStatePW::rhoBandK(const psi::Psi& psi) &psi(ibnd, 0), this->wfcr); - this->basis->recip_to_real(this->ctx, this->wfcr, this->wfcr, ik); + this->basis->recip_to_real( this->wfcr, this->wfcr, ik); elecstate_pw_op()(this->ctx, current_spin, this->charge->nrxx, w1, this->kin_r, this->wfcr); } diff --git a/source/module_elecstate/elecstate_pw_cal_tau.cpp b/source/module_elecstate/elecstate_pw_cal_tau.cpp index 628dd25aeff..1cf4f57b700 100644 --- a/source/module_elecstate/elecstate_pw_cal_tau.cpp +++ b/source/module_elecstate/elecstate_pw_cal_tau.cpp @@ -23,7 +23,7 @@ void ElecStatePW::cal_tau(const psi::Psi& psi) int nbands = psi.get_nbands(); for (int ibnd = 0; ibnd < nbands; ibnd++) { - this->basis->recip_to_real(this->ctx, &psi(ibnd,0), this->wfcr, ik); + this->basis->recip_to_real(&psi(ibnd,0), this->wfcr, ik); const auto w1 = static_cast(this->wg(ik, ibnd) / ucell->omega); @@ -43,7 +43,7 @@ void ElecStatePW::cal_tau(const psi::Psi& psi) &psi(ibnd, 0), this->wfcr); - this->basis->recip_to_real(this->ctx, this->wfcr, this->wfcr, ik); + this->basis->recip_to_real(this->wfcr, this->wfcr, ik); elecstate_pw_op()(this->ctx, current_spin, this->charge->nrxx, w1, this->kin_r, this->wfcr); } diff --git a/source/module_hamilt_general/module_xc/test/xc3_mock.h b/source/module_hamilt_general/module_xc/test/xc3_mock.h index ffca76c70cb..ebbc139d75c 100644 --- a/source/module_hamilt_general/module_xc/test/xc3_mock.h +++ b/source/module_hamilt_general/module_xc/test/xc3_mock.h @@ -78,8 +78,7 @@ namespace ModulePW template - void PW_Basis_K::real_to_recip(const Device* ctx, - const std::complex* in, + void PW_Basis_K::real_to_recip(const std::complex* in, std::complex* out, const int ik, const bool add, @@ -91,8 +90,7 @@ namespace ModulePW } } template - void PW_Basis_K::recip_to_real(const Device* ctx, - const std::complex* in, + void PW_Basis_K::recip_to_real(const std::complex* in, std::complex* out, const int ik, const bool add, @@ -104,28 +102,24 @@ namespace ModulePW } } - template void PW_Basis_K::real_to_recip(const base_device::DEVICE_CPU* ctx, - const std::complex* in, + template void PW_Basis_K::real_to_recip(const std::complex* in, std::complex* out, const int ik, const bool add, const double factor) const; - template void PW_Basis_K::recip_to_real(const base_device::DEVICE_CPU* ctx, - const std::complex* in, + template void PW_Basis_K::recip_to_real(const std::complex* in, std::complex* out, const int ik, const bool add, const double factor) const; #if __CUDA || __ROCM - template void PW_Basis_K::real_to_recip(const base_device::DEVICE_GPU* ctx, - const std::complex* in, + template void PW_Basis_K::real_to_recip(const std::complex* in, std::complex* out, const int ik, const bool add, const double factor) const; - template void PW_Basis_K::recip_to_real(const base_device::DEVICE_GPU* ctx, - const std::complex* in, + template void PW_Basis_K::recip_to_real(const std::complex* in, std::complex* out, const int ik, const bool add, diff --git a/source/module_hamilt_general/module_xc/xc_functional_gradcorr.cpp b/source/module_hamilt_general/module_xc/xc_functional_gradcorr.cpp index 42501465108..b3cad37c448 100644 --- a/source/module_hamilt_general/module_xc/xc_functional_gradcorr.cpp +++ b/source/module_hamilt_general/module_xc/xc_functional_gradcorr.cpp @@ -644,8 +644,7 @@ void XC_Functional::grad_wfc( rhog, porter.data()); // Array of std::complex // bring the gdr from G --> R - Device * ctx = nullptr; - wfc_basis->recip_to_real(ctx, porter.data(), porter.data(), ik); + wfc_basis->recip_to_real(porter.data(), porter.data(), ik); xc_functional_grad_wfc_solver( ipol, wfc_basis->nrxx, // Integers diff --git a/source/module_hamilt_pw/hamilt_ofdft/ml_data_descriptor.cpp b/source/module_hamilt_pw/hamilt_ofdft/ml_data_descriptor.cpp index 85135d789f7..e9ed4e7268e 100644 --- a/source/module_hamilt_pw/hamilt_ofdft/ml_data_descriptor.cpp +++ b/source/module_hamilt_pw/hamilt_ofdft/ml_data_descriptor.cpp @@ -197,7 +197,7 @@ void ML_data::getF_KS1( continue; } - pw_psi->recip_to_real(ctx, &psi->operator()(ibnd,0), wfcr, ik); + pw_psi->recip_to_real( &psi->operator()(ibnd,0), wfcr, ik); const double w1 = pelec->wg(ik, ibnd) / ucell.omega; // output one wf, to check KS equation @@ -308,7 +308,7 @@ void ML_data::getF_KS2( continue; } - pw_psi->recip_to_real(ctx, &psi->operator()(ibnd,0), wfcr, ik); + pw_psi->recip_to_real( &psi->operator()(ibnd,0), wfcr, ik); const double w1 = pelec->wg(ik, ibnd) / ucell.omega; if (pelec->ekb(ik,ibnd) > epsilonM) diff --git a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/meta_pw.cpp b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/meta_pw.cpp index 817dba61cea..538819f4522 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/meta_pw.cpp +++ b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/meta_pw.cpp @@ -67,13 +67,13 @@ void Meta>::act( for (int j = 0; j < 3; j++) { meta_op()(this->ctx, this->ik, j, ngk_ik, this->wfcpw->npwk_max, this->tpiba, wfcpw->get_gcar_data(), wfcpw->get_kvec_c_data(), tmpsi_in, this->porter); - wfcpw->recip_to_real(this->ctx, this->porter, this->porter, this->ik); + wfcpw->recip_to_real(this->porter, this->porter, this->ik); if(this->vk_col != 0) { vector_mul_vector_op()(this->vk_col, this->porter, this->porter, this->vk + current_spin * this->vk_col); } - wfcpw->real_to_recip(this->ctx, this->porter, this->porter, this->ik); + wfcpw->real_to_recip(this->porter, this->porter, this->ik); meta_op()(this->ctx, this->ik, j, ngk_ik, this->wfcpw->npwk_max, this->tpiba, wfcpw->get_gcar_data(), wfcpw->get_kvec_c_data(), this->porter, tmhpsi, true); } // x,y,z directions diff --git a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.cpp b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.cpp index 4b6d36908a0..202800026c7 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.cpp +++ b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.cpp @@ -188,7 +188,7 @@ void OperatorEXXPW::act_op(const int nbands, { const T *psi_nk = tmpsi_in + n_iband * nbasis; // retrieve \psi_nk in real space - wfcpw->recip_to_real(ctx, psi_nk, psi_nk_real, this->ik); + wfcpw->recip_to_real( psi_nk, psi_nk_real, this->ik); // for \psi_nk, get the pw of iq and band m auto q_points = get_q_points(this->ik); @@ -208,7 +208,7 @@ void OperatorEXXPW::act_op(const int nbands, // if (has_real.find({iq, m_iband}) == has_real.end()) // { const T* psi_mq = get_pw(m_iband, iq); - wfcpw->recip_to_real(ctx, psi_mq, psi_mq_real, iq); + wfcpw->recip_to_real( psi_mq, psi_mq_real, iq); // syncmem_complex_op()(this->ctx, this->ctx, psi_all_real + m_iband * wfcpw->nrxx, psi_mq_real, wfcpw->nrxx); // has_real[{iq, m_iband}] = true; // } @@ -271,7 +271,7 @@ void OperatorEXXPW::act_op(const int nbands, } // end of iq auto h_psi_nk = tmhpsi + n_iband * nbasis; Real hybrid_alpha = GlobalC::exx_info.info_global.hybrid_alpha; - wfcpw->real_to_recip(ctx, h_psi_real, h_psi_nk, this->ik, true, hybrid_alpha); + wfcpw->real_to_recip(h_psi_real, h_psi_nk, this->ik, true, hybrid_alpha); setmem_complex_op()(h_psi_real, 0, rhopw->nrxx); } @@ -810,7 +810,7 @@ double OperatorEXXPW::cal_exx_energy_op(psi::Psi *ppsi_) c psi.fix_kb(ik, n_iband); const T* psi_nk = psi.get_pointer(); // retrieve \psi_nk in real space - wfcpw->recip_to_real(ctx, psi_nk, psi_nk_real, ik); + wfcpw->recip_to_real( psi_nk, psi_nk_real, ik); // for \psi_nk, get the pw of iq and band m // q_points is a vector of integers, 0 to nks-1 @@ -839,7 +839,7 @@ double OperatorEXXPW::cal_exx_energy_op(psi::Psi *ppsi_) c psi_.fix_kb(iq, m_iband); const T* psi_mq = psi_.get_pointer(); // const T* psi_mq = get_pw(m_iband, iq); - wfcpw->recip_to_real(ctx, psi_mq, psi_mq_real, iq); + wfcpw->recip_to_real(psi_mq, psi_mq_real, iq); T omega_inv = 1.0 / ucell->omega; diff --git a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/veff_pw.cpp b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/veff_pw.cpp index 54e1a052be2..c931d61f2f8 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/veff_pw.cpp +++ b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/veff_pw.cpp @@ -62,7 +62,7 @@ void Veff>::act( if (npol == 1) { // wfcpw->recip2real(tmpsi_in, porter, this->ik); - wfcpw->recip_to_real(this->ctx, tmpsi_in, this->porter, this->ik); + wfcpw->recip_to_real(tmpsi_in, this->porter, this->ik); // NOTICE: when MPI threads are larger than number of Z grids // veff would contain nothing, and nothing should be done in real space // but the 3DFFT can not be skipped, it will cause hanging @@ -76,7 +76,7 @@ void Veff>::act( // } } // wfcpw->real2recip(porter, tmhpsi, this->ik, true); - wfcpw->real_to_recip(this->ctx, this->porter, tmhpsi, this->ik, true); + wfcpw->real_to_recip(this->porter, tmhpsi, this->ik, true); // wfcpw->convolution(this->ctx, // this->ik, // this->veff_col, @@ -89,8 +89,8 @@ void Veff>::act( { // T *porter1 = new T[wfcpw->nmaxgr]; // fft to real space and doing things. - wfcpw->recip_to_real(this->ctx, tmpsi_in, this->porter, this->ik); - wfcpw->recip_to_real(this->ctx, tmpsi_in + max_npw, this->porter1, this->ik); + wfcpw->recip_to_real(tmpsi_in, this->porter, this->ik); + wfcpw->recip_to_real(tmpsi_in + max_npw, this->porter1, this->ik); if(this->veff_col != 0) { /// denghui added at 20221109 @@ -114,8 +114,8 @@ void Veff>::act( // } } // (3) fft back to G space. - wfcpw->real_to_recip(this->ctx, this->porter, tmhpsi, this->ik, true); - wfcpw->real_to_recip(this->ctx, this->porter1, tmhpsi + max_npw, this->ik, true); + wfcpw->real_to_recip(this->porter, tmhpsi, this->ik, true); + wfcpw->real_to_recip(this->porter1, tmhpsi + max_npw, this->ik, true); } tmhpsi += max_npw * npol; tmpsi_in += max_npw * npol; diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_iter.cpp b/source/module_hamilt_pw/hamilt_stodft/sto_iter.cpp index 1134c7777be..e8eddca95fc 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_iter.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/sto_iter.cpp @@ -642,7 +642,7 @@ void Stochastic_Iter::cal_storho(const UnitCell& ucell, T* tmpout = stowf.shchi->get_pointer(); for (int ichi = 0; ichi < nchip_ik; ++ichi) { - wfc_basis->recip_to_real(this->ctx, tmpout, porter, ik); + wfc_basis->recip_to_real(tmpout, porter, ik); const auto w1 = static_cast(this->pkv->wk[ik]); elecstate::elecstate_pw_op()(this->ctx, current_spin, nrxx, w1, pes->rho, porter); // for (int ir = 0; ir < nrxx; ++ir) diff --git a/source/module_io/get_pchg_pw.h b/source/module_io/get_pchg_pw.h index 58f82f4a9c9..7fd1eff9ad2 100644 --- a/source/module_io/get_pchg_pw.h +++ b/source/module_io/get_pchg_pw.h @@ -96,7 +96,7 @@ void get_pchg_pw(const std::vector& bands_to_print, << ik % (nks / nspin) + 1 << ", spin " << spin_index + 1 << std::endl; psi->fix_k(ik); - pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfcr.data(), ik); + pw_wfc->recip_to_real,Device>(&psi[0](ib, 0), wfcr.data(), ik); // To ensure the normalization of charge density in multi-k calculation (if if_separate_k is true) double wg_sum_k = 0; @@ -139,7 +139,7 @@ void get_pchg_pw(const std::vector& bands_to_print, << ik % (nks / nspin) + 1 << ", spin " << spin_index + 1 << std::endl; psi->fix_k(ik); - pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfcr.data(), ik); + pw_wfc->recip_to_real,Device>( &psi[0](ib, 0), wfcr.data(), ik); double w1 = static_cast(wk[ik] / ucell->omega); From 3301b52570dde9888b2bc119c875c8fbc5da8e49 Mon Sep 17 00:00:00 2001 From: ubuntu <3158793232@qq.com> Date: Sun, 23 Mar 2025 20:06:22 +0800 Subject: [PATCH 7/8] moidfy back the func --- .../module_basis/module_pw/test/test-other.cpp | 8 ++++---- source/module_elecstate/elecstate_pw.cpp | 8 ++++---- .../module_elecstate/elecstate_pw_cal_tau.cpp | 4 ++-- .../module_xc/test/xc3_mock.h | 18 ++++++++++++------ .../module_xc/xc_functional_gradcorr.cpp | 3 ++- .../hamilt_ofdft/ml_data_descriptor.cpp | 4 ++-- .../hamilt_pwdft/operator_pw/meta_pw.cpp | 4 ++-- .../hamilt_pwdft/operator_pw/op_exx_pw.cpp | 10 +++++----- .../hamilt_stodft/sto_iter.cpp | 2 +- source/module_io/get_pchg_pw.h | 4 ++-- 10 files changed, 36 insertions(+), 29 deletions(-) diff --git a/source/module_basis/module_pw/test/test-other.cpp b/source/module_basis/module_pw/test/test-other.cpp index c5eb0093b27..308f4c6c689 100644 --- a/source/module_basis/module_pw/test/test-other.cpp +++ b/source/module_basis/module_pw/test/test-other.cpp @@ -76,26 +76,26 @@ TEST_F(PWTEST,test_other) } #endif - pwktest.recip_to_real,base_device::DEVICE_CPU>(rhog1, rhor1, ik); + pwktest.recip_to_real(ctx, rhog1, rhor1, ik); pwktest.recip2real(rhog2, rhor2, ik); for(int ir = 0 ; ir < nrxx; ++ir) { EXPECT_NEAR(std::abs(rhor1[ir]),std::abs(rhor2[ir]),1e-8); } - pwktest.real_to_recip,base_device::DEVICE_CPU>(rhor1, rhog1, ik); + pwktest.real_to_recip(ctx, rhor1, rhog1, ik); pwktest.real2recip(rhor2, rhog2, ik); for(int ig = 0 ; ig < npwk; ++ig) { EXPECT_NEAR(std::abs(rhog1[ig]),std::abs(rhog2[ig]),1e-8); } #ifdef __ENABLE_FLOAT_FFTW - pwktest.recip_to_real,base_device::DEVICE_CPU>(rhofg1, rhofr1, ik); + pwktest.recip_to_real(ctx, rhofg1, rhofr1, ik); pwktest.recip2real(rhofg2, rhofr2, ik); for(int ir = 0 ; ir < nrxx; ++ir) { EXPECT_NEAR(std::abs(rhofr1[ir]),std::abs(rhofr2[ir]),1e-6); } - pwktest.real_to_recip,base_device::DEVICE_CPU>(rhofr1, rhofg1, ik); + pwktest.real_to_recip(ctx, rhofr1, rhofg1, ik); pwktest.real2recip(rhofr2, rhofg2, ik); for(int ig = 0 ; ig < npwk; ++ig) { diff --git a/source/module_elecstate/elecstate_pw.cpp b/source/module_elecstate/elecstate_pw.cpp index 76a4ce6c42b..ef7a168752e 100644 --- a/source/module_elecstate/elecstate_pw.cpp +++ b/source/module_elecstate/elecstate_pw.cpp @@ -202,9 +202,9 @@ void ElecStatePW::rhoBandK(const psi::Psi& psi) /// be care of when smearing_sigma is large, wg would less than 0 /// - this->basis->recip_to_real( &psi(ibnd,0), this->wfcr, ik); + this->basis->recip_to_real(this->ctx, &psi(ibnd,0), this->wfcr, ik); - this->basis->recip_to_real( &psi(ibnd,npwx), this->wfcr_another_spin, ik); + this->basis->recip_to_real(this->ctx, &psi(ibnd,npwx), this->wfcr_another_spin, ik); const auto w1 = static_cast(this->wg(ik, ibnd) / ucell->omega); @@ -230,7 +230,7 @@ void ElecStatePW::rhoBandK(const psi::Psi& psi) /// only occupied band should be calculated. /// - this->basis->recip_to_real(&psi(ibnd,0), this->wfcr, ik); + this->basis->recip_to_real(this->ctx, &psi(ibnd,0), this->wfcr, ik); const auto w1 = static_cast(this->wg(ik, ibnd) / ucell->omega); @@ -258,7 +258,7 @@ void ElecStatePW::rhoBandK(const psi::Psi& psi) &psi(ibnd, 0), this->wfcr); - this->basis->recip_to_real( this->wfcr, this->wfcr, ik); + this->basis->recip_to_real(this->ctx, this->wfcr, this->wfcr, ik); elecstate_pw_op()(this->ctx, current_spin, this->charge->nrxx, w1, this->kin_r, this->wfcr); } diff --git a/source/module_elecstate/elecstate_pw_cal_tau.cpp b/source/module_elecstate/elecstate_pw_cal_tau.cpp index 1cf4f57b700..628dd25aeff 100644 --- a/source/module_elecstate/elecstate_pw_cal_tau.cpp +++ b/source/module_elecstate/elecstate_pw_cal_tau.cpp @@ -23,7 +23,7 @@ void ElecStatePW::cal_tau(const psi::Psi& psi) int nbands = psi.get_nbands(); for (int ibnd = 0; ibnd < nbands; ibnd++) { - this->basis->recip_to_real(&psi(ibnd,0), this->wfcr, ik); + this->basis->recip_to_real(this->ctx, &psi(ibnd,0), this->wfcr, ik); const auto w1 = static_cast(this->wg(ik, ibnd) / ucell->omega); @@ -43,7 +43,7 @@ void ElecStatePW::cal_tau(const psi::Psi& psi) &psi(ibnd, 0), this->wfcr); - this->basis->recip_to_real(this->wfcr, this->wfcr, ik); + this->basis->recip_to_real(this->ctx, this->wfcr, this->wfcr, ik); elecstate_pw_op()(this->ctx, current_spin, this->charge->nrxx, w1, this->kin_r, this->wfcr); } diff --git a/source/module_hamilt_general/module_xc/test/xc3_mock.h b/source/module_hamilt_general/module_xc/test/xc3_mock.h index ebbc139d75c..ffca76c70cb 100644 --- a/source/module_hamilt_general/module_xc/test/xc3_mock.h +++ b/source/module_hamilt_general/module_xc/test/xc3_mock.h @@ -78,7 +78,8 @@ namespace ModulePW template - void PW_Basis_K::real_to_recip(const std::complex* in, + void PW_Basis_K::real_to_recip(const Device* ctx, + const std::complex* in, std::complex* out, const int ik, const bool add, @@ -90,7 +91,8 @@ namespace ModulePW } } template - void PW_Basis_K::recip_to_real(const std::complex* in, + void PW_Basis_K::recip_to_real(const Device* ctx, + const std::complex* in, std::complex* out, const int ik, const bool add, @@ -102,24 +104,28 @@ namespace ModulePW } } - template void PW_Basis_K::real_to_recip(const std::complex* in, + template void PW_Basis_K::real_to_recip(const base_device::DEVICE_CPU* ctx, + const std::complex* in, std::complex* out, const int ik, const bool add, const double factor) const; - template void PW_Basis_K::recip_to_real(const std::complex* in, + template void PW_Basis_K::recip_to_real(const base_device::DEVICE_CPU* ctx, + const std::complex* in, std::complex* out, const int ik, const bool add, const double factor) const; #if __CUDA || __ROCM - template void PW_Basis_K::real_to_recip(const std::complex* in, + template void PW_Basis_K::real_to_recip(const base_device::DEVICE_GPU* ctx, + const std::complex* in, std::complex* out, const int ik, const bool add, const double factor) const; - template void PW_Basis_K::recip_to_real(const std::complex* in, + template void PW_Basis_K::recip_to_real(const base_device::DEVICE_GPU* ctx, + const std::complex* in, std::complex* out, const int ik, const bool add, diff --git a/source/module_hamilt_general/module_xc/xc_functional_gradcorr.cpp b/source/module_hamilt_general/module_xc/xc_functional_gradcorr.cpp index b3cad37c448..42501465108 100644 --- a/source/module_hamilt_general/module_xc/xc_functional_gradcorr.cpp +++ b/source/module_hamilt_general/module_xc/xc_functional_gradcorr.cpp @@ -644,7 +644,8 @@ void XC_Functional::grad_wfc( rhog, porter.data()); // Array of std::complex // bring the gdr from G --> R - wfc_basis->recip_to_real(porter.data(), porter.data(), ik); + Device * ctx = nullptr; + wfc_basis->recip_to_real(ctx, porter.data(), porter.data(), ik); xc_functional_grad_wfc_solver( ipol, wfc_basis->nrxx, // Integers diff --git a/source/module_hamilt_pw/hamilt_ofdft/ml_data_descriptor.cpp b/source/module_hamilt_pw/hamilt_ofdft/ml_data_descriptor.cpp index e9ed4e7268e..85135d789f7 100644 --- a/source/module_hamilt_pw/hamilt_ofdft/ml_data_descriptor.cpp +++ b/source/module_hamilt_pw/hamilt_ofdft/ml_data_descriptor.cpp @@ -197,7 +197,7 @@ void ML_data::getF_KS1( continue; } - pw_psi->recip_to_real( &psi->operator()(ibnd,0), wfcr, ik); + pw_psi->recip_to_real(ctx, &psi->operator()(ibnd,0), wfcr, ik); const double w1 = pelec->wg(ik, ibnd) / ucell.omega; // output one wf, to check KS equation @@ -308,7 +308,7 @@ void ML_data::getF_KS2( continue; } - pw_psi->recip_to_real( &psi->operator()(ibnd,0), wfcr, ik); + pw_psi->recip_to_real(ctx, &psi->operator()(ibnd,0), wfcr, ik); const double w1 = pelec->wg(ik, ibnd) / ucell.omega; if (pelec->ekb(ik,ibnd) > epsilonM) diff --git a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/meta_pw.cpp b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/meta_pw.cpp index 538819f4522..817dba61cea 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/meta_pw.cpp +++ b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/meta_pw.cpp @@ -67,13 +67,13 @@ void Meta>::act( for (int j = 0; j < 3; j++) { meta_op()(this->ctx, this->ik, j, ngk_ik, this->wfcpw->npwk_max, this->tpiba, wfcpw->get_gcar_data(), wfcpw->get_kvec_c_data(), tmpsi_in, this->porter); - wfcpw->recip_to_real(this->porter, this->porter, this->ik); + wfcpw->recip_to_real(this->ctx, this->porter, this->porter, this->ik); if(this->vk_col != 0) { vector_mul_vector_op()(this->vk_col, this->porter, this->porter, this->vk + current_spin * this->vk_col); } - wfcpw->real_to_recip(this->porter, this->porter, this->ik); + wfcpw->real_to_recip(this->ctx, this->porter, this->porter, this->ik); meta_op()(this->ctx, this->ik, j, ngk_ik, this->wfcpw->npwk_max, this->tpiba, wfcpw->get_gcar_data(), wfcpw->get_kvec_c_data(), this->porter, tmhpsi, true); } // x,y,z directions diff --git a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.cpp b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.cpp index 202800026c7..4b6d36908a0 100644 --- a/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.cpp +++ b/source/module_hamilt_pw/hamilt_pwdft/operator_pw/op_exx_pw.cpp @@ -188,7 +188,7 @@ void OperatorEXXPW::act_op(const int nbands, { const T *psi_nk = tmpsi_in + n_iband * nbasis; // retrieve \psi_nk in real space - wfcpw->recip_to_real( psi_nk, psi_nk_real, this->ik); + wfcpw->recip_to_real(ctx, psi_nk, psi_nk_real, this->ik); // for \psi_nk, get the pw of iq and band m auto q_points = get_q_points(this->ik); @@ -208,7 +208,7 @@ void OperatorEXXPW::act_op(const int nbands, // if (has_real.find({iq, m_iband}) == has_real.end()) // { const T* psi_mq = get_pw(m_iband, iq); - wfcpw->recip_to_real( psi_mq, psi_mq_real, iq); + wfcpw->recip_to_real(ctx, psi_mq, psi_mq_real, iq); // syncmem_complex_op()(this->ctx, this->ctx, psi_all_real + m_iband * wfcpw->nrxx, psi_mq_real, wfcpw->nrxx); // has_real[{iq, m_iband}] = true; // } @@ -271,7 +271,7 @@ void OperatorEXXPW::act_op(const int nbands, } // end of iq auto h_psi_nk = tmhpsi + n_iband * nbasis; Real hybrid_alpha = GlobalC::exx_info.info_global.hybrid_alpha; - wfcpw->real_to_recip(h_psi_real, h_psi_nk, this->ik, true, hybrid_alpha); + wfcpw->real_to_recip(ctx, h_psi_real, h_psi_nk, this->ik, true, hybrid_alpha); setmem_complex_op()(h_psi_real, 0, rhopw->nrxx); } @@ -810,7 +810,7 @@ double OperatorEXXPW::cal_exx_energy_op(psi::Psi *ppsi_) c psi.fix_kb(ik, n_iband); const T* psi_nk = psi.get_pointer(); // retrieve \psi_nk in real space - wfcpw->recip_to_real( psi_nk, psi_nk_real, ik); + wfcpw->recip_to_real(ctx, psi_nk, psi_nk_real, ik); // for \psi_nk, get the pw of iq and band m // q_points is a vector of integers, 0 to nks-1 @@ -839,7 +839,7 @@ double OperatorEXXPW::cal_exx_energy_op(psi::Psi *ppsi_) c psi_.fix_kb(iq, m_iband); const T* psi_mq = psi_.get_pointer(); // const T* psi_mq = get_pw(m_iband, iq); - wfcpw->recip_to_real(psi_mq, psi_mq_real, iq); + wfcpw->recip_to_real(ctx, psi_mq, psi_mq_real, iq); T omega_inv = 1.0 / ucell->omega; diff --git a/source/module_hamilt_pw/hamilt_stodft/sto_iter.cpp b/source/module_hamilt_pw/hamilt_stodft/sto_iter.cpp index e8eddca95fc..1134c7777be 100644 --- a/source/module_hamilt_pw/hamilt_stodft/sto_iter.cpp +++ b/source/module_hamilt_pw/hamilt_stodft/sto_iter.cpp @@ -642,7 +642,7 @@ void Stochastic_Iter::cal_storho(const UnitCell& ucell, T* tmpout = stowf.shchi->get_pointer(); for (int ichi = 0; ichi < nchip_ik; ++ichi) { - wfc_basis->recip_to_real(tmpout, porter, ik); + wfc_basis->recip_to_real(this->ctx, tmpout, porter, ik); const auto w1 = static_cast(this->pkv->wk[ik]); elecstate::elecstate_pw_op()(this->ctx, current_spin, nrxx, w1, pes->rho, porter); // for (int ir = 0; ir < nrxx; ++ir) diff --git a/source/module_io/get_pchg_pw.h b/source/module_io/get_pchg_pw.h index 7fd1eff9ad2..58f82f4a9c9 100644 --- a/source/module_io/get_pchg_pw.h +++ b/source/module_io/get_pchg_pw.h @@ -96,7 +96,7 @@ void get_pchg_pw(const std::vector& bands_to_print, << ik % (nks / nspin) + 1 << ", spin " << spin_index + 1 << std::endl; psi->fix_k(ik); - pw_wfc->recip_to_real,Device>(&psi[0](ib, 0), wfcr.data(), ik); + pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfcr.data(), ik); // To ensure the normalization of charge density in multi-k calculation (if if_separate_k is true) double wg_sum_k = 0; @@ -139,7 +139,7 @@ void get_pchg_pw(const std::vector& bands_to_print, << ik % (nks / nspin) + 1 << ", spin " << spin_index + 1 << std::endl; psi->fix_k(ik); - pw_wfc->recip_to_real,Device>( &psi[0](ib, 0), wfcr.data(), ik); + pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfcr.data(), ik); double w1 = static_cast(wk[ik] / ucell->omega); From 283c6745c162fc98c1f9af27131cf092350f0b4f Mon Sep 17 00:00:00 2001 From: ubuntu <3158793232@qq.com> Date: Mon, 24 Mar 2025 14:53:22 +0800 Subject: [PATCH 8/8] update func --- source/module_basis/module_pw/pw_basis.h | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/source/module_basis/module_pw/pw_basis.h b/source/module_basis/module_pw/pw_basis.h index 52a59d2fb32..8bc45518a2b 100644 --- a/source/module_basis/module_pw/pw_basis.h +++ b/source/module_basis/module_pw/pw_basis.h @@ -352,10 +352,7 @@ class PW_Basis void recip_to_real(TK* in, TR* out, const bool add = false, - const typename GetTypeReal::type factor = 1.0) const - { - this->recip2real_gpu(in, out, add, factor); - }; + const typename GetTypeReal::type factor = 1.0) const; // template