Skip to content

Commit b1141ba

Browse files
chengleizhengchengleizhengclaude
authored
Fix/gpu elecstate pw op spin stride (#7776)
* 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. * Fix(Test): add rho_stride argument to elecstate_op unit test calls The elecstate_pw_op operator signature gained a rho_stride parameter in bc21c53, 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 <noreply@anthropic.com> * rename rho_stride to nrxx_dense in elecstate_pw_op --------- Co-authored-by: chengleizheng <you@example.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
1 parent ff23741 commit b1141ba

8 files changed

Lines changed: 54 additions & 34 deletions

File tree

‎source/source_estate/elecstate_pw.cpp‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,7 @@ void ElecStatePW<T, Device>::rhoBandK(const psi::Psi<T, Device>& psi)
227227
PARAM.globalv.domag,
228228
PARAM.globalv.domag_z,
229229
this->basis->nrxx,
230+
this->charge->nrxx,
230231
w1,
231232
this->rho,
232233
this->wfcr,
@@ -249,7 +250,7 @@ void ElecStatePW<T, Device>::rhoBandK(const psi::Psi<T, Device>& psi)
249250
if (w1 != 0.0)
250251
{
251252
// replaced by denghui at 20221110
252-
elecstate_pw_op()(this->ctx, current_spin, this->basis->nrxx, w1, this->rho, this->wfcr);
253+
elecstate_pw_op()(this->ctx, current_spin, this->basis->nrxx, this->charge->nrxx, w1, this->rho, this->wfcr);
253254
}
254255

255256
// kinetic energy density
@@ -272,7 +273,7 @@ void ElecStatePW<T, Device>::rhoBandK(const psi::Psi<T, Device>& psi)
272273

273274
this->basis->recip_to_real(this->ctx, this->wfcr, this->wfcr, ik);
274275

275-
elecstate_pw_op()(this->ctx, current_spin, this->charge->nrxx, w1, this->kin_r, this->wfcr);
276+
elecstate_pw_op()(this->ctx, current_spin, this->charge->nrxx, this->charge->nrxx, w1, this->kin_r, this->wfcr);
276277
}
277278
}
278279
}

‎source/source_estate/elecstate_pw_cal_tau.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ void ElecStatePW<T, Device>::cal_tau(const psi::Psi<T, Device>& psi)
4545

4646
this->basis->recip_to_real(this->ctx, this->wfcr, this->wfcr, ik);
4747

48-
elecstate_pw_op()(this->ctx, current_spin, this->charge->nrxx, w1, this->kin_r, this->wfcr);
48+
elecstate_pw_op()(this->ctx, current_spin, this->charge->nrxx, this->charge->nrxx, w1, this->kin_r, this->wfcr);
4949
}
5050
}
5151
}

‎source/source_estate/kernels/cuda/elecstate_op.cu‎

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,12 @@ __global__ void elecstate_pw(
1414
const int nrxx,
1515
const FPTYPE w1,
1616
FPTYPE* rho,
17-
const thrust::complex<FPTYPE>* wfcr)
17+
const thrust::complex<FPTYPE>* wfcr,
18+
const int nrxx_dense)
1819
{
1920
int idx = blockIdx.x * blockDim.x + threadIdx.x;
2021
if(idx >= nrxx) {return;}
21-
rho[spin * nrxx + idx] += w1 * norm(wfcr[idx]);
22+
rho[spin * nrxx_dense + idx] += w1 * norm(wfcr[idx]);
2223
}
2324

2425
template<typename FPTYPE>
@@ -29,46 +30,49 @@ __global__ void elecstate_pw(
2930
const FPTYPE w1,
3031
FPTYPE* rho,
3132
const thrust::complex<FPTYPE>* wfcr,
32-
const thrust::complex<FPTYPE>* wfcr_another_spin)
33+
const thrust::complex<FPTYPE>* wfcr_another_spin,
34+
const int nrxx_dense)
3335
{
3436
int idx = blockIdx.x * blockDim.x + threadIdx.x;
3537
if(idx >= nrxx) {return;}
36-
rho[0 * nrxx + idx] += w1 * (norm(wfcr[idx]) + norm(wfcr_another_spin[idx]));
38+
rho[0 * nrxx_dense + idx] += w1 * (norm(wfcr[idx]) + norm(wfcr_another_spin[idx]));
3739

3840
if (DOMAG) {
39-
rho[1 * nrxx + idx] += w1 * 2.0
41+
rho[1 * nrxx_dense + idx] += w1 * 2.0
4042
* (wfcr[idx].real() * wfcr_another_spin[idx].real()
4143
+ wfcr[idx].imag() * wfcr_another_spin[idx].imag());
42-
rho[2 * nrxx + idx] += w1 * 2.0
44+
rho[2 * nrxx_dense + idx] += w1 * 2.0
4345
* (wfcr[idx].real() * wfcr_another_spin[idx].imag()
4446
- wfcr_another_spin[idx].real() * wfcr[idx].imag());
45-
rho[3 * nrxx + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx]));
47+
rho[3 * nrxx_dense + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx]));
4648
}
4749
else if(DOMAG_Z) {
48-
rho[1 * nrxx + idx] = 0;
49-
rho[2 * nrxx + idx] = 0;
50-
rho[3 * nrxx + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx]));
50+
rho[1 * nrxx_dense + idx] = 0;
51+
rho[2 * nrxx_dense + idx] = 0;
52+
rho[3 * nrxx_dense + idx] += w1 * (norm(wfcr[idx]) - norm(wfcr_another_spin[idx]));
5153
}
5254
else {
53-
rho[0 * nrxx + idx] = 0;
54-
rho[1 * nrxx + idx] = 0;
55-
rho[2 * nrxx + idx] = 0;
56-
rho[3 * nrxx + idx] = 0;
55+
rho[0 * nrxx_dense + idx] = 0;
56+
rho[1 * nrxx_dense + idx] = 0;
57+
rho[2 * nrxx_dense + idx] = 0;
58+
rho[3 * nrxx_dense + idx] = 0;
5759
}
5860
}
5961

6062
template <typename FPTYPE>
6163
void elecstate_pw_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const base_device::DEVICE_GPU* ctx,
6264
const int& spin,
6365
const int& nrxx,
66+
const int& nrxx_dense,
6467
const FPTYPE& w1,
6568
FPTYPE** rho,
6669
const std::complex<FPTYPE>* wfcr)
6770
{
6871
const int block = (nrxx + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
6972
elecstate_pw<FPTYPE><<<block, THREADS_PER_BLOCK>>>(
7073
spin, nrxx, w1, rho[0],
71-
reinterpret_cast<const thrust::complex<FPTYPE>*>(wfcr)
74+
reinterpret_cast<const thrust::complex<FPTYPE>*>(wfcr),
75+
nrxx_dense
7276
);
7377

7478
CHECK_CUDA_SYNC();
@@ -79,6 +83,7 @@ void elecstate_pw_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const base_dev
7983
const bool& DOMAG,
8084
const bool& DOMAG_Z,
8185
const int& nrxx,
86+
const int& nrxx_dense,
8287
const FPTYPE& w1,
8388
FPTYPE** rho,
8489
const std::complex<FPTYPE>* wfcr,
@@ -88,7 +93,8 @@ void elecstate_pw_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const base_dev
8893
elecstate_pw<FPTYPE><<<block, THREADS_PER_BLOCK>>>(
8994
DOMAG, DOMAG_Z, nrxx, w1, rho[0],
9095
reinterpret_cast<const thrust::complex<FPTYPE>*>(wfcr),
91-
reinterpret_cast<const thrust::complex<FPTYPE>*>(wfcr_another_spin)
96+
reinterpret_cast<const thrust::complex<FPTYPE>*>(wfcr_another_spin),
97+
nrxx_dense
9298
);
9399

94100
CHECK_CUDA_SYNC();

‎source/source_estate/kernels/elecstate_op.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ struct elecstate_pw_op<FPTYPE, base_device::DEVICE_CPU>
88
void operator()(const base_device::DEVICE_CPU* /*ctx*/,
99
const int& spin,
1010
const int& nrxx,
11+
const int& /*nrxx_dense*/,
1112
const FPTYPE& w1,
1213
FPTYPE** rho,
1314
const std::complex<FPTYPE>* wfcr)
@@ -29,6 +30,7 @@ struct elecstate_pw_op<FPTYPE, base_device::DEVICE_CPU>
2930
const bool& DOMAG,
3031
const bool& DOMAG_Z,
3132
const int& nrxx,
33+
const int& /*nrxx_dense*/,
3234
const FPTYPE& w1,
3335
FPTYPE** rho,
3436
const std::complex<FPTYPE>* wfcr,

‎source/source_estate/kernels/elecstate_op.h‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,8 @@ struct elecstate_pw_op {
1414
/// Input Parameters
1515
/// @param ctx - which device this function runs on
1616
/// @param spin - current spin
17-
/// @param nrxx - number of planewaves
17+
/// @param nrxx - number of real-space grid points on this process (basis/wfc grid)
18+
/// @param nrxx_dense - nrxx of the dense (charge) grid, used as the stride between spin components of rho
1819
/// @param weight - input constant
1920
/// @param wfcr - input array, psi in real space
2021
///
@@ -24,6 +25,7 @@ struct elecstate_pw_op {
2425
const Device* ctx,
2526
const int& spin,
2627
const int& nrxx,
28+
const int& nrxx_dense,
2729
const FPTYPE& weight,
2830
FPTYPE** rho,
2931
const std::complex<FPTYPE>* wfcr);
@@ -34,7 +36,8 @@ struct elecstate_pw_op {
3436
/// @param ctx - which device this function runs on
3537
/// @param DOMAG - PARAM.globalv.domag
3638
/// @param DOMAG_Z - PARAM.globalv.domag_z
37-
/// @param nrxx - number of planewaves
39+
/// @param nrxx - number of real-space grid points on this process (basis/wfc grid)
40+
/// @param nrxx_dense - nrxx of the dense (charge) grid, used as the stride between spin components of rho
3841
/// @param weight - input constant
3942
/// @param wfcr - input array, psi in real space
4043
/// @param wfcr_another_spin - input array, psi in real space
@@ -46,6 +49,7 @@ struct elecstate_pw_op {
4649
const bool& DOMAG,
4750
const bool& DOMAG_Z,
4851
const int& nrxx,
52+
const int& nrxx_dense,
4953
const FPTYPE& weight,
5054
FPTYPE** rho,
5155
const std::complex<FPTYPE>* wfcr,
@@ -59,6 +63,7 @@ struct elecstate_pw_op<FPTYPE, base_device::DEVICE_GPU>
5963
void operator()(const base_device::DEVICE_GPU* ctx,
6064
const int& spin,
6165
const int& nrxx,
66+
const int& nrxx_dense,
6267
const FPTYPE& w1,
6368
FPTYPE** rho,
6469
const std::complex<FPTYPE>* wfcr);
@@ -67,6 +72,7 @@ struct elecstate_pw_op<FPTYPE, base_device::DEVICE_GPU>
6772
const bool& DOMAG,
6873
const bool& DOMAG_Z,
6974
const int& nrxx,
75+
const int& nrxx_dense,
7076
const FPTYPE& w1,
7177
FPTYPE** rho,
7278
const std::complex<FPTYPE>* wfcr,

‎source/source_estate/kernels/test/elecstate_op_test.cpp‎

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -63,10 +63,11 @@ TEST_F(TestModuleElecstateMultiDevice, elecstate_pw_op_cpu)
6363
double ** rho = new double* [1];
6464
rho[0] = rho_data.data();
6565
elecstate_cpu_op()(
66-
this->cpu_ctx,
66+
this->cpu_ctx,
6767
this->spin, this->nrxx,
68-
this->w1,
69-
rho,
68+
this->nrxx,
69+
this->w1,
70+
rho,
7071
this->wfcr.data());
7172

7273
// check the result
@@ -85,12 +86,13 @@ TEST_F(TestModuleElecstateMultiDevice, elecstate_pw_spin_op_cpu)
8586
rho[2] = rho_data.data() + this->nrxx * 2;
8687
rho[3] = rho_data.data() + this->nrxx * 3;
8788
elecstate_cpu_op()(
88-
this->cpu_ctx,
89+
this->cpu_ctx,
8990
this->DOMAG,
9091
this->DOMAG_Z,
9192
this->nrxx,
92-
this->w1,
93-
rho,
93+
this->nrxx,
94+
this->w1,
95+
rho,
9496
this->wfcr_2.data(),
9597
this->wfcr_another_spin_2.data());
9698

@@ -114,10 +116,11 @@ TEST_F(TestModuleElecstateMultiDevice, elecstate_pw_op_gpu)
114116
double ** rho = new double* [1];
115117
rho[0] = d_rho_data;
116118
elecstate_gpu_op()(
117-
this->gpu_ctx,
119+
this->gpu_ctx,
118120
this->spin, this->nrxx,
121+
this->nrxx,
119122
this->w1,
120-
rho,
123+
rho,
121124
d_wfcr);
122125

123126
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)
149152
rho[3] = d_rho_data_2 + this->nrxx * 3;
150153

151154
elecstate_gpu_op()(
152-
this->gpu_ctx,
155+
this->gpu_ctx,
153156
this->DOMAG,
154157
this->DOMAG_Z,
155158
this->nrxx,
156-
this->w1,
157-
rho,
159+
this->nrxx,
160+
this->w1,
161+
rho,
158162
d_wfcr_2,
159163
d_wfcr_another_spin_2);
160164

‎source/source_io/module_wf/read_wf2rho_pw.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,7 @@ void ModuleIO::read_wf2rho_pw(
139139
PARAM.globalv.domag,
140140
PARAM.globalv.domag_z,
141141
nrxx,
142+
nrxx,
142143
w1,
143144
chg.rho,
144145
rho_tmp.data(),
@@ -158,7 +159,7 @@ void ModuleIO::read_wf2rho_pw(
158159
if (w1 != 0.0)
159160
{
160161
base_device::DEVICE_CPU* ctx = nullptr;
161-
elecstate::elecstate_pw_op<double, base_device::DEVICE_CPU>()(ctx, is, nrxx,
162+
elecstate::elecstate_pw_op<double, base_device::DEVICE_CPU>()(ctx, is, nrxx, nrxx,
162163
w1, chg.rho, rho_tmp.data());
163164
}
164165
}

‎source/source_pw/module_stodft/sto_iter.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -645,7 +645,7 @@ void Stochastic_Iter<T, Device>::cal_storho(const UnitCell& ucell,
645645
{
646646
wfc_basis->recip_to_real(this->ctx, tmpout, porter, ik);
647647
const auto w1 = static_cast<Real>(this->pkv->wk[ik]);
648-
elecstate::elecstate_pw_op<Real, Device>()(this->ctx, current_spin, nrxx, w1, pes->rho, porter);
648+
elecstate::elecstate_pw_op<Real, Device>()(this->ctx, current_spin, nrxx, pes->charge->nrxx, w1, pes->rho, porter);
649649
// for (int ir = 0; ir < nrxx; ++ir)
650650
// {
651651
// pes->charge->rho[0][ir] += norm(porter[ir]) * this->pkv->wk[ik];

0 commit comments

Comments
 (0)