From b1bb7592d5eb05180b2072619bc1fad5a77c25df Mon Sep 17 00:00:00 2001 From: chengleizheng Date: Thu, 16 Jul 2026 12:49:15 +0800 Subject: [PATCH 1/5] Fix GPU pointer accesses from CPU code in USPP PW calculation --- source/source_estate/elecstate_pw.cpp | 264 +++++++++++--------- source/source_estate/elecstate_pw.h | 11 + source/source_pw/module_pwdft/hamilt_pw.cpp | 14 +- source/source_pw/module_pwdft/hamilt_pw.h | 1 + 4 files changed, 170 insertions(+), 120 deletions(-) diff --git a/source/source_estate/elecstate_pw.cpp b/source/source_estate/elecstate_pw.cpp index de05d441b58..4c76bfec02d 100644 --- a/source/source_estate/elecstate_pw.cpp +++ b/source/source_estate/elecstate_pw.cpp @@ -50,7 +50,7 @@ ElecStatePW::~ElecStatePW() } if (PARAM.globalv.use_uspp) { - delmem_var_op()(this->becsum); + delmem_var_h_op()(this->becsum); } delmem_complex_op()(this->wfcr); delmem_complex_op()(this->wfcr_another_spin); @@ -277,11 +277,19 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) const int nbands = psi.get_nbands() * npol; const int nkb = this->ppcell->nkb; this->vkb = this->ppcell->template get_vkb_data(); - T* becp = nullptr; - resmem_complex_op()(becp, nbands * nkb, "ElecState::becp"); const int nh_tot = this->ppcell->nhm * (this->ppcell->nhm + 1) / 2; - resmem_var_op()(becsum, nh_tot * ucell->nat * PARAM.inp.nspin, "ElecState::becsum"); - setmem_var_op()(becsum, 0, nh_tot * ucell->nat * PARAM.inp.nspin); + // becsum on CPU (forces_us / stress_us use CPU dgemm) + resmem_var_h_op()(becsum, nh_tot * ucell->nat * PARAM.inp.nspin, "ElecState::becsum"); + setmem_var_h_op()(becsum, 0, nh_tot * ucell->nat * PARAM.inp.nspin); + + // becp: GPU for first gemm, then D2H for CPU loops + T* becp_gpu = nullptr; + std::vector becp; + if (nkb > 0) + { + resmem_complex_op()(becp_gpu, nbands * nkb, "ElecState::becp"); + becp.resize(nbands * nkb); + } for (int ik = 0; ik < psi.get_nk(); ++ik) { @@ -296,41 +304,46 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) this->ppcell->getvnl(this->ctx, *ucell,ik, this->vkb); } - // becp = + // becp = (GPU gemm) char transa = 'C'; char transb = 'N'; - if (nbands == 1) - { - int inc = 1; - gemv_op()(transa, - npw, - this->ppcell->nkb, - &one, - this->vkb, - this->ppcell->vkb.nc, - psi_now, - inc, - &zero, - becp, - inc); - } - else + if (this->ppcell->nkb > 0) { - gemm_op()(transa, - transb, - this->ppcell->nkb, - nbands, - npw, - &one, - this->vkb, - this->ppcell->vkb.nc, - psi_now, - npwx, - &zero, - becp, - this->ppcell->nkb); + if (nbands == 1) + { + int inc = 1; + gemv_op()(transa, + npw, + this->ppcell->nkb, + &one, + this->vkb, + this->ppcell->vkbnc, + psi_now, + inc, + &zero, + becp_gpu, + inc); + } + else + { + gemm_op()(transa, + transb, + this->ppcell->nkb, + nbands, + npw, + &one, + this->vkb, + this->ppcell->vkbnc, + psi_now, + npwx, + &zero, + becp_gpu, + this->ppcell->nkb); + } + // D2H: GPU becp → CPU becp + syncmem_complex_d2h_op()(becp.data(), becp_gpu, nbands * nkb); } - Parallel_Reduce::reduce_pool(becp, this->ppcell->nkb * nbands); + Parallel_Reduce::reduce_pool(becp.data(), this->ppcell->nkb * nbands); // sum over bands: \sum_i w_i for (int it = 0; it < ucell->ntype; it++) @@ -338,12 +351,10 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) Atom* atom = &ucell->atoms[it]; if (atom->ncpp.tvanp) { - T *auxk1 = nullptr, *auxk2 = nullptr, *aux_gk = nullptr; - resmem_complex_op()(auxk1, nbands * atom->ncpp.nh, "ElecState::auxk1"); - resmem_complex_op()(auxk2, nbands * atom->ncpp.nh, "ElecState::auxk2"); - resmem_complex_op()(aux_gk, - atom->ncpp.nh * atom->ncpp.nh * npol * npol, - "ElecState::aux_gk"); + const int nh_atom = atom->ncpp.nh; + // auxk1, auxk2 on CPU for filling loops + std::vector auxk1(nbands * nh_atom); + std::vector auxk2(nbands * nh_atom); for (int ia = 0; ia < atom->na; ia++) { const int iat = ucell->itia2iat(it, ia); @@ -353,7 +364,7 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) } else { - for (int ih = 0; ih < atom->ncpp.nh; ih++) + for (int ih = 0; ih < nh_atom; ih++) { const int ikb = this->ppcell->indv_ijkb0[iat] + ih; for (int ib = 0; ib < nbands; ib++) @@ -364,60 +375,62 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) } } - char transa = 'C'; - char transb = 'N'; - gemm_op()(transa, - transb, - atom->ncpp.nh, - atom->ncpp.nh, + // GPU gemm: aux_gk = auxk1^H * auxk2 + T *aux_gk_gpu = nullptr, *auxk1_gpu = nullptr, *auxk2_gpu = nullptr; + resmem_complex_op()(auxk1_gpu, nbands * nh_atom, "ElecState::auxk1"); + resmem_complex_op()(auxk2_gpu, nbands * nh_atom, "ElecState::auxk2"); + resmem_complex_op()(aux_gk_gpu, nh_atom * nh_atom * npol * npol, "ElecState::aux_gk"); + syncmem_complex_h2d_op()(auxk1_gpu, auxk1.data(), nbands * nh_atom); + syncmem_complex_h2d_op()(auxk2_gpu, auxk2.data(), nbands * nh_atom); + + char transa2 = 'C'; + char transb2 = 'N'; + gemm_op()(transa2, + transb2, + nh_atom, + nh_atom, nbands, &one, - auxk1, + auxk1_gpu, nbands, - auxk2, + auxk2_gpu, nbands, &zero, - aux_gk, - atom->ncpp.nh); - } + aux_gk_gpu, + nh_atom); - // copy output from GEMM into desired format - if (PARAM.inp.noncolin && !atom->ncpp.has_so) - { - } - else if (PARAM.inp.noncolin && atom->ncpp.has_so) - { - } - else - { + // D2H: GPU aux_gk → CPU + std::vector aux_gk(nh_atom * nh_atom); + syncmem_complex_d2h_op()(aux_gk.data(), aux_gk_gpu, nh_atom * nh_atom); + + delmem_complex_op()(auxk1_gpu); + delmem_complex_op()(auxk2_gpu); + delmem_complex_op()(aux_gk_gpu); + + // copy output from GEMM into desired format int ijh = 0; const int index = currect_spin * ucell->nat * nh_tot + iat * nh_tot; - for (int ih = 0; ih < atom->ncpp.nh; ih++) + for (int ih = 0; ih < nh_atom; ih++) { - for (int jh = ih; jh < atom->ncpp.nh; jh++) + for (int jh = ih; jh < nh_atom; jh++) { - // nondiagonal terms summed and collapsed into a - // single index (matrix is symmetric wrt (ih,jh)) if (ih == jh) { - becsum[index + ijh] += std::real(aux_gk[ih * atom->ncpp.nh + jh]); + becsum[index + ijh] += std::real(aux_gk[ih * nh_atom + jh]); } else { - becsum[index + ijh] += 2.0 * std::real(aux_gk[ih * atom->ncpp.nh + jh]); + becsum[index + ijh] += 2.0 * std::real(aux_gk[ih * nh_atom + jh]); } ijh++; } } } } - delmem_complex_op()(auxk1); - delmem_complex_op()(auxk2); - delmem_complex_op()(aux_gk); } } } - delmem_complex_op()(becp); + if (becp_gpu != nullptr) delmem_complex_op()(becp_gpu); } template @@ -433,7 +446,16 @@ void ElecStatePW::add_usrho(const psi::Psi& psi) { for (int is = 0; is < PARAM.inp.nspin; is++) { - this->rhopw_smooth->real2recip(this->rho[is], this->rhog[is]); + // D2H rho → CPU buffer (real2recip reads rho via CPU loops) + std::vector rho_host(this->rhopw_smooth->nrxx); + syncmem_var_d2h_op()(rho_host.data(), this->rho[is], this->rhopw_smooth->nrxx); + // D2H rhog → CPU temp + std::vector rhog_host(this->rhopw_smooth->npw); + syncmem_complex_d2h_op()(rhog_host.data(), this->rhog[is], this->rhopw_smooth->npw); + // CPU real2recip + this->rhopw_smooth->real2recip(rho_host.data(), rhog_host.data()); + // H2D rhog back to GPU + syncmem_complex_h2d_op()(this->rhog[is], rhog_host.data(), this->rhopw_smooth->npw); } } @@ -448,7 +470,16 @@ void ElecStatePW::add_usrho(const psi::Psi& psi) { for (int is = 0; is < PARAM.inp.nspin; is++) { - this->charge->rhopw->recip2real(this->rhog[is], this->rho[is]); + // D2H rhog → CPU temp (recip2real reads rhog via CPU loops) + std::vector rhog_host(this->charge->rhopw->npw); + syncmem_complex_d2h_op()(rhog_host.data(), this->rhog[is], this->charge->rhopw->npw); + // D2H rho → CPU buffer + std::vector rho_host(this->charge->rhopw->nrxx); + syncmem_var_d2h_op()(rho_host.data(), this->rho[is], this->charge->rhopw->nrxx); + // CPU recip2real + this->charge->rhopw->recip2real(rhog_host.data(), rho_host.data()); + // H2D rho back to GPU + syncmem_var_h2d_op()(this->rho[is], rho_host.data(), this->charge->rhopw->nrxx); } } } @@ -464,33 +495,38 @@ void ElecStatePW::addusdens_g(const Real* becsum, T** rhog) Structure_Factor* psf = this->ppcell->psf; const std::complex ci_tpi = ModuleBase::NEG_IMAG_UNIT * ModuleBase::TWO_PI; - Real* qmod = nullptr; - resmem_var_op()(qmod, npw, "ElecState::qmod"); - T* qgm = nullptr; - resmem_complex_op()(qgm, npw, "ElecState::qgm"); - Real* ylmk0 = nullptr; - resmem_var_op()(ylmk0, npw * lmaxq * lmaxq, "ElecState::ylmk0"); - Real* g = reinterpret_cast(this->charge->rhopw->gcar); - - ModuleBase::YlmReal::Ylm_Real(this->ctx, lmaxq * lmaxq, npw, g, ylmk0); - + // ---------- all on CPU ---------- + // Use double for CPU computations (radial_fft_q non-template version takes double) + std::vector qmod_host(npw); + std::vector> qgm_host(npw); for (int ig = 0; ig < npw; ig++) { - qmod[ig] = static_cast(this->charge->rhopw->gcar[ig].norm() * ucell->tpiba); + qmod_host[ig] = static_cast(this->charge->rhopw->gcar[ig].norm() * ucell->tpiba); } + // ylmk0: compute on GPU then D2H + Real* ylmk0_gpu = nullptr; + resmem_var_op()(ylmk0_gpu, npw * lmaxq * lmaxq, "ElecState::ylmk0"); + Real* g_gpu = nullptr; + resmem_var_op()(g_gpu, npw * 3, "ElecState::g_gpu"); + syncmem_var_h2d_op()(g_gpu, reinterpret_cast(this->charge->rhopw->gcar), npw * 3); + ModuleBase::YlmReal::Ylm_Real(this->ctx, lmaxq * lmaxq, npw, g_gpu, ylmk0_gpu); + delmem_var_op()(g_gpu); + std::vector ylmk0_host(npw * lmaxq * lmaxq); + syncmem_var_d2h_op()(ylmk0_host.data(), ylmk0_gpu, npw * lmaxq * lmaxq); + ModuleBase::matrix ylm_mat(lmaxq * lmaxq, npw); + for (int i = 0; i < npw * lmaxq * lmaxq; i++) ylm_mat.c[i] = static_cast(ylmk0_host[i]); + for (int it = 0; it < ucell->ntype; it++) { Atom* atom = &ucell->atoms[it]; if (atom->ncpp.tvanp) { - // nij = max number of (ih,jh) pairs per atom type nt const int nij = atom->ncpp.nh * (atom->ncpp.nh + 1) / 2; - T *skk = nullptr, *aux2 = nullptr, *tbecsum = nullptr; - resmem_complex_op()(skk, atom->na * npw, "ElecState::skk"); - resmem_complex_op()(aux2, nij * npw, "ElecState::aux2"); - resmem_complex_op()(tbecsum, PARAM.inp.nspin * atom->na * nij, "ElecState::tbecsum"); + // skk, tbecsum: CPU vectors + std::vector> skk_host(atom->na * npw); + std::vector> tbecsum_host(PARAM.inp.nspin * atom->na * nij); for (int ia = 0; ia < atom->na; ia++) { const int iat = ucell->itia2iat(it, ia); @@ -498,60 +534,54 @@ void ElecStatePW::addusdens_g(const Real* becsum, T** rhog) { for (int ij = 0; ij < nij; ij++) { - tbecsum[is * atom->na * nij + ia * nij + ij] - = static_cast(becsum[is * ucell->nat * nh_tot + iat * nh_tot + ij]); + tbecsum_host[is * atom->na * nij + ia * nij + ij] + = static_cast>(becsum[is * ucell->nat * nh_tot + iat * nh_tot + ij]); } } for (int ig = 0; ig < npw; ig++) { double arg = this->charge->rhopw->gcar[ig] * atom->tau[ia]; - skk[ia * npw + ig] = static_cast(ModuleBase::libm::exp(ci_tpi * arg)); + skk_host[ia * npw + ig] = ModuleBase::libm::exp(ci_tpi * arg); } } for (int is = 0; is < PARAM.inp.nspin; is++) { - // sum over atoms + // CPU BLAS: aux2 = skk * tbecsum^T + std::vector> aux2_host(nij * npw); + const std::complex one_d(1, 0), zero_d(0, 0); char transa = 'N'; char transb = 'T'; - gemm_op()(transa, - transb, - npw, - nij, - atom->na, - &one, - skk, - npw, - &tbecsum[is * atom->na * nij], - nij, - &zero, - aux2, - npw); + zgemm_(&transa, &transb, &npw, &nij, &atom->na, + &one_d, skk_host.data(), &npw, + &tbecsum_host[is * atom->na * nij], &nij, + &zero_d, aux2_host.data(), &npw); + + // D2H rhog + std::vector rhog_host(npw); + syncmem_complex_d2h_op()(rhog_host.data(), rhog[is], npw); - // sum over lm indices of Q_{lm} int ijh = 0; for (int ih = 0; ih < atom->ncpp.nh; ih++) { for (int jh = ih; jh < atom->ncpp.nh; jh++) { - this->ppcell->radial_fft_q(this->ctx, npw, ih, jh, it, qmod, ylmk0, qgm); + // CPU radial_fft_q (non-template version, double) + this->ppcell->radial_fft_q(npw, ih, jh, it, qmod_host.data(), ylm_mat, qgm_host.data()); for (int ig = 0; ig < npw; ig++) { - rhog[is][ig] += qgm[ig] * aux2[ijh * npw + ig]; + rhog_host[ig] += static_cast(qgm_host[ig] * aux2_host[ijh * npw + ig]); } ijh++; } } + // H2D rhog back + syncmem_complex_h2d_op()(rhog[is], rhog_host.data(), npw); } - delmem_complex_op()(skk); - delmem_complex_op()(aux2); - delmem_complex_op()(tbecsum); } } - delmem_var_op()(qmod); - delmem_complex_op()(qgm); - delmem_var_op()(ylmk0); + delmem_var_op()(ylmk0_gpu); } template class ElecStatePW, base_device::DEVICE_CPU>; diff --git a/source/source_estate/elecstate_pw.h b/source/source_estate/elecstate_pw.h index e8e4b95af36..5612cd2fee2 100644 --- a/source/source_estate/elecstate_pw.h +++ b/source/source_estate/elecstate_pw.h @@ -99,6 +99,17 @@ class ElecStatePW : public ElecState using resmem_complex_op = base_device::memory::resize_memory_op; using delmem_complex_op = base_device::memory::delete_memory_op; + using resmem_complex_h_op = base_device::memory::resize_memory_op; + using delmem_complex_h_op = base_device::memory::delete_memory_op; + using syncmem_complex_d2h_op = base_device::memory::synchronize_memory_op; + using syncmem_complex_h2d_op = base_device::memory::synchronize_memory_op; + + using resmem_var_h_op = base_device::memory::resize_memory_op; + using delmem_var_h_op = base_device::memory::delete_memory_op; + using setmem_var_h_op = base_device::memory::set_memory_op; + using syncmem_var_h2d_op = base_device::memory::synchronize_memory_op; + using syncmem_var_d2h_op = base_device::memory::synchronize_memory_op; + using gemv_op = ModuleBase::gemv_op; using gemm_op = ModuleBase::gemm_op; }; diff --git a/source/source_pw/module_pwdft/hamilt_pw.cpp b/source/source_pw/module_pwdft/hamilt_pw.cpp index 7c06f971690..ca30589f2b4 100644 --- a/source/source_pw/module_pwdft/hamilt_pw.cpp +++ b/source/source_pw/module_pwdft/hamilt_pw.cpp @@ -252,15 +252,23 @@ void HamiltPW::sPsi(const T* psi_in, // psi const int nh = atoms->ncpp.nh; T* qqc = nullptr; resmem_complex_op()(qqc, nh * nh, "Hamilt::qqc"); - Real* qq_now = &qq_nt[it * this->ppcell->nhm * this->ppcell->nhm]; + + std::vector qqc_host(nh*nh); + const double* qq_now_host = &this->ppcell->qq_nt.ptr[it * this->ppcell->nhm * this->ppcell->nhm]; + for (int i = 0; i < nh; i++) { for (int j = 0; j < nh; j++) { - int index = i * this->ppcell->nhm + j; - qqc[i * nh + j] = qq_now[index] * one; + const int source_index = i * this->ppcell->nhm + j; + const int target_index = i * nh + j; + + qqc_host[target_index] = static_cast(qq_now_host[source_index]) * one; } } + + syncmem_complex_h2d_op()(qqc, qqc_host.data(), qqc_host.size()); + for (int ia = 0; ia < atoms->na; ia++) { const int iat = ucell->itia2iat(it, ia); diff --git a/source/source_pw/module_pwdft/hamilt_pw.h b/source/source_pw/module_pwdft/hamilt_pw.h index ee7d228087a..898c57a9859 100644 --- a/source/source_pw/module_pwdft/hamilt_pw.h +++ b/source/source_pw/module_pwdft/hamilt_pw.h @@ -21,6 +21,7 @@ class HamiltPW : public Hamilt // return T if T is real type(float, double), // otherwise return the real type of T(complex, std::complex) using Real = typename GetTypeReal::type; + using syncmem_complex_h2d_op = base_device::memory::synchronize_memory_op; public: HamiltPW(elecstate::Potential* pot_in, From 5b44ef2113348a05c825b862fa46e096c514e912 Mon Sep 17 00:00:00 2001 From: chengleizheng Date: Thu, 16 Jul 2026 20:57:32 +0800 Subject: [PATCH 2/5] Switch radial_fft_q to template version in addusdens_g to fix linker error Co-Authored-By: Claude Opus 4.7 --- source/source_estate/elecstate_pw.cpp | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/source/source_estate/elecstate_pw.cpp b/source/source_estate/elecstate_pw.cpp index 4c76bfec02d..cbf45880411 100644 --- a/source/source_estate/elecstate_pw.cpp +++ b/source/source_estate/elecstate_pw.cpp @@ -496,7 +496,6 @@ void ElecStatePW::addusdens_g(const Real* becsum, T** rhog) const std::complex ci_tpi = ModuleBase::NEG_IMAG_UNIT * ModuleBase::TWO_PI; // ---------- all on CPU ---------- - // Use double for CPU computations (radial_fft_q non-template version takes double) std::vector qmod_host(npw); std::vector> qgm_host(npw); for (int ig = 0; ig < npw; ig++) @@ -514,8 +513,8 @@ void ElecStatePW::addusdens_g(const Real* becsum, T** rhog) delmem_var_op()(g_gpu); std::vector ylmk0_host(npw * lmaxq * lmaxq); syncmem_var_d2h_op()(ylmk0_host.data(), ylmk0_gpu, npw * lmaxq * lmaxq); - ModuleBase::matrix ylm_mat(lmaxq * lmaxq, npw); - for (int i = 0; i < npw * lmaxq * lmaxq; i++) ylm_mat.c[i] = static_cast(ylmk0_host[i]); + std::vector ylmk0_double(npw * lmaxq * lmaxq); + for (int i = 0; i < npw * lmaxq * lmaxq; i++) ylmk0_double[i] = static_cast(ylmk0_host[i]); for (int it = 0; it < ucell->ntype; it++) { @@ -566,8 +565,9 @@ void ElecStatePW::addusdens_g(const Real* becsum, T** rhog) { for (int jh = ih; jh < atom->ncpp.nh; jh++) { - // CPU radial_fft_q (non-template version, double) - this->ppcell->radial_fft_q(npw, ih, jh, it, qmod_host.data(), ylm_mat, qgm_host.data()); + // CPU radial_fft_q (template version, DEVICE_CPU) + this->ppcell->template radial_fft_q( + nullptr, npw, ih, jh, it, qmod_host.data(), ylmk0_double.data(), qgm_host.data()); for (int ig = 0; ig < npw; ig++) { rhog_host[ig] += static_cast(qgm_host[ig] * aux2_host[ijh * npw + ig]); From a0ad57a5f8bce8f87e934e1a2908f7f15eb99fed Mon Sep 17 00:00:00 2001 From: chengleizheng Date: Fri, 17 Jul 2026 17:51:45 +0800 Subject: [PATCH 3/5] Refactor: rename device/host pointers in USPP PW code --- source/source_estate/elecstate_pw.cpp | 96 ++++++++++----------- source/source_pw/module_pwdft/hamilt_pw.cpp | 9 +- 2 files changed, 52 insertions(+), 53 deletions(-) diff --git a/source/source_estate/elecstate_pw.cpp b/source/source_estate/elecstate_pw.cpp index cbf45880411..b15ce8e158d 100644 --- a/source/source_estate/elecstate_pw.cpp +++ b/source/source_estate/elecstate_pw.cpp @@ -282,13 +282,13 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) resmem_var_h_op()(becsum, nh_tot * ucell->nat * PARAM.inp.nspin, "ElecState::becsum"); setmem_var_h_op()(becsum, 0, nh_tot * ucell->nat * PARAM.inp.nspin); - // becp: GPU for first gemm, then D2H for CPU loops - T* becp_gpu = nullptr; - std::vector becp; + // becp: device buffer for gemm, then D2H for host loops + T* becp = nullptr; + std::vector becp_host; if (nkb > 0) { - resmem_complex_op()(becp_gpu, nbands * nkb, "ElecState::becp"); - becp.resize(nbands * nkb); + resmem_complex_op()(becp, nbands * nkb, "ElecState::becp"); + becp_host.resize(nbands * nkb); } for (int ik = 0; ik < psi.get_nk(); ++ik) @@ -304,7 +304,7 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) this->ppcell->getvnl(this->ctx, *ucell,ik, this->vkb); } - // becp = (GPU gemm) + // becp = (device gemm) char transa = 'C'; char transb = 'N'; if (this->ppcell->nkb > 0) @@ -321,7 +321,7 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) psi_now, inc, &zero, - becp_gpu, + becp, inc); } else @@ -337,13 +337,13 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) psi_now, npwx, &zero, - becp_gpu, + becp, this->ppcell->nkb); } - // D2H: GPU becp → CPU becp - syncmem_complex_d2h_op()(becp.data(), becp_gpu, nbands * nkb); + // D2H: device becp → host becp_host + syncmem_complex_d2h_op()(becp_host.data(), becp, nbands * nkb); } - Parallel_Reduce::reduce_pool(becp.data(), this->ppcell->nkb * nbands); + Parallel_Reduce::reduce_pool(becp_host.data(), this->ppcell->nkb * nbands); // sum over bands: \sum_i w_i for (int it = 0; it < ucell->ntype; it++) @@ -352,9 +352,14 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) if (atom->ncpp.tvanp) { const int nh_atom = atom->ncpp.nh; - // auxk1, auxk2 on CPU for filling loops - std::vector auxk1(nbands * nh_atom); - std::vector auxk2(nbands * nh_atom); + // auxk1, auxk2: host buffers for filling loops + std::vector auxk1_host(nbands * nh_atom); + std::vector auxk2_host(nbands * nh_atom); + // device buffers allocated once per atom type + T *aux_gk = nullptr, *auxk1 = nullptr, *auxk2 = nullptr; + resmem_complex_op()(auxk1, nbands * nh_atom, "ElecState::auxk1"); + resmem_complex_op()(auxk2, nbands * nh_atom, "ElecState::auxk2"); + resmem_complex_op()(aux_gk, nh_atom * nh_atom * npol * npol, "ElecState::aux_gk"); for (int ia = 0; ia < atom->na; ia++) { const int iat = ucell->itia2iat(it, ia); @@ -369,19 +374,15 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) const int ikb = this->ppcell->indv_ijkb0[iat] + ih; for (int ib = 0; ib < nbands; ib++) { - auxk1[ih * nbands + ib] = becp[ib * this->ppcell->nkb + ikb]; - auxk2[ih * nbands + ib] - = becp[ib * this->ppcell->nkb + ikb] * static_cast(this->wg(ik, ib)); + auxk1_host[ih * nbands + ib] = becp_host[ib * this->ppcell->nkb + ikb]; + auxk2_host[ih * nbands + ib] + = becp_host[ib * this->ppcell->nkb + ikb] * static_cast(this->wg(ik, ib)); } } - // GPU gemm: aux_gk = auxk1^H * auxk2 - T *aux_gk_gpu = nullptr, *auxk1_gpu = nullptr, *auxk2_gpu = nullptr; - resmem_complex_op()(auxk1_gpu, nbands * nh_atom, "ElecState::auxk1"); - resmem_complex_op()(auxk2_gpu, nbands * nh_atom, "ElecState::auxk2"); - resmem_complex_op()(aux_gk_gpu, nh_atom * nh_atom * npol * npol, "ElecState::aux_gk"); - syncmem_complex_h2d_op()(auxk1_gpu, auxk1.data(), nbands * nh_atom); - syncmem_complex_h2d_op()(auxk2_gpu, auxk2.data(), nbands * nh_atom); + // device gemm: aux_gk = auxk1^H * auxk2 + syncmem_complex_h2d_op()(auxk1, auxk1_host.data(), nbands * nh_atom); + syncmem_complex_h2d_op()(auxk2, auxk2_host.data(), nbands * nh_atom); char transa2 = 'C'; char transb2 = 'N'; @@ -391,21 +392,17 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) nh_atom, nbands, &one, - auxk1_gpu, + auxk1, nbands, - auxk2_gpu, + auxk2, nbands, &zero, - aux_gk_gpu, + aux_gk, nh_atom); - // D2H: GPU aux_gk → CPU - std::vector aux_gk(nh_atom * nh_atom); - syncmem_complex_d2h_op()(aux_gk.data(), aux_gk_gpu, nh_atom * nh_atom); - - delmem_complex_op()(auxk1_gpu); - delmem_complex_op()(auxk2_gpu); - delmem_complex_op()(aux_gk_gpu); + // D2H: device aux_gk → host + std::vector aux_gk_host(nh_atom * nh_atom); + syncmem_complex_d2h_op()(aux_gk_host.data(), aux_gk, nh_atom * nh_atom); // copy output from GEMM into desired format int ijh = 0; @@ -416,21 +413,24 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) { if (ih == jh) { - becsum[index + ijh] += std::real(aux_gk[ih * nh_atom + jh]); + becsum[index + ijh] += std::real(aux_gk_host[ih * nh_atom + jh]); } else { - becsum[index + ijh] += 2.0 * std::real(aux_gk[ih * nh_atom + jh]); + becsum[index + ijh] += 2.0 * std::real(aux_gk_host[ih * nh_atom + jh]); } ijh++; } } } } + delmem_complex_op()(auxk1); + delmem_complex_op()(auxk2); + delmem_complex_op()(aux_gk); } } } - if (becp_gpu != nullptr) delmem_complex_op()(becp_gpu); + delmem_complex_op()(becp); } template @@ -454,7 +454,7 @@ void ElecStatePW::add_usrho(const psi::Psi& psi) syncmem_complex_d2h_op()(rhog_host.data(), this->rhog[is], this->rhopw_smooth->npw); // CPU real2recip this->rhopw_smooth->real2recip(rho_host.data(), rhog_host.data()); - // H2D rhog back to GPU + // H2D rhog back to device syncmem_complex_h2d_op()(this->rhog[is], rhog_host.data(), this->rhopw_smooth->npw); } } @@ -503,16 +503,16 @@ void ElecStatePW::addusdens_g(const Real* becsum, T** rhog) qmod_host[ig] = static_cast(this->charge->rhopw->gcar[ig].norm() * ucell->tpiba); } - // ylmk0: compute on GPU then D2H - Real* ylmk0_gpu = nullptr; - resmem_var_op()(ylmk0_gpu, npw * lmaxq * lmaxq, "ElecState::ylmk0"); - Real* g_gpu = nullptr; - resmem_var_op()(g_gpu, npw * 3, "ElecState::g_gpu"); - syncmem_var_h2d_op()(g_gpu, reinterpret_cast(this->charge->rhopw->gcar), npw * 3); - ModuleBase::YlmReal::Ylm_Real(this->ctx, lmaxq * lmaxq, npw, g_gpu, ylmk0_gpu); - delmem_var_op()(g_gpu); + // ylmk0: compute on device then D2H + Real* ylmk0 = nullptr; + resmem_var_op()(ylmk0, npw * lmaxq * lmaxq, "ElecState::ylmk0"); + Real* g = nullptr; + resmem_var_op()(g, npw * 3, "ElecState::g"); + syncmem_var_h2d_op()(g, reinterpret_cast(this->charge->rhopw->gcar), npw * 3); + ModuleBase::YlmReal::Ylm_Real(this->ctx, lmaxq * lmaxq, npw, g, ylmk0); + delmem_var_op()(g); std::vector ylmk0_host(npw * lmaxq * lmaxq); - syncmem_var_d2h_op()(ylmk0_host.data(), ylmk0_gpu, npw * lmaxq * lmaxq); + syncmem_var_d2h_op()(ylmk0_host.data(), ylmk0, npw * lmaxq * lmaxq); std::vector ylmk0_double(npw * lmaxq * lmaxq); for (int i = 0; i < npw * lmaxq * lmaxq; i++) ylmk0_double[i] = static_cast(ylmk0_host[i]); @@ -581,7 +581,7 @@ void ElecStatePW::addusdens_g(const Real* becsum, T** rhog) } } - delmem_var_op()(ylmk0_gpu); + delmem_var_op()(ylmk0); } template class ElecStatePW, base_device::DEVICE_CPU>; diff --git a/source/source_pw/module_pwdft/hamilt_pw.cpp b/source/source_pw/module_pwdft/hamilt_pw.cpp index ca30589f2b4..e095630b6a2 100644 --- a/source/source_pw/module_pwdft/hamilt_pw.cpp +++ b/source/source_pw/module_pwdft/hamilt_pw.cpp @@ -252,7 +252,6 @@ void HamiltPW::sPsi(const T* psi_in, // psi const int nh = atoms->ncpp.nh; T* qqc = nullptr; resmem_complex_op()(qqc, nh * nh, "Hamilt::qqc"); - std::vector qqc_host(nh*nh); const double* qq_now_host = &this->ppcell->qq_nt.ptr[it * this->ppcell->nhm * this->ppcell->nhm]; @@ -260,13 +259,13 @@ void HamiltPW::sPsi(const T* psi_in, // psi { for (int j = 0; j < nh; j++) { - const int source_index = i * this->ppcell->nhm + j; - const int target_index = i * nh + j; + const int source_index = i * this->ppcell->nhm + j; + const int target_index = i * nh + j; - qqc_host[target_index] = static_cast(qq_now_host[source_index]) * one; + qqc_host[target_index] = static_cast(qq_now_host[source_index]) * one; } } - + syncmem_complex_h2d_op()(qqc, qqc_host.data(), qqc_host.size()); for (int ia = 0; ia < atoms->na; ia++) From abb36d46f23068ed7302cb0c6cb64977c2f2e7e2 Mon Sep 17 00:00:00 2001 From: chengleizheng Date: Sat, 18 Jul 2026 10:41:42 +0800 Subject: [PATCH 4/5] minor formatting cleanup --- source/source_pw/module_pwdft/hamilt_pw.cpp | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/source/source_pw/module_pwdft/hamilt_pw.cpp b/source/source_pw/module_pwdft/hamilt_pw.cpp index e095630b6a2..6b858428f54 100644 --- a/source/source_pw/module_pwdft/hamilt_pw.cpp +++ b/source/source_pw/module_pwdft/hamilt_pw.cpp @@ -252,21 +252,21 @@ void HamiltPW::sPsi(const T* psi_in, // psi const int nh = atoms->ncpp.nh; T* qqc = nullptr; resmem_complex_op()(qqc, nh * nh, "Hamilt::qqc"); - std::vector qqc_host(nh*nh); - const double* qq_now_host = &this->ppcell->qq_nt.ptr[it * this->ppcell->nhm * this->ppcell->nhm]; - + std::vector qqc_host(nh*nh); + const double* qq_now_host = &this->ppcell->qq_nt.ptr[it * this->ppcell->nhm * this->ppcell->nhm]; + for (int i = 0; i < nh; i++) { for (int j = 0; j < nh; j++) { - const int source_index = i * this->ppcell->nhm + j; - const int target_index = i * nh + j; + const int source_index = i * this->ppcell->nhm + j; + const int target_index = i * nh + j; - qqc_host[target_index] = static_cast(qq_now_host[source_index]) * one; + qqc_host[target_index] = static_cast(qq_now_host[source_index]) * one; } } - - syncmem_complex_h2d_op()(qqc, qqc_host.data(), qqc_host.size()); + + syncmem_complex_h2d_op()(qqc, qqc_host.data(), qqc_host.size()); for (int ia = 0; ia < atoms->na; ia++) { From 27d7badb0707536e09cc6989ef65ec75df963695 Mon Sep 17 00:00:00 2001 From: chengleizheng Date: Sat, 18 Jul 2026 10:50:14 +0800 Subject: [PATCH 5/5] Add TODO comments for future USPP GPU optimizations --- source/source_estate/elecstate_pw.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/source/source_estate/elecstate_pw.cpp b/source/source_estate/elecstate_pw.cpp index b15ce8e158d..befcfe162d7 100644 --- a/source/source_estate/elecstate_pw.cpp +++ b/source/source_estate/elecstate_pw.cpp @@ -369,6 +369,7 @@ void ElecStatePW::cal_becsum(const psi::Psi& psi) } else { + // TODO: gather auxk from becp on device to skip the H2D→fill→D2H round-trip for (int ih = 0; ih < nh_atom; ih++) { const int ikb = this->ppcell->indv_ijkb0[iat] + ih; @@ -453,6 +454,7 @@ void ElecStatePW::add_usrho(const psi::Psi& psi) std::vector rhog_host(this->rhopw_smooth->npw); syncmem_complex_d2h_op()(rhog_host.data(), this->rhog[is], this->rhopw_smooth->npw); // CPU real2recip + // TODO: replace with cufft to keep rho/rhog on device this->rhopw_smooth->real2recip(rho_host.data(), rhog_host.data()); // H2D rhog back to device syncmem_complex_h2d_op()(this->rhog[is], rhog_host.data(), this->rhopw_smooth->npw); @@ -496,6 +498,7 @@ void ElecStatePW::addusdens_g(const Real* becsum, T** rhog) const std::complex ci_tpi = ModuleBase::NEG_IMAG_UNIT * ModuleBase::TWO_PI; // ---------- all on CPU ---------- + // TODO: port skk/tbecsum construction and radial_fft_q to device std::vector qmod_host(npw); std::vector> qgm_host(npw); for (int ig = 0; ig < npw; ig++)