Skip to content

Commit a0ad57a

Browse files
author
chengleizheng
committed
Refactor: rename device/host pointers in USPP PW code
1 parent 5b44ef2 commit a0ad57a

2 files changed

Lines changed: 52 additions & 53 deletions

File tree

‎source/source_estate/elecstate_pw.cpp‎

Lines changed: 48 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -282,13 +282,13 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
282282
resmem_var_h_op()(becsum, nh_tot * ucell->nat * PARAM.inp.nspin, "ElecState<PW>::becsum");
283283
setmem_var_h_op()(becsum, 0, nh_tot * ucell->nat * PARAM.inp.nspin);
284284

285-
// becp: GPU for first gemm, then D2H for CPU loops
286-
T* becp_gpu = nullptr;
287-
std::vector<T> becp;
285+
// becp: device buffer for gemm, then D2H for host loops
286+
T* becp = nullptr;
287+
std::vector<T> becp_host;
288288
if (nkb > 0)
289289
{
290-
resmem_complex_op()(becp_gpu, nbands * nkb, "ElecState<PW>::becp");
291-
becp.resize(nbands * nkb);
290+
resmem_complex_op()(becp, nbands * nkb, "ElecState<PW>::becp");
291+
becp_host.resize(nbands * nkb);
292292
}
293293

294294
for (int ik = 0; ik < psi.get_nk(); ++ik)
@@ -304,7 +304,7 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
304304
this->ppcell->getvnl(this->ctx, *ucell,ik, this->vkb);
305305
}
306306

307-
// becp = <beta|psi> (GPU gemm)
307+
// becp = <beta|psi> (device gemm)
308308
char transa = 'C';
309309
char transb = 'N';
310310
if (this->ppcell->nkb > 0)
@@ -321,7 +321,7 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
321321
psi_now,
322322
inc,
323323
&zero,
324-
becp_gpu,
324+
becp,
325325
inc);
326326
}
327327
else
@@ -337,13 +337,13 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
337337
psi_now,
338338
npwx,
339339
&zero,
340-
becp_gpu,
340+
becp,
341341
this->ppcell->nkb);
342342
}
343-
// D2H: GPU becp → CPU becp
344-
syncmem_complex_d2h_op()(becp.data(), becp_gpu, nbands * nkb);
343+
// D2H: device becp → host becp_host
344+
syncmem_complex_d2h_op()(becp_host.data(), becp, nbands * nkb);
345345
}
346-
Parallel_Reduce::reduce_pool(becp.data(), this->ppcell->nkb * nbands);
346+
Parallel_Reduce::reduce_pool(becp_host.data(), this->ppcell->nkb * nbands);
347347

348348
// sum over bands: \sum_i <psi_i|beta_l><beta_m|psi_i> w_i
349349
for (int it = 0; it < ucell->ntype; it++)
@@ -352,9 +352,14 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
352352
if (atom->ncpp.tvanp)
353353
{
354354
const int nh_atom = atom->ncpp.nh;
355-
// auxk1, auxk2 on CPU for filling loops
356-
std::vector<T> auxk1(nbands * nh_atom);
357-
std::vector<T> auxk2(nbands * nh_atom);
355+
// auxk1, auxk2: host buffers for filling loops
356+
std::vector<T> auxk1_host(nbands * nh_atom);
357+
std::vector<T> auxk2_host(nbands * nh_atom);
358+
// device buffers allocated once per atom type
359+
T *aux_gk = nullptr, *auxk1 = nullptr, *auxk2 = nullptr;
360+
resmem_complex_op()(auxk1, nbands * nh_atom, "ElecState<PW>::auxk1");
361+
resmem_complex_op()(auxk2, nbands * nh_atom, "ElecState<PW>::auxk2");
362+
resmem_complex_op()(aux_gk, nh_atom * nh_atom * npol * npol, "ElecState<PW>::aux_gk");
358363
for (int ia = 0; ia < atom->na; ia++)
359364
{
360365
const int iat = ucell->itia2iat(it, ia);
@@ -369,19 +374,15 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
369374
const int ikb = this->ppcell->indv_ijkb0[iat] + ih;
370375
for (int ib = 0; ib < nbands; ib++)
371376
{
372-
auxk1[ih * nbands + ib] = becp[ib * this->ppcell->nkb + ikb];
373-
auxk2[ih * nbands + ib]
374-
= becp[ib * this->ppcell->nkb + ikb] * static_cast<Real>(this->wg(ik, ib));
377+
auxk1_host[ih * nbands + ib] = becp_host[ib * this->ppcell->nkb + ikb];
378+
auxk2_host[ih * nbands + ib]
379+
= becp_host[ib * this->ppcell->nkb + ikb] * static_cast<Real>(this->wg(ik, ib));
375380
}
376381
}
377382

378-
// GPU gemm: aux_gk = auxk1^H * auxk2
379-
T *aux_gk_gpu = nullptr, *auxk1_gpu = nullptr, *auxk2_gpu = nullptr;
380-
resmem_complex_op()(auxk1_gpu, nbands * nh_atom, "ElecState<PW>::auxk1");
381-
resmem_complex_op()(auxk2_gpu, nbands * nh_atom, "ElecState<PW>::auxk2");
382-
resmem_complex_op()(aux_gk_gpu, nh_atom * nh_atom * npol * npol, "ElecState<PW>::aux_gk");
383-
syncmem_complex_h2d_op()(auxk1_gpu, auxk1.data(), nbands * nh_atom);
384-
syncmem_complex_h2d_op()(auxk2_gpu, auxk2.data(), nbands * nh_atom);
383+
// device gemm: aux_gk = auxk1^H * auxk2
384+
syncmem_complex_h2d_op()(auxk1, auxk1_host.data(), nbands * nh_atom);
385+
syncmem_complex_h2d_op()(auxk2, auxk2_host.data(), nbands * nh_atom);
385386

386387
char transa2 = 'C';
387388
char transb2 = 'N';
@@ -391,21 +392,17 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
391392
nh_atom,
392393
nbands,
393394
&one,
394-
auxk1_gpu,
395+
auxk1,
395396
nbands,
396-
auxk2_gpu,
397+
auxk2,
397398
nbands,
398399
&zero,
399-
aux_gk_gpu,
400+
aux_gk,
400401
nh_atom);
401402

402-
// D2H: GPU aux_gk → CPU
403-
std::vector<T> aux_gk(nh_atom * nh_atom);
404-
syncmem_complex_d2h_op()(aux_gk.data(), aux_gk_gpu, nh_atom * nh_atom);
405-
406-
delmem_complex_op()(auxk1_gpu);
407-
delmem_complex_op()(auxk2_gpu);
408-
delmem_complex_op()(aux_gk_gpu);
403+
// D2H: device aux_gk → host
404+
std::vector<T> aux_gk_host(nh_atom * nh_atom);
405+
syncmem_complex_d2h_op()(aux_gk_host.data(), aux_gk, nh_atom * nh_atom);
409406

410407
// copy output from GEMM into desired format
411408
int ijh = 0;
@@ -416,21 +413,24 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
416413
{
417414
if (ih == jh)
418415
{
419-
becsum[index + ijh] += std::real(aux_gk[ih * nh_atom + jh]);
416+
becsum[index + ijh] += std::real(aux_gk_host[ih * nh_atom + jh]);
420417
}
421418
else
422419
{
423-
becsum[index + ijh] += 2.0 * std::real(aux_gk[ih * nh_atom + jh]);
420+
becsum[index + ijh] += 2.0 * std::real(aux_gk_host[ih * nh_atom + jh]);
424421
}
425422
ijh++;
426423
}
427424
}
428425
}
429426
}
427+
delmem_complex_op()(auxk1);
428+
delmem_complex_op()(auxk2);
429+
delmem_complex_op()(aux_gk);
430430
}
431431
}
432432
}
433-
if (becp_gpu != nullptr) delmem_complex_op()(becp_gpu);
433+
delmem_complex_op()(becp);
434434
}
435435

436436
template <typename T, typename Device>
@@ -454,7 +454,7 @@ void ElecStatePW<T, Device>::add_usrho(const psi::Psi<T, Device>& psi)
454454
syncmem_complex_d2h_op()(rhog_host.data(), this->rhog[is], this->rhopw_smooth->npw);
455455
// CPU real2recip
456456
this->rhopw_smooth->real2recip(rho_host.data(), rhog_host.data());
457-
// H2D rhog back to GPU
457+
// H2D rhog back to device
458458
syncmem_complex_h2d_op()(this->rhog[is], rhog_host.data(), this->rhopw_smooth->npw);
459459
}
460460
}
@@ -503,16 +503,16 @@ void ElecStatePW<T, Device>::addusdens_g(const Real* becsum, T** rhog)
503503
qmod_host[ig] = static_cast<double>(this->charge->rhopw->gcar[ig].norm() * ucell->tpiba);
504504
}
505505

506-
// ylmk0: compute on GPU then D2H
507-
Real* ylmk0_gpu = nullptr;
508-
resmem_var_op()(ylmk0_gpu, npw * lmaxq * lmaxq, "ElecState<PW>::ylmk0");
509-
Real* g_gpu = nullptr;
510-
resmem_var_op()(g_gpu, npw * 3, "ElecState<PW>::g_gpu");
511-
syncmem_var_h2d_op()(g_gpu, reinterpret_cast<Real*>(this->charge->rhopw->gcar), npw * 3);
512-
ModuleBase::YlmReal::Ylm_Real(this->ctx, lmaxq * lmaxq, npw, g_gpu, ylmk0_gpu);
513-
delmem_var_op()(g_gpu);
506+
// ylmk0: compute on device then D2H
507+
Real* ylmk0 = nullptr;
508+
resmem_var_op()(ylmk0, npw * lmaxq * lmaxq, "ElecState<PW>::ylmk0");
509+
Real* g = nullptr;
510+
resmem_var_op()(g, npw * 3, "ElecState<PW>::g");
511+
syncmem_var_h2d_op()(g, reinterpret_cast<Real*>(this->charge->rhopw->gcar), npw * 3);
512+
ModuleBase::YlmReal::Ylm_Real(this->ctx, lmaxq * lmaxq, npw, g, ylmk0);
513+
delmem_var_op()(g);
514514
std::vector<Real> ylmk0_host(npw * lmaxq * lmaxq);
515-
syncmem_var_d2h_op()(ylmk0_host.data(), ylmk0_gpu, npw * lmaxq * lmaxq);
515+
syncmem_var_d2h_op()(ylmk0_host.data(), ylmk0, npw * lmaxq * lmaxq);
516516
std::vector<double> ylmk0_double(npw * lmaxq * lmaxq);
517517
for (int i = 0; i < npw * lmaxq * lmaxq; i++) ylmk0_double[i] = static_cast<double>(ylmk0_host[i]);
518518

@@ -581,7 +581,7 @@ void ElecStatePW<T, Device>::addusdens_g(const Real* becsum, T** rhog)
581581
}
582582
}
583583

584-
delmem_var_op()(ylmk0_gpu);
584+
delmem_var_op()(ylmk0);
585585
}
586586

587587
template class ElecStatePW<std::complex<float>, base_device::DEVICE_CPU>;

‎source/source_pw/module_pwdft/hamilt_pw.cpp‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -252,21 +252,20 @@ void HamiltPW<T, Device>::sPsi(const T* psi_in, // psi
252252
const int nh = atoms->ncpp.nh;
253253
T* qqc = nullptr;
254254
resmem_complex_op()(qqc, nh * nh, "Hamilt<PW>::qqc");
255-
256255
std::vector<T> qqc_host(nh*nh);
257256
const double* qq_now_host = &this->ppcell->qq_nt.ptr[it * this->ppcell->nhm * this->ppcell->nhm];
258257

259258
for (int i = 0; i < nh; i++)
260259
{
261260
for (int j = 0; j < nh; j++)
262261
{
263-
const int source_index = i * this->ppcell->nhm + j;
264-
const int target_index = i * nh + j;
262+
const int source_index = i * this->ppcell->nhm + j;
263+
const int target_index = i * nh + j;
265264

266-
qqc_host[target_index] = static_cast<Real>(qq_now_host[source_index]) * one;
265+
qqc_host[target_index] = static_cast<Real>(qq_now_host[source_index]) * one;
267266
}
268267
}
269-
268+
270269
syncmem_complex_h2d_op()(qqc, qqc_host.data(), qqc_host.size());
271270

272271
for (int ia = 0; ia < atoms->na; ia++)

0 commit comments

Comments
 (0)