From cba00f5961e340a9ecc0b939fbcc2ce4e176b698 Mon Sep 17 00:00:00 2001 From: Chen Nuo <49788094+Cstandardlib@users.noreply.github.com> Date: Tue, 28 Apr 2026 23:54:02 +0800 Subject: [PATCH 1/5] [Fix] Skip unnecessary vkb CPU allocation on GPU path in pseudopot_cell_vnl MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On GPU path, vkb.create(nkb, npwx) allocates CPU ComplexMatrix memory that is never used — getvnl() writes directly to GPU buffers (c_vkb/z_vkb). The only consumer of vkb.nc metadata is the leading dimension in gemm/gemv. This wastes nkb*npwx*16 bytes of CPU memory (~3.2 GB for large systems). Changes: - Add vkbnc member to store column dimension independently - Guard vkb.create() behind !use_gpu_ in init() - Replace all ppcell->vkb.nc with ppcell->vkbnc (op_pw_nl.cpp, hamilt_pw.cpp) - Add lazy-allocation guard in getgradq_vnl() for GPU Velocity path Tested: GPU build + 28/28 kernel UTs + 38/40 GPU integration tests (2 pre-existing failures: scf_bpcg, scf_out_wf) --- source/source_pw/module_pwdft/hamilt_pw.cpp | 8 ++++---- source/source_pw/module_pwdft/op_pw_nl.cpp | 8 ++++---- source/source_pw/module_pwdft/vnl_pw.cpp | 11 +++++++++-- source/source_pw/module_pwdft/vnl_pw.h | 4 ++++ source/source_pw/module_pwdft/vnl_pw_grad.cpp | 5 +++++ 5 files changed, 26 insertions(+), 10 deletions(-) diff --git a/source/source_pw/module_pwdft/hamilt_pw.cpp b/source/source_pw/module_pwdft/hamilt_pw.cpp index 27a56cbe11b..152e3451428 100644 --- a/source/source_pw/module_pwdft/hamilt_pw.cpp +++ b/source/source_pw/module_pwdft/hamilt_pw.cpp @@ -272,7 +272,7 @@ void HamiltPW::sPsi(const T* psi_in, // psi this->ppcell->nkb, &one, this->vkb, - this->ppcell->vkb.nc, + this->ppcell->vkbnc, psi_in, inc, &zero, @@ -288,7 +288,7 @@ void HamiltPW::sPsi(const T* psi_in, // psi npw, &one, this->vkb, - this->ppcell->vkb.nc, + this->ppcell->vkbnc, psi_in, nrow, &zero, @@ -360,7 +360,7 @@ void HamiltPW::sPsi(const T* psi_in, // psi this->ppcell->nkb, &one, this->vkb, - this->ppcell->vkb.nc, + this->ppcell->vkbnc, ps, inc, &one, @@ -376,7 +376,7 @@ void HamiltPW::sPsi(const T* psi_in, // psi this->ppcell->nkb, &one, this->vkb, - this->ppcell->vkb.nc, + this->ppcell->vkbnc, ps, this->ppcell->nkb, &one, diff --git a/source/source_pw/module_pwdft/op_pw_nl.cpp b/source/source_pw/module_pwdft/op_pw_nl.cpp index 0e5de357b1d..d3551808ea7 100644 --- a/source/source_pw/module_pwdft/op_pw_nl.cpp +++ b/source/source_pw/module_pwdft/op_pw_nl.cpp @@ -172,7 +172,7 @@ void Nonlocal>::add_nonlocal_pp(T *hpsi_in, const T *becp, this->ppcell->nkb, &this->one, this->vkb, - this->ppcell->vkb.nc, + this->ppcell->vkbnc, this->ps, inc, &this->one, @@ -197,7 +197,7 @@ void Nonlocal>::add_nonlocal_pp(T *hpsi_in, const T *becp, this->ppcell->nkb, &this->one, this->vkb, - this->ppcell->vkb.nc, + this->ppcell->vkbnc, this->ps, npm, &this->one, @@ -251,7 +251,7 @@ void Nonlocal>::act( nkb, &this->one, this->vkb, - this->ppcell->vkb.nc, + this->ppcell->vkbnc, tmpsi_in, inc, &this->zero, @@ -276,7 +276,7 @@ void Nonlocal>::act( this->npw, &this->one, this->vkb, - this->ppcell->vkb.nc, + this->ppcell->vkbnc, tmpsi_in, max_npw, &this->zero, diff --git a/source/source_pw/module_pwdft/vnl_pw.cpp b/source/source_pw/module_pwdft/vnl_pw.cpp index d6b5274d51f..83c42adac89 100644 --- a/source/source_pw/module_pwdft/vnl_pw.cpp +++ b/source/source_pw/module_pwdft/vnl_pw.cpp @@ -214,10 +214,17 @@ void pseudopot_cell_vnl::init(const UnitCell& ucell, // dq+4)*cell_factor; this->lmaxq = 2 * this->lmaxkb + 1; int npwx = this->wfcpw->npwk_max; + this->vkbnc = npwx; if (nkb > 0 && allocate_vkb) { - vkb.create(nkb, npwx); - ModuleBase::Memory::record("VNL::vkb", nkb * npwx * sizeof(std::complex)); + if (!this->use_gpu_) + { + vkb.create(nkb, npwx); + ModuleBase::Memory::record("VNL::vkb", nkb * npwx * sizeof(std::complex)); + } + // GPU path: vkb ComplexMatrix is not allocated. + // Column dimension is stored in vkbnc for gemm/gemv leading dimension. + // Actual GPU buffers (c_vkb/z_vkb) are allocated below. } // this->nqx = 10000; // calculted in allocate_nlpot.f90 diff --git a/source/source_pw/module_pwdft/vnl_pw.h b/source/source_pw/module_pwdft/vnl_pw.h index 93a593e9257..6282b138657 100644 --- a/source/source_pw/module_pwdft/vnl_pw.h +++ b/source/source_pw/module_pwdft/vnl_pw.h @@ -108,6 +108,10 @@ class pseudopot_cell_vnl std::complex*** vkb_alpha; Structure_Factor* psf = nullptr; + // Column dimension of vkb matrix (= npwx), used as leading dimension in gemm/gemv. + // On GPU path vkb ComplexMatrix is not allocated to save CPU memory; this stores the dimension. + int vkbnc = 0; + // other variables std::complex Cal_C(int alpha, int lu, int mu, int L, int M); diff --git a/source/source_pw/module_pwdft/vnl_pw_grad.cpp b/source/source_pw/module_pwdft/vnl_pw_grad.cpp index 135fe475944..65984cf581a 100644 --- a/source/source_pw/module_pwdft/vnl_pw_grad.cpp +++ b/source/source_pw/module_pwdft/vnl_pw_grad.cpp @@ -91,6 +91,11 @@ void pseudopot_cell_vnl::getgradq_vnl(const UnitCell& ucell, ModuleBase::YlmReal::grad_Ylm_Real(x1, npw, gk, ylm, dylm[0], dylm[1], dylm[2]); + // GPU path skips vkb allocation in init(); allocate now if needed + if (this->vkb.nc == 0 && this->nkb > 0 && this->vkbnc > 0) { + this->vkb.create(this->nkb, this->vkbnc); + } + int jkb = 0; for(int it = 0;it < ucell.ntype;it++) { From a0f43dc23578a3c3de2dd19b31c4b207f5b05809 Mon Sep 17 00:00:00 2001 From: Chen Nuo <49788094+Cstandardlib@users.noreply.github.com> Date: Wed, 29 Apr 2026 01:06:41 +0800 Subject: [PATCH 2/5] Update source/source_pw/module_pwdft/vnl_pw_grad.cpp Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- source/source_pw/module_pwdft/vnl_pw_grad.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/source/source_pw/module_pwdft/vnl_pw_grad.cpp b/source/source_pw/module_pwdft/vnl_pw_grad.cpp index 65984cf581a..df02364fc6d 100644 --- a/source/source_pw/module_pwdft/vnl_pw_grad.cpp +++ b/source/source_pw/module_pwdft/vnl_pw_grad.cpp @@ -94,6 +94,7 @@ void pseudopot_cell_vnl::getgradq_vnl(const UnitCell& ucell, // GPU path skips vkb allocation in init(); allocate now if needed if (this->vkb.nc == 0 && this->nkb > 0 && this->vkbnc > 0) { this->vkb.create(this->nkb, this->vkbnc); + ModuleBase::Memory::record("VNL::vkb", static_cast(this->nkb) * this->vkbnc * sizeof(double)); } int jkb = 0; From b71e3fe091535bab9e9dd1e35b69be8feb54e56e Mon Sep 17 00:00:00 2001 From: Chen Nuo <49788094+Cstandardlib@users.noreply.github.com> Date: Wed, 29 Apr 2026 01:07:16 +0800 Subject: [PATCH 3/5] Update source/source_pw/module_pwdft/vnl_pw_grad.cpp Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- source/source_pw/module_pwdft/vnl_pw_grad.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/source/source_pw/module_pwdft/vnl_pw_grad.cpp b/source/source_pw/module_pwdft/vnl_pw_grad.cpp index df02364fc6d..6b134abe568 100644 --- a/source/source_pw/module_pwdft/vnl_pw_grad.cpp +++ b/source/source_pw/module_pwdft/vnl_pw_grad.cpp @@ -91,8 +91,10 @@ void pseudopot_cell_vnl::getgradq_vnl(const UnitCell& ucell, ModuleBase::YlmReal::grad_Ylm_Real(x1, npw, gk, ylm, dylm[0], dylm[1], dylm[2]); - // GPU path skips vkb allocation in init(); allocate now if needed - if (this->vkb.nc == 0 && this->nkb > 0 && this->vkbnc > 0) { + // GPU path skips vkb allocation in init(); allocate now if needed, + // and resize if the required dimensions changed after a later init(). + if (this->nkb > 0 && this->vkbnc > 0 + && (this->vkb.nr != this->nkb || this->vkb.nc != this->vkbnc)) { this->vkb.create(this->nkb, this->vkbnc); ModuleBase::Memory::record("VNL::vkb", static_cast(this->nkb) * this->vkbnc * sizeof(double)); } From b963e3bf99782eedd0bcef97167b9f4c5d19925c Mon Sep 17 00:00:00 2001 From: Chen Nuo <49788094+Cstandardlib@users.noreply.github.com> Date: Wed, 29 Apr 2026 11:45:00 +0800 Subject: [PATCH 4/5] Revert "Update source/source_pw/module_pwdft/vnl_pw_grad.cpp" This reverts commit b71e3fe091535bab9e9dd1e35b69be8feb54e56e. --- source/source_pw/module_pwdft/vnl_pw_grad.cpp | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/source/source_pw/module_pwdft/vnl_pw_grad.cpp b/source/source_pw/module_pwdft/vnl_pw_grad.cpp index 6b134abe568..df02364fc6d 100644 --- a/source/source_pw/module_pwdft/vnl_pw_grad.cpp +++ b/source/source_pw/module_pwdft/vnl_pw_grad.cpp @@ -91,10 +91,8 @@ void pseudopot_cell_vnl::getgradq_vnl(const UnitCell& ucell, ModuleBase::YlmReal::grad_Ylm_Real(x1, npw, gk, ylm, dylm[0], dylm[1], dylm[2]); - // GPU path skips vkb allocation in init(); allocate now if needed, - // and resize if the required dimensions changed after a later init(). - if (this->nkb > 0 && this->vkbnc > 0 - && (this->vkb.nr != this->nkb || this->vkb.nc != this->vkbnc)) { + // GPU path skips vkb allocation in init(); allocate now if needed + if (this->vkb.nc == 0 && this->nkb > 0 && this->vkbnc > 0) { this->vkb.create(this->nkb, this->vkbnc); ModuleBase::Memory::record("VNL::vkb", static_cast(this->nkb) * this->vkbnc * sizeof(double)); } From 0999a78c24120ccd480843ccc7e54982a54930c4 Mon Sep 17 00:00:00 2001 From: Chen Nuo <49788094+Cstandardlib@users.noreply.github.com> Date: Wed, 29 Apr 2026 12:01:58 +0800 Subject: [PATCH 5/5] Revert "Update source/source_pw/module_pwdft/vnl_pw_grad.cpp" This reverts commit a0f43dc23578a3c3de2dd19b31c4b207f5b05809. --- source/source_pw/module_pwdft/vnl_pw_grad.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/source/source_pw/module_pwdft/vnl_pw_grad.cpp b/source/source_pw/module_pwdft/vnl_pw_grad.cpp index df02364fc6d..65984cf581a 100644 --- a/source/source_pw/module_pwdft/vnl_pw_grad.cpp +++ b/source/source_pw/module_pwdft/vnl_pw_grad.cpp @@ -94,7 +94,6 @@ void pseudopot_cell_vnl::getgradq_vnl(const UnitCell& ucell, // GPU path skips vkb allocation in init(); allocate now if needed if (this->vkb.nc == 0 && this->nkb > 0 && this->vkbnc > 0) { this->vkb.create(this->nkb, this->vkbnc); - ModuleBase::Memory::record("VNL::vkb", static_cast(this->nkb) * this->vkbnc * sizeof(double)); } int jkb = 0;