From bc21c5300e6fb4103421ce243c434a9ada747c17 Mon Sep 17 00:00:00 2001 From: chengleizheng Date: Tue, 4 Aug 2026 23:26:19 +0800 Subject: [PATCH 1/3] Fix(GPU): correct spin stride of rho in elecstate_pw_op kernel The GPU kernel indexed the spin components of rho with spin*nrxx, where nrxx is the wavefunction real-space grid size, while rho is allocated with charge->nrxx (density grid) as the per-spin stride. For USPP tests using the double grid the two grids differ, so the spin-down density was written to the wrong offset and effectively lost. Pass an explicit rho_stride parameter instead. Fixes 007_PW_UPF201_USPP_Fe GPU SCF etot being off by 6.07 eV. --- source/source_estate/elecstate_pw.cpp | 5 ++- source/source_estate/elecstate_pw_cal_tau.cpp | 2 +- .../kernels/cuda/elecstate_op.cu | 38 +++++++++++-------- source/source_estate/kernels/elecstate_op.cpp | 2 + source/source_estate/kernels/elecstate_op.h | 6 +++ source/source_io/module_wf/read_wf2rho_pw.cpp | 3 +- source/source_pw/module_stodft/sto_iter.cpp | 2 +- 7 files changed, 37 insertions(+), 21 deletions(-) diff --git a/source/source_estate/elecstate_pw.cpp b/source/source_estate/elecstate_pw.cpp index 7edb7b6bcf..32e1e84d62 100644 --- a/source/source_estate/elecstate_pw.cpp +++ b/source/source_estate/elecstate_pw.cpp @@ -218,6 +218,7 @@ void ElecStatePW::rhoBandK(const psi::Psi& psi) PARAM.globalv.domag, PARAM.globalv.domag_z, this->basis->nrxx, + this->charge->nrxx, w1, this->rho, this->wfcr, @@ -240,7 +241,7 @@ void ElecStatePW::rhoBandK(const psi::Psi& psi) if (w1 != 0.0) { // replaced by denghui at 20221110 - elecstate_pw_op()(this->ctx, current_spin, this->basis->nrxx, w1, this->rho, this->wfcr); + elecstate_pw_op()(this->ctx, current_spin, this->basis->nrxx, this->charge->nrxx, w1, this->rho, this->wfcr); } // kinetic energy density @@ -263,7 +264,7 @@ void ElecStatePW::rhoBandK(const psi::Psi& psi) 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); + elecstate_pw_op()(this->ctx, current_spin, this->charge->nrxx, this->charge->nrxx, w1, this->kin_r, this->wfcr); } } } diff --git a/source/source_estate/elecstate_pw_cal_tau.cpp b/source/source_estate/elecstate_pw_cal_tau.cpp index a59990600a..5c0b6ae9f5 100644 --- a/source/source_estate/elecstate_pw_cal_tau.cpp +++ b/source/source_estate/elecstate_pw_cal_tau.cpp @@ -45,7 +45,7 @@ void ElecStatePW::cal_tau(const psi::Psi& psi) 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); + elecstate_pw_op()(this->ctx, current_spin, this->charge->nrxx, this->charge->nrxx, w1, this->kin_r, this->wfcr); } } } diff --git a/source/source_estate/kernels/cuda/elecstate_op.cu b/source/source_estate/kernels/cuda/elecstate_op.cu index 4597a1f1ad..7ffd65e3d4 100644 --- a/source/source_estate/kernels/cuda/elecstate_op.cu +++ b/source/source_estate/kernels/cuda/elecstate_op.cu @@ -14,11 +14,12 @@ __global__ void elecstate_pw( const int nrxx, const FPTYPE w1, FPTYPE* rho, - const thrust::complex* wfcr) + const thrust::complex* wfcr, + const int rho_stride) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if(idx >= nrxx) {return;} - rho[spin * nrxx + idx] += w1 * norm(wfcr[idx]); + rho[spin * rho_stride + idx] += w1 * norm(wfcr[idx]); } template @@ -29,31 +30,32 @@ __global__ void elecstate_pw( const FPTYPE w1, FPTYPE* rho, const thrust::complex* wfcr, - const thrust::complex* wfcr_another_spin) + const thrust::complex* wfcr_another_spin, + const int rho_stride) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if(idx >= nrxx) {return;} - rho[0 * nrxx + idx] += w1 * (norm(wfcr[idx]) + norm(wfcr_another_spin[idx])); + rho[0 * rho_stride + idx] += w1 * (norm(wfcr[idx]) + norm(wfcr_another_spin[idx])); if (DOMAG) { - rho[1 * nrxx + idx] += w1 * 2.0 + rho[1 * rho_stride + idx] += w1 * 2.0 * (wfcr[idx].real() * wfcr_another_spin[idx].real() + wfcr[idx].imag() * wfcr_another_spin[idx].imag()); - rho[2 * nrxx + idx] += w1 * 2.0 + rho[2 * rho_stride + idx] += w1 * 2.0 * (wfcr[idx].real() * wfcr_another_spin[idx].imag() - wfcr_another_spin[idx].real() * wfcr[idx].imag()); - rho[3 * nrxx + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx])); + rho[3 * rho_stride + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx])); } else if(DOMAG_Z) { - rho[1 * nrxx + idx] = 0; - rho[2 * nrxx + idx] = 0; - rho[3 * nrxx + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx])); + rho[1 * rho_stride + idx] = 0; + rho[2 * rho_stride + idx] = 0; + rho[3 * rho_stride + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx])); } else { - rho[0 * nrxx + idx] = 0; - rho[1 * nrxx + idx] = 0; - rho[2 * nrxx + idx] = 0; - rho[3 * nrxx + idx] = 0; + rho[0 * rho_stride + idx] = 0; + rho[1 * rho_stride + idx] = 0; + rho[2 * rho_stride + idx] = 0; + rho[3 * rho_stride + idx] = 0; } } @@ -61,6 +63,7 @@ template void elecstate_pw_op::operator()(const base_device::DEVICE_GPU* ctx, const int& spin, const int& nrxx, + const int& rho_stride, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr) @@ -68,7 +71,8 @@ void elecstate_pw_op::operator()(const base_dev const int block = (nrxx + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK; elecstate_pw<<>>( spin, nrxx, w1, rho[0], - reinterpret_cast*>(wfcr) + reinterpret_cast*>(wfcr), + rho_stride ); CHECK_CUDA_SYNC(); @@ -79,6 +83,7 @@ void elecstate_pw_op::operator()(const base_dev const bool& DOMAG, const bool& DOMAG_Z, const int& nrxx, + const int& rho_stride, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr, @@ -88,7 +93,8 @@ void elecstate_pw_op::operator()(const base_dev elecstate_pw<<>>( DOMAG, DOMAG_Z, nrxx, w1, rho[0], reinterpret_cast*>(wfcr), - reinterpret_cast*>(wfcr_another_spin) + reinterpret_cast*>(wfcr_another_spin), + rho_stride ); CHECK_CUDA_SYNC(); diff --git a/source/source_estate/kernels/elecstate_op.cpp b/source/source_estate/kernels/elecstate_op.cpp index c3933319fc..9960b44358 100644 --- a/source/source_estate/kernels/elecstate_op.cpp +++ b/source/source_estate/kernels/elecstate_op.cpp @@ -8,6 +8,7 @@ struct elecstate_pw_op void operator()(const base_device::DEVICE_CPU* /*ctx*/, const int& spin, const int& nrxx, + const int& /*rho_stride*/, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr) @@ -29,6 +30,7 @@ struct elecstate_pw_op const bool& DOMAG, const bool& DOMAG_Z, const int& nrxx, + const int& /*rho_stride*/, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr, diff --git a/source/source_estate/kernels/elecstate_op.h b/source/source_estate/kernels/elecstate_op.h index f7b8d48c2e..090b7e00e9 100644 --- a/source/source_estate/kernels/elecstate_op.h +++ b/source/source_estate/kernels/elecstate_op.h @@ -15,6 +15,7 @@ struct elecstate_pw_op { /// @param ctx - which device this function runs on /// @param spin - current spin /// @param nrxx - number of planewaves + /// @param rho_stride - stride between spin components of rho /// @param weight - input constant /// @param wfcr - input array, psi in real space /// @@ -24,6 +25,7 @@ struct elecstate_pw_op { const Device* ctx, const int& spin, const int& nrxx, + const int& rho_stride, const FPTYPE& weight, FPTYPE** rho, const std::complex* wfcr); @@ -35,6 +37,7 @@ struct elecstate_pw_op { /// @param DOMAG - PARAM.globalv.domag /// @param DOMAG_Z - PARAM.globalv.domag_z /// @param nrxx - number of planewaves + /// @param rho_stride - stride between spin components of rho /// @param weight - input constant /// @param wfcr - input array, psi in real space /// @param wfcr_another_spin - input array, psi in real space @@ -46,6 +49,7 @@ struct elecstate_pw_op { const bool& DOMAG, const bool& DOMAG_Z, const int& nrxx, + const int& rho_stride, const FPTYPE& weight, FPTYPE** rho, const std::complex* wfcr, @@ -59,6 +63,7 @@ struct elecstate_pw_op void operator()(const base_device::DEVICE_GPU* ctx, const int& spin, const int& nrxx, + const int& rho_stride, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr); @@ -67,6 +72,7 @@ struct elecstate_pw_op const bool& DOMAG, const bool& DOMAG_Z, const int& nrxx, + const int& rho_stride, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr, diff --git a/source/source_io/module_wf/read_wf2rho_pw.cpp b/source/source_io/module_wf/read_wf2rho_pw.cpp index 7122ee3f12..98592eed1b 100644 --- a/source/source_io/module_wf/read_wf2rho_pw.cpp +++ b/source/source_io/module_wf/read_wf2rho_pw.cpp @@ -139,6 +139,7 @@ void ModuleIO::read_wf2rho_pw( PARAM.globalv.domag, PARAM.globalv.domag_z, nrxx, + nrxx, w1, chg.rho, rho_tmp.data(), @@ -158,7 +159,7 @@ void ModuleIO::read_wf2rho_pw( if (w1 != 0.0) { base_device::DEVICE_CPU* ctx = nullptr; - elecstate::elecstate_pw_op()(ctx, is, nrxx, + elecstate::elecstate_pw_op()(ctx, is, nrxx, nrxx, w1, chg.rho, rho_tmp.data()); } } diff --git a/source/source_pw/module_stodft/sto_iter.cpp b/source/source_pw/module_stodft/sto_iter.cpp index a46de4cc15..1b3709bdb5 100644 --- a/source/source_pw/module_stodft/sto_iter.cpp +++ b/source/source_pw/module_stodft/sto_iter.cpp @@ -645,7 +645,7 @@ void Stochastic_Iter::cal_storho(const UnitCell& ucell, { 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); + elecstate::elecstate_pw_op()(this->ctx, current_spin, nrxx, pes->charge->nrxx, w1, pes->rho, porter); // for (int ir = 0; ir < nrxx; ++ir) // { // pes->charge->rho[0][ir] += norm(porter[ir]) * this->pkv->wk[ik]; From 96ef4cd6241a4a4a95d7814aee519593c06604f2 Mon Sep 17 00:00:00 2001 From: chengleizheng Date: Wed, 5 Aug 2026 09:53:06 +0800 Subject: [PATCH 2/3] Fix(Test): add rho_stride argument to elecstate_op unit test calls The elecstate_pw_op operator signature gained a rho_stride parameter in bc21c5300, but the unit test still called it with the old signatures, breaking the test build. Pass this->nrxx, which matches the stride of the rho layout used in the test. Co-Authored-By: Claude Opus 4.7 --- .../kernels/test/elecstate_op_test.cpp | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/source/source_estate/kernels/test/elecstate_op_test.cpp b/source/source_estate/kernels/test/elecstate_op_test.cpp index 24f2097cf0..ae441bed7b 100644 --- a/source/source_estate/kernels/test/elecstate_op_test.cpp +++ b/source/source_estate/kernels/test/elecstate_op_test.cpp @@ -63,10 +63,11 @@ TEST_F(TestModuleElecstateMultiDevice, elecstate_pw_op_cpu) double ** rho = new double* [1]; rho[0] = rho_data.data(); elecstate_cpu_op()( - this->cpu_ctx, + this->cpu_ctx, this->spin, this->nrxx, - this->w1, - rho, + this->nrxx, + this->w1, + rho, this->wfcr.data()); // check the result @@ -85,12 +86,13 @@ TEST_F(TestModuleElecstateMultiDevice, elecstate_pw_spin_op_cpu) rho[2] = rho_data.data() + this->nrxx * 2; rho[3] = rho_data.data() + this->nrxx * 3; elecstate_cpu_op()( - this->cpu_ctx, + this->cpu_ctx, this->DOMAG, this->DOMAG_Z, this->nrxx, - this->w1, - rho, + this->nrxx, + this->w1, + rho, this->wfcr_2.data(), this->wfcr_another_spin_2.data()); @@ -114,10 +116,11 @@ TEST_F(TestModuleElecstateMultiDevice, elecstate_pw_op_gpu) double ** rho = new double* [1]; rho[0] = d_rho_data; elecstate_gpu_op()( - this->gpu_ctx, + this->gpu_ctx, this->spin, this->nrxx, + this->nrxx, this->w1, - rho, + rho, d_wfcr); syncmem_var_d2h_op()(rho_data.data(), d_rho_data, rho_data.size()); @@ -149,12 +152,13 @@ TEST_F(TestModuleElecstateMultiDevice, elecstate_pw_spin_op_gpu) rho[3] = d_rho_data_2 + this->nrxx * 3; elecstate_gpu_op()( - this->gpu_ctx, + this->gpu_ctx, this->DOMAG, this->DOMAG_Z, this->nrxx, - this->w1, - rho, + this->nrxx, + this->w1, + rho, d_wfcr_2, d_wfcr_another_spin_2); From 07bc74a570e495052e67a0ae03785b3d2a1662d7 Mon Sep 17 00:00:00 2001 From: chengleizheng Date: Wed, 5 Aug 2026 16:58:20 +0800 Subject: [PATCH 3/3] rename rho_stride to nrxx_dense in elecstate_pw_op --- .../kernels/cuda/elecstate_op.cu | 36 +++++++++---------- source/source_estate/kernels/elecstate_op.cpp | 4 +-- source/source_estate/kernels/elecstate_op.h | 16 ++++----- 3 files changed, 28 insertions(+), 28 deletions(-) diff --git a/source/source_estate/kernels/cuda/elecstate_op.cu b/source/source_estate/kernels/cuda/elecstate_op.cu index 7ffd65e3d4..4e6feedb7e 100644 --- a/source/source_estate/kernels/cuda/elecstate_op.cu +++ b/source/source_estate/kernels/cuda/elecstate_op.cu @@ -15,11 +15,11 @@ __global__ void elecstate_pw( const FPTYPE w1, FPTYPE* rho, const thrust::complex* wfcr, - const int rho_stride) + const int nrxx_dense) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if(idx >= nrxx) {return;} - rho[spin * rho_stride + idx] += w1 * norm(wfcr[idx]); + rho[spin * nrxx_dense + idx] += w1 * norm(wfcr[idx]); } template @@ -31,31 +31,31 @@ __global__ void elecstate_pw( FPTYPE* rho, const thrust::complex* wfcr, const thrust::complex* wfcr_another_spin, - const int rho_stride) + const int nrxx_dense) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if(idx >= nrxx) {return;} - rho[0 * rho_stride + idx] += w1 * (norm(wfcr[idx]) + norm(wfcr_another_spin[idx])); + rho[0 * nrxx_dense + idx] += w1 * (norm(wfcr[idx]) + norm(wfcr_another_spin[idx])); if (DOMAG) { - rho[1 * rho_stride + idx] += w1 * 2.0 + rho[1 * nrxx_dense + idx] += w1 * 2.0 * (wfcr[idx].real() * wfcr_another_spin[idx].real() + wfcr[idx].imag() * wfcr_another_spin[idx].imag()); - rho[2 * rho_stride + idx] += w1 * 2.0 + rho[2 * nrxx_dense + idx] += w1 * 2.0 * (wfcr[idx].real() * wfcr_another_spin[idx].imag() - wfcr_another_spin[idx].real() * wfcr[idx].imag()); - rho[3 * rho_stride + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx])); + rho[3 * nrxx_dense + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx])); } else if(DOMAG_Z) { - rho[1 * rho_stride + idx] = 0; - rho[2 * rho_stride + idx] = 0; - rho[3 * rho_stride + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx])); + rho[1 * nrxx_dense + idx] = 0; + rho[2 * nrxx_dense + idx] = 0; + rho[3 * nrxx_dense + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx])); } else { - rho[0 * rho_stride + idx] = 0; - rho[1 * rho_stride + idx] = 0; - rho[2 * rho_stride + idx] = 0; - rho[3 * rho_stride + idx] = 0; + rho[0 * nrxx_dense + idx] = 0; + rho[1 * nrxx_dense + idx] = 0; + rho[2 * nrxx_dense + idx] = 0; + rho[3 * nrxx_dense + idx] = 0; } } @@ -63,7 +63,7 @@ template void elecstate_pw_op::operator()(const base_device::DEVICE_GPU* ctx, const int& spin, const int& nrxx, - const int& rho_stride, + const int& nrxx_dense, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr) @@ -72,7 +72,7 @@ void elecstate_pw_op::operator()(const base_dev elecstate_pw<<>>( spin, nrxx, w1, rho[0], reinterpret_cast*>(wfcr), - rho_stride + nrxx_dense ); CHECK_CUDA_SYNC(); @@ -83,7 +83,7 @@ void elecstate_pw_op::operator()(const base_dev const bool& DOMAG, const bool& DOMAG_Z, const int& nrxx, - const int& rho_stride, + const int& nrxx_dense, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr, @@ -94,7 +94,7 @@ void elecstate_pw_op::operator()(const base_dev DOMAG, DOMAG_Z, nrxx, w1, rho[0], reinterpret_cast*>(wfcr), reinterpret_cast*>(wfcr_another_spin), - rho_stride + nrxx_dense ); CHECK_CUDA_SYNC(); diff --git a/source/source_estate/kernels/elecstate_op.cpp b/source/source_estate/kernels/elecstate_op.cpp index 9960b44358..14976b10c4 100644 --- a/source/source_estate/kernels/elecstate_op.cpp +++ b/source/source_estate/kernels/elecstate_op.cpp @@ -8,7 +8,7 @@ struct elecstate_pw_op void operator()(const base_device::DEVICE_CPU* /*ctx*/, const int& spin, const int& nrxx, - const int& /*rho_stride*/, + const int& /*nrxx_dense*/, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr) @@ -30,7 +30,7 @@ struct elecstate_pw_op const bool& DOMAG, const bool& DOMAG_Z, const int& nrxx, - const int& /*rho_stride*/, + const int& /*nrxx_dense*/, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr, diff --git a/source/source_estate/kernels/elecstate_op.h b/source/source_estate/kernels/elecstate_op.h index 090b7e00e9..fcd1002cdb 100644 --- a/source/source_estate/kernels/elecstate_op.h +++ b/source/source_estate/kernels/elecstate_op.h @@ -14,8 +14,8 @@ struct elecstate_pw_op { /// Input Parameters /// @param ctx - which device this function runs on /// @param spin - current spin - /// @param nrxx - number of planewaves - /// @param rho_stride - stride between spin components of rho + /// @param nrxx - number of real-space grid points on this process (basis/wfc grid) + /// @param nrxx_dense - nrxx of the dense (charge) grid, used as the stride between spin components of rho /// @param weight - input constant /// @param wfcr - input array, psi in real space /// @@ -25,7 +25,7 @@ struct elecstate_pw_op { const Device* ctx, const int& spin, const int& nrxx, - const int& rho_stride, + const int& nrxx_dense, const FPTYPE& weight, FPTYPE** rho, const std::complex* wfcr); @@ -36,8 +36,8 @@ struct elecstate_pw_op { /// @param ctx - which device this function runs on /// @param DOMAG - PARAM.globalv.domag /// @param DOMAG_Z - PARAM.globalv.domag_z - /// @param nrxx - number of planewaves - /// @param rho_stride - stride between spin components of rho + /// @param nrxx - number of real-space grid points on this process (basis/wfc grid) + /// @param nrxx_dense - nrxx of the dense (charge) grid, used as the stride between spin components of rho /// @param weight - input constant /// @param wfcr - input array, psi in real space /// @param wfcr_another_spin - input array, psi in real space @@ -49,7 +49,7 @@ struct elecstate_pw_op { const bool& DOMAG, const bool& DOMAG_Z, const int& nrxx, - const int& rho_stride, + const int& nrxx_dense, const FPTYPE& weight, FPTYPE** rho, const std::complex* wfcr, @@ -63,7 +63,7 @@ struct elecstate_pw_op void operator()(const base_device::DEVICE_GPU* ctx, const int& spin, const int& nrxx, - const int& rho_stride, + const int& nrxx_dense, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr); @@ -72,7 +72,7 @@ struct elecstate_pw_op const bool& DOMAG, const bool& DOMAG_Z, const int& nrxx, - const int& rho_stride, + const int& nrxx_dense, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr,