From 2fb910542eddc4b6bc92e2d5c2e47fa7e2e104ff Mon Sep 17 00:00:00 2001 From: AsTonyshment Date: Thu, 17 Jul 2025 21:02:20 +0800 Subject: [PATCH 1/2] Fix GPU output of out_pchg and out_wfc_norm, out_wfc_re_im --- source/source_esolver/esolver_ks_pw.cpp | 98 ++++++++++++++----------- source/source_io/get_pchg_pw.h | 44 +++++++++-- source/source_io/get_wf_pw.h | 71 ++++++++++++------ 3 files changed, 144 insertions(+), 69 deletions(-) diff --git a/source/source_esolver/esolver_ks_pw.cpp b/source/source_esolver/esolver_ks_pw.cpp index acc752d0cfe..2d3e3f0164d 100644 --- a/source/source_esolver/esolver_ks_pw.cpp +++ b/source/source_esolver/esolver_ks_pw.cpp @@ -11,13 +11,6 @@ #include "source_estate/module_charge/symmetry_rho.h" #include "source_hamilt/module_ewald/H_Ewald_pw.h" #include "source_hamilt/module_vdw/vdw.h" -#include "source_lcao/module_deltaspin/spin_constrain.h" -#include "source_lcao/module_dftu/dftu.h" -#include "source_pw/module_pwdft/elecond.h" -#include "source_pw/module_pwdft/forces.h" -#include "source_pw/module_pwdft/hamilt_pw.h" -#include "source_pw/module_pwdft/onsite_projector.h" -#include "source_pw/module_pwdft/stress_pw.h" #include "source_hsolver/diago_iter_assist.h" #include "source_hsolver/hsolver_pw.h" #include "source_hsolver/kernels/dngvd_op.h" @@ -25,13 +18,20 @@ #include "source_io/cal_ldos.h" #include "source_io/get_pchg_pw.h" #include "source_io/get_wf_pw.h" +#include "source_io/module_parameter/parameter.h" #include "source_io/numerical_basis.h" #include "source_io/numerical_descriptor.h" #include "source_io/to_wannier90_pw.h" #include "source_io/winput.h" #include "source_io/write_dos_pw.h" #include "source_io/write_wfc_pw.h" -#include "source_io/module_parameter/parameter.h" +#include "source_lcao/module_deltaspin/spin_constrain.h" +#include "source_lcao/module_dftu/dftu.h" +#include "source_pw/module_pwdft/elecond.h" +#include "source_pw/module_pwdft/forces.h" +#include "source_pw/module_pwdft/hamilt_pw.h" +#include "source_pw/module_pwdft/onsite_projector.h" +#include "source_pw/module_pwdft/stress_pw.h" #include @@ -717,16 +717,25 @@ void ESolver_KS_PW::after_scf(UnitCell& ucell, const int istep, const //------------------------------------------------------------------ // 4) calculate band-decomposed (partial) charge density in pw basis //------------------------------------------------------------------ - const std::vector out_pchg = PARAM.inp.out_pchg; - if (out_pchg.size() > 0) + if (PARAM.inp.out_pchg.size() > 0) { - ModuleIO::get_pchg_pw(out_pchg, + if (this->__kspw_psi != nullptr && PARAM.inp.precision == "single") + { + delete reinterpret_cast, Device>*>(this->__kspw_psi); + } + + // Refresh __kspw_psi + this->__kspw_psi = PARAM.inp.precision == "single" + ? new psi::Psi, Device>(this->kspw_psi[0]) + : reinterpret_cast, Device>*>(this->kspw_psi); + + ModuleIO::get_pchg_pw(PARAM.inp.out_pchg, this->kspw_psi->get_nbands(), PARAM.inp.nspin, this->pw_rhod->nxyz, this->chr.ngmc, &ucell, - this->psi, + this->__kspw_psi, this->pw_rhod, this->pw_wfc, this->ctx, @@ -943,20 +952,25 @@ void ESolver_KS_PW::after_all_runners(UnitCell& ucell) //---------------------------------------------------------- //! 5) Print out electronic wave functions in real space //---------------------------------------------------------- - const std::vector out_wfc_norm = PARAM.inp.out_wfc_norm; - const std::vector out_wfc_re_im = PARAM.inp.out_wfc_re_im; - if (out_wfc_norm.size() > 0 || out_wfc_re_im.size() > 0) + if (PARAM.inp.out_wfc_norm.size() > 0 || PARAM.inp.out_wfc_re_im.size() > 0) { - ModuleIO::get_wf_pw(out_wfc_norm, - out_wfc_re_im, + if (this->__kspw_psi != nullptr && PARAM.inp.precision == "single") + { + delete reinterpret_cast, Device>*>(this->__kspw_psi); + } + + // Refresh __kspw_psi + this->__kspw_psi = PARAM.inp.precision == "single" + ? new psi::Psi, Device>(this->kspw_psi[0]) + : reinterpret_cast, Device>*>(this->kspw_psi); + + ModuleIO::get_wf_pw(PARAM.inp.out_wfc_norm, + PARAM.inp.out_wfc_re_im, this->kspw_psi->get_nbands(), PARAM.inp.nspin, - this->pw_rhod->nx, - this->pw_rhod->ny, - this->pw_rhod->nz, this->pw_rhod->nxyz, &ucell, - this->psi, + this->__kspw_psi, this->pw_wfc, this->ctx, this->Pgrid, @@ -991,29 +1005,29 @@ void ESolver_KS_PW::after_all_runners(UnitCell& ucell) ModuleIO::Write_MLKEDF_Descriptors write_mlkedf_desc; write_mlkedf_desc.cal_tool->set_para(this->chr.nrxx, - PARAM.inp.nelec, - PARAM.inp.of_tf_weight, - PARAM.inp.of_vw_weight, - PARAM.inp.of_ml_chi_p, - PARAM.inp.of_ml_chi_q, - PARAM.inp.of_ml_chi_xi, - PARAM.inp.of_ml_chi_pnl, - PARAM.inp.of_ml_chi_qnl, - PARAM.inp.of_ml_nkernel, - PARAM.inp.of_ml_kernel, - PARAM.inp.of_ml_kernel_scaling, - PARAM.inp.of_ml_yukawa_alpha, - PARAM.inp.of_ml_kernel_file, - ucell.omega, - this->pw_rho); + PARAM.inp.nelec, + PARAM.inp.of_tf_weight, + PARAM.inp.of_vw_weight, + PARAM.inp.of_ml_chi_p, + PARAM.inp.of_ml_chi_q, + PARAM.inp.of_ml_chi_xi, + PARAM.inp.of_ml_chi_pnl, + PARAM.inp.of_ml_chi_qnl, + PARAM.inp.of_ml_nkernel, + PARAM.inp.of_ml_kernel, + PARAM.inp.of_ml_kernel_scaling, + PARAM.inp.of_ml_yukawa_alpha, + PARAM.inp.of_ml_kernel_file, + ucell.omega, + this->pw_rho); write_mlkedf_desc.generateTrainData_KS(PARAM.globalv.global_mlkedf_descriptor_dir, - this->kspw_psi, - this->pelec, - this->pw_wfc, - this->pw_rho, - ucell, - this->pelec->pot->get_effective_v(0)); + this->kspw_psi, + this->pelec, + this->pw_wfc, + this->pw_rho, + ucell, + this->pelec->pot->get_effective_v(0)); } #endif } diff --git a/source/source_io/get_pchg_pw.h b/source/source_io/get_pchg_pw.h index 845a32b0bd5..7c2d14fe6ac 100644 --- a/source/source_io/get_pchg_pw.h +++ b/source/source_io/get_pchg_pw.h @@ -12,7 +12,7 @@ void get_pchg_pw(const std::vector& out_pchg, const int nxyz, const int chr_ngmc, UnitCell* ucell, - const psi::Psi>* psi, + const psi::Psi, Device>* kspw_psi, const ModulePW::PW_Basis* pw_rhod, const ModulePW::PW_Basis_K* pw_wfc, const Device* ctx, @@ -67,9 +67,17 @@ void get_pchg_pw(const std::vector& out_pchg, bands_picked[i] = static_cast(out_pchg[i]); } + // Allocate host memory std::vector> wfcr(nxyz); std::vector> rho_band(nspin, std::vector(nxyz)); + // Allocate device memory + std::complex* wfcr_device = nullptr; + if (!std::is_same::value) + { + base_device::memory::resize_memory_op, Device>()(wfcr_device, nxyz); + } + for (int ib = 0; ib < nbands; ++ib) { // Skip the loop iteration if bands_picked[ib] is 0 @@ -91,8 +99,21 @@ void get_pchg_pw(const std::vector& out_pchg, const int spin_index = kv.isk[ik]; // spin index const int k_number = ikstot % (nkstot / nspin) + 1; // k-point number, starting from 1 - psi->fix_k(ik); - pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfcr.data(), ik); + kspw_psi->fix_k(ik); + + // FFT on device and copy result back to host + if (std::is_same::value) + { + pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr.data(), ik); + } + else + { + pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr_device, ik); + + base_device::memory::synchronize_memory_op, + base_device::DEVICE_CPU, + Device>()(wfcr.data(), wfcr_device, nxyz); + } // To ensure the normalization of charge density in multi-k calculation (if if_separate_k is true) double wg_sum_k = 0.0; @@ -142,8 +163,21 @@ void get_pchg_pw(const std::vector& out_pchg, const int spin_index = kv.isk[ik]; // spin index const int k_number = ikstot % (nkstot / nspin) + 1; // k-point number, starting from 1 - psi->fix_k(ik); - pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfcr.data(), ik); + kspw_psi->fix_k(ik); + + // FFT on device and copy result back to host + if (std::is_same::value) + { + pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr.data(), ik); + } + else + { + pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr_device, ik); + + base_device::memory::synchronize_memory_op, + base_device::DEVICE_CPU, + Device>()(wfcr.data(), wfcr_device, nxyz); + } double w1 = static_cast(kv.wk[ik] / ucell->omega); diff --git a/source/source_io/get_wf_pw.h b/source/source_io/get_wf_pw.h index de16c8def98..9262ad181b6 100644 --- a/source/source_io/get_wf_pw.h +++ b/source/source_io/get_wf_pw.h @@ -1,20 +1,6 @@ #ifndef GET_WF_PW_H #define GET_WF_PW_H -#include "cube_io.h" -#include "source_base/module_device/device.h" -#include "source_base/tool_quit.h" -#include "source_basis/module_pw/pw_basis.h" -#include "source_basis/module_pw/pw_basis_k.h" -#include "source_cell/unitcell.h" -#include "source_estate/elecstate.h" -#include "source_estate/module_charge/symmetry_rho.h" -#include "source_pw/module_pwdft/parallel_grid.h" -#include "source_psi/psi.h" - -#include -#include - namespace ModuleIO { template @@ -22,12 +8,9 @@ void get_wf_pw(const std::vector& out_wfc_norm, const std::vector& out_wfc_re_im, const int nbands, const int nspin, - const int nx, - const int ny, - const int nz, const int nxyz, UnitCell* ucell, - const psi::Psi>* psi, + const psi::Psi, Device>* kspw_psi, const ModulePW::PW_Basis_K* pw_wfc, const Device* ctx, const Parallel_Grid& pgrid, @@ -94,9 +77,17 @@ void get_wf_pw(const std::vector& out_wfc_norm, bands_picked_re_im[i] = static_cast(out_wfc_re_im[i]); } + // Allocate host memory std::vector> wfcr_norm(nxyz); std::vector> rho_band_norm(nspin, std::vector(nxyz)); + // Allocate device memory + std::complex* wfcr_norm_device = nullptr; + if (!std::is_same::value) + { + base_device::memory::resize_memory_op, Device>()(wfcr_norm_device, nxyz); + } + for (int ib = 0; ib < nbands; ++ib) { // Skip the loop iteration if bands_picked[ib] is 0 @@ -115,8 +106,22 @@ void get_wf_pw(const std::vector& out_wfc_norm, const int spin_index = kv.isk[ik]; // spin index const int k_number = ikstot % (nkstot / nspin) + 1; // k-point number, starting from 1 - psi->fix_k(ik); - pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfcr_norm.data(), ik); + kspw_psi->fix_k(ik); + + // FFT on device and copy result back to host + if (std::is_same::value) + { + pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr_norm.data(), ik); + } + else + { + pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfcr_norm_device, ik); + + base_device::memory::synchronize_memory_op, base_device::DEVICE_CPU, Device>()( + wfcr_norm.data(), + wfcr_norm_device, + nxyz); + } // To ensure the normalization of charge density in multi-k calculation double wg_sum_k = 0.0; @@ -159,10 +164,18 @@ void get_wf_pw(const std::vector& out_wfc_norm, } } + // Allocate host memory std::vector> wfc_re_im(nxyz); std::vector> rho_band_re(nspin, std::vector(nxyz)); std::vector> rho_band_im(nspin, std::vector(nxyz)); + // Allocate device memory + std::complex* wfc_re_im_device = nullptr; + if (!std::is_same::value) + { + base_device::memory::resize_memory_op, Device>()(wfc_re_im_device, nxyz); + } + for (int ib = 0; ib < nbands; ++ib) { // Skip the loop iteration if bands_picked[ib] is 0 @@ -182,8 +195,22 @@ void get_wf_pw(const std::vector& out_wfc_norm, const int spin_index = kv.isk[ik]; // spin index const int k_number = ikstot % (nkstot / nspin) + 1; // k-point number, starting from 1 - psi->fix_k(ik); - pw_wfc->recip_to_real(ctx, &psi[0](ib, 0), wfc_re_im.data(), ik); + kspw_psi->fix_k(ik); + + // FFT on device and copy result back to host + if (std::is_same::value) + { + pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfc_re_im.data(), ik); + } + else + { + pw_wfc->recip_to_real(ctx, &kspw_psi[0](ib, 0), wfc_re_im_device, ik); + + base_device::memory::synchronize_memory_op, base_device::DEVICE_CPU, Device>()( + wfc_re_im.data(), + wfc_re_im_device, + nxyz); + } // To ensure the normalization of charge density in multi-k calculation double wg_sum_k = 0.0; From d2b98693a86773ef782ff8bd8a09d1963380dc0c Mon Sep 17 00:00:00 2001 From: AsTonyshment Date: Thu, 17 Jul 2025 21:21:16 +0800 Subject: [PATCH 2/2] GPU integrate test is functional again --- tests/11_PW_GPU/006_PW_get_wf_GPU/INPUT | 22 +++++++++++++++++++ .../KPT | 0 tests/11_PW_GPU/006_PW_get_wf_GPU/README | 1 + .../STRU | 0 tests/11_PW_GPU/006_PW_get_wf_GPU/result.ref | 11 ++++++++++ tests/11_PW_GPU/BUG_PW_outwfcr_GPU/INPUT | 22 ------------------- tests/11_PW_GPU/BUG_PW_outwfcr_GPU/README | 1 - tests/11_PW_GPU/BUG_PW_outwfcr_GPU/result.ref | 12 ---------- tests/11_PW_GPU/CASES_GPU.txt | 1 + 9 files changed, 35 insertions(+), 35 deletions(-) create mode 100644 tests/11_PW_GPU/006_PW_get_wf_GPU/INPUT rename tests/11_PW_GPU/{BUG_PW_outwfcr_GPU => 006_PW_get_wf_GPU}/KPT (100%) create mode 100644 tests/11_PW_GPU/006_PW_get_wf_GPU/README rename tests/11_PW_GPU/{BUG_PW_outwfcr_GPU => 006_PW_get_wf_GPU}/STRU (100%) create mode 100644 tests/11_PW_GPU/006_PW_get_wf_GPU/result.ref delete mode 100644 tests/11_PW_GPU/BUG_PW_outwfcr_GPU/INPUT delete mode 100644 tests/11_PW_GPU/BUG_PW_outwfcr_GPU/README delete mode 100644 tests/11_PW_GPU/BUG_PW_outwfcr_GPU/result.ref diff --git a/tests/11_PW_GPU/006_PW_get_wf_GPU/INPUT b/tests/11_PW_GPU/006_PW_get_wf_GPU/INPUT new file mode 100644 index 00000000000..c6775e868bc --- /dev/null +++ b/tests/11_PW_GPU/006_PW_get_wf_GPU/INPUT @@ -0,0 +1,22 @@ +INPUT_PARAMETERS +#Parameters (1.General) +suffix autotest +calculation scf + +nbands 6 +symmetry 0 +pseudo_dir ../../PP_ORB + +device gpu + +#Parameters (2.Iteration) +ecutwfc 100 +scf_thr 1e-9 +scf_nmax 100 + +#Parameters (3.Basis) +basis_type pw + +out_wfc_norm 1*1 + +pw_seed 1 diff --git a/tests/11_PW_GPU/BUG_PW_outwfcr_GPU/KPT b/tests/11_PW_GPU/006_PW_get_wf_GPU/KPT similarity index 100% rename from tests/11_PW_GPU/BUG_PW_outwfcr_GPU/KPT rename to tests/11_PW_GPU/006_PW_get_wf_GPU/KPT diff --git a/tests/11_PW_GPU/006_PW_get_wf_GPU/README b/tests/11_PW_GPU/006_PW_get_wf_GPU/README new file mode 100644 index 00000000000..2fbda197696 --- /dev/null +++ b/tests/11_PW_GPU/006_PW_get_wf_GPU/README @@ -0,0 +1 @@ +GPU version / test the output of out_wfc_norm diff --git a/tests/11_PW_GPU/BUG_PW_outwfcr_GPU/STRU b/tests/11_PW_GPU/006_PW_get_wf_GPU/STRU similarity index 100% rename from tests/11_PW_GPU/BUG_PW_outwfcr_GPU/STRU rename to tests/11_PW_GPU/006_PW_get_wf_GPU/STRU diff --git a/tests/11_PW_GPU/006_PW_get_wf_GPU/result.ref b/tests/11_PW_GPU/006_PW_get_wf_GPU/result.ref new file mode 100644 index 00000000000..5a48b2a377a --- /dev/null +++ b/tests/11_PW_GPU/006_PW_get_wf_GPU/result.ref @@ -0,0 +1,11 @@ +etotref -211.878925376034 +etotperatomref -105.9394626880 +wfi1s1k1.cube 21.94721795 +wfi1s1k2.cube 19.65847203 +wfi1s1k3.cube 19.65846259 +wfi1s1k4.cube 19.16322464 +wfi1s1k5.cube 19.65846987 +wfi1s1k6.cube 19.17198153 +wfi1s1k7.cube 18.92365819 +wfi1s1k8.cube 19.6584683 +totaltimeref 3.26 \ No newline at end of file diff --git a/tests/11_PW_GPU/BUG_PW_outwfcr_GPU/INPUT b/tests/11_PW_GPU/BUG_PW_outwfcr_GPU/INPUT deleted file mode 100644 index 653bb7b7c91..00000000000 --- a/tests/11_PW_GPU/BUG_PW_outwfcr_GPU/INPUT +++ /dev/null @@ -1,22 +0,0 @@ -INPUT_PARAMETERS -#Parameters (1.General) -suffix autotest -calculation scf - -nbands 6 -symmetry 1 -pseudo_dir ../../PP_ORB - -device gpu - -#Parameters (2.Iteration) -ecutwfc 20 -scf_thr 1e-9 -scf_nmax 100 - -#Parameters (3.Basis) -basis_type pw - -out_wfc_r 1 - -pw_seed 1 diff --git a/tests/11_PW_GPU/BUG_PW_outwfcr_GPU/README b/tests/11_PW_GPU/BUG_PW_outwfcr_GPU/README deleted file mode 100644 index c4bd30f7fdf..00000000000 --- a/tests/11_PW_GPU/BUG_PW_outwfcr_GPU/README +++ /dev/null @@ -1 +0,0 @@ -GPU version / test the output of wfc_r, and compare the variance diff --git a/tests/11_PW_GPU/BUG_PW_outwfcr_GPU/result.ref b/tests/11_PW_GPU/BUG_PW_outwfcr_GPU/result.ref deleted file mode 100644 index 1880e707f78..00000000000 --- a/tests/11_PW_GPU/BUG_PW_outwfcr_GPU/result.ref +++ /dev/null @@ -1,12 +0,0 @@ -etotref -197.1405644417868 -etotperatomref -98.5702822209 -variance_wfc_r_0_0 0.31340 -variance_wfc_r_0_1 1.71055 -variance_wfc_r_0_2 2.39603 -variance_wfc_r_0_3 1.66607 -variance_wfc_r_0_4 1.05190 -variance_wfc_r_0_5 1.29386 -pointgroupref T_d -spacegroupref O_h -nksibzref 1 -totaltimeref 0.31 diff --git a/tests/11_PW_GPU/CASES_GPU.txt b/tests/11_PW_GPU/CASES_GPU.txt index e1dc4701b26..bb0cf0b776d 100644 --- a/tests/11_PW_GPU/CASES_GPU.txt +++ b/tests/11_PW_GPU/CASES_GPU.txt @@ -2,3 +2,4 @@ 002_PW_CG_GPU 003_PW_DA_GPU 005_PW_CG_GPU_float +006_PW_get_wf_GPU