diff --git a/source/source_estate/elecstate_pw.cpp b/source/source_estate/elecstate_pw.cpp index 641def5dc4..9ff18d7c1d 100644 --- a/source/source_estate/elecstate_pw.cpp +++ b/source/source_estate/elecstate_pw.cpp @@ -227,6 +227,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, @@ -249,7 +250,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 @@ -272,7 +273,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..4e6feedb7e 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 nrxx_dense) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if(idx >= nrxx) {return;} - rho[spin * nrxx + idx] += w1 * norm(wfcr[idx]); + rho[spin * nrxx_dense + 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 nrxx_dense) { 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 * nrxx_dense + idx] += w1 * (norm(wfcr[idx]) + norm(wfcr_another_spin[idx])); if (DOMAG) { - rho[1 * nrxx + 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 * nrxx + 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 * nrxx + 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 * nrxx + idx] = 0; - rho[2 * nrxx + idx] = 0; - rho[3 * nrxx + 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 * nrxx + idx] = 0; - rho[1 * nrxx + idx] = 0; - rho[2 * nrxx + idx] = 0; - rho[3 * nrxx + 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; } } @@ -61,6 +63,7 @@ template void elecstate_pw_op::operator()(const base_device::DEVICE_GPU* ctx, const int& spin, const int& nrxx, + const int& nrxx_dense, 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), + nrxx_dense ); 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& nrxx_dense, 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), + nrxx_dense ); CHECK_CUDA_SYNC(); diff --git a/source/source_estate/kernels/elecstate_op.cpp b/source/source_estate/kernels/elecstate_op.cpp index c3933319fc..14976b10c4 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& /*nrxx_dense*/, 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& /*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 f7b8d48c2e..fcd1002cdb 100644 --- a/source/source_estate/kernels/elecstate_op.h +++ b/source/source_estate/kernels/elecstate_op.h @@ -14,7 +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 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 /// @@ -24,6 +25,7 @@ struct elecstate_pw_op { const Device* ctx, const int& spin, const int& nrxx, + const int& nrxx_dense, const FPTYPE& weight, FPTYPE** rho, const std::complex* wfcr); @@ -34,7 +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 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 @@ -46,6 +49,7 @@ struct elecstate_pw_op { const bool& DOMAG, const bool& DOMAG_Z, const int& nrxx, + const int& nrxx_dense, 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& nrxx_dense, 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& nrxx_dense, const FPTYPE& w1, FPTYPE** rho, const std::complex* wfcr, 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); 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];