Skip to content

Commit cda0a73

Browse files
chengleizhengchengleizhengclaude
authored
fix(gpu): avoid host-side FFT crash in LDOS output and wfc-read charge density (deepmodeling#7847)
* fix(gpu): avoid host-side FFT crash in LDOS output and wfc-read charge density On GPU builds the host-side recip2real/real2recip (fftw path) returns nullptr aux buffers, causing segfaults in cal_ldos_pw (out_ldos) and read_wf2rho_pw (init_chg=wfc). Route these calls through the recip_to_real<Device> dispatch with device-side buffers and explicit CPU<->device synchronization. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * update CASES_GPU * replace the unsafe host-side FFTW path with device-aware FFT dispatch and explicit host-device synchronization. --------- Co-authored-by: chengleizheng <you@example.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
1 parent 75145da commit cda0a73

6 files changed

Lines changed: 256 additions & 31 deletions

File tree

‎source/source_io/module_ctrl/ctrl_output_pw.cpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -249,12 +249,16 @@ void ModuleIO::ctrl_runner_pw(UnitCell& ucell,
249249
ModuleBase::TITLE("ModuleIO", "ctrl_runner_pw");
250250
ModuleBase::timer::start("ModuleIO", "ctrl_runner_pw");
251251

252+
// Create local ctx for device type deduction
253+
Device* ctx = nullptr;
254+
252255
//----------------------------------------------------------
253256
//! 1) Compute LDOS
254257
//----------------------------------------------------------
255258
if (inp.out_ldos[0])
256259
{
257-
ModuleIO::cal_ldos_pw(reinterpret_cast<elecstate::ElecStatePW<std::complex<double>>*>(pelec), stp.psi_cpu[0], para_grid, ucell);
260+
stp.update_psi_d();
261+
ModuleIO::cal_ldos_pw(reinterpret_cast<elecstate::ElecStatePW<std::complex<double>>*>(pelec), *stp.template get_psi_d<T, Device>(), ctx, para_grid, ucell);
258262
}
259263

260264
//----------------------------------------------------------

‎source/source_io/module_dos/cal_ldos.cpp‎

Lines changed: 76 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include "../module_output/cube_io.h"
55
#include "source_estate/module_dm/cal_dm_psi.h"
66
#include "source_hamilt/module_gint/gint_interface.h"
7+
#include "source_base/module_device/memory_op.h"
78

89
#include <type_traits>
910

@@ -105,23 +106,27 @@ template class Cal_ldos<double>; // Gamma_only case
105106
template class Cal_ldos<std::complex<double>>; // multi-k case
106107

107108
// pw case
109+
template <typename Device>
108110
void cal_ldos_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
109-
const psi::Psi<std::complex<double>>& psi,
111+
const psi::Psi<std::complex<double>, Device>& psi,
112+
const Device* ctx,
110113
const Parallel_Grid& pgrid,
111114
const UnitCell& ucell)
112115
{
113116
if (PARAM.inp.out_ldos[0] == 1 || PARAM.inp.out_ldos[0] == 3)
114117
{
115-
ModuleIO::stm_mode_pw(pelec, psi, pgrid, ucell);
118+
ModuleIO::stm_mode_pw<Device>(pelec, psi, ctx, pgrid, ucell);
116119
}
117120
if (PARAM.inp.out_ldos[0] == 2 || PARAM.inp.out_ldos[0] == 3)
118121
{
119-
ModuleIO::ldos_mode_pw(pelec, psi, pgrid, ucell);
122+
ModuleIO::ldos_mode_pw<Device>(pelec, psi, ctx, pgrid, ucell);
120123
}
121124
}
122125

126+
template <typename Device>
123127
void stm_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
124-
const psi::Psi<std::complex<double>>& psi,
128+
const psi::Psi<std::complex<double>, Device>& psi,
129+
const Device* ctx,
125130
const Parallel_Grid& pgrid,
126131
const UnitCell& ucell)
127132
{
@@ -135,6 +140,13 @@ void stm_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
135140
std::vector<double> ldos(pelec->charge->nrxx);
136141
std::vector<std::complex<double>> wfcr(pelec->basis->nrxx);
137142

143+
// Allocate device memory
144+
std::complex<double>* wfcr_dev = nullptr;
145+
if (!std::is_same<Device, base_device::DEVICE_CPU>::value)
146+
{
147+
base_device::memory::resize_memory_op<std::complex<double>, Device>()(wfcr_dev, pelec->basis->nrxx);
148+
}
149+
138150
for (int ik = 0; ik < pelec->klist->get_nks(); ++ik)
139151
{
140152
psi.fix_k(ik);
@@ -143,7 +155,18 @@ void stm_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
143155

144156
for (int ib = 0; ib < nbands; ib++)
145157
{
146-
pelec->basis->recip2real(&psi(ib, 0), wfcr.data(), ik);
158+
// FFT on device and copy result back to host
159+
if (std::is_same<Device, base_device::DEVICE_CPU>::value)
160+
{
161+
pelec->basis->recip2real(&psi(ib, 0), wfcr.data(), ik);
162+
}
163+
else
164+
{
165+
pelec->basis->recip_to_real(ctx, &psi(ib, 0), wfcr_dev, ik);
166+
167+
base_device::memory::synchronize_memory_op<std::complex<double>, base_device::DEVICE_CPU, Device>()(
168+
wfcr.data(), wfcr_dev, pelec->basis->nrxx);
169+
}
147170

148171
const double eigenval = (pelec->ekb(ik, ib) - efermi) * ModuleBase::Ry_to_eV;
149172
double weight = en > 0 ? pelec->klist->wk[ik] - pelec->wg(ik, ib) : pelec->wg(ik, ib);
@@ -159,6 +182,12 @@ void stm_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
159182
}
160183
}
161184

185+
// Free device memory
186+
if (!std::is_same<Device, base_device::DEVICE_CPU>::value)
187+
{
188+
base_device::memory::delete_memory_op<std::complex<double>, Device>()(wfcr_dev);
189+
}
190+
162191
std::stringstream fn;
163192
fn << PARAM.globalv.global_out_dir << "LDOS_" << en << "eV"
164193
<< ".cube";
@@ -168,8 +197,10 @@ void stm_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
168197
}
169198
}
170199

200+
template <typename Device>
171201
void ldos_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
172-
const psi::Psi<std::complex<double>>& psi,
202+
const psi::Psi<std::complex<double>, Device>& psi,
203+
const Device* ctx,
173204
const Parallel_Grid& pgrid,
174205
const UnitCell& ucell)
175206
{
@@ -205,6 +236,14 @@ void ldos_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
205236
// calculate ldos
206237
std::vector<double> tmp(pelec->charge->nrxx);
207238
std::vector<std::complex<double>> wfcr(pelec->basis->nrxx);
239+
240+
// Allocate device memory
241+
std::complex<double>* wfcr_dev = nullptr;
242+
if (!std::is_same<Device, base_device::DEVICE_CPU>::value)
243+
{
244+
base_device::memory::resize_memory_op<std::complex<double>, Device>()(wfcr_dev, pelec->basis->nrxx);
245+
}
246+
208247
for (int ik = 0; ik < pelec->klist->get_nks(); ++ik)
209248
{
210249
psi.fix_k(ik);
@@ -213,7 +252,18 @@ void ldos_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
213252

214253
for (int ib = 0; ib < nbands; ib++)
215254
{
216-
pelec->basis->recip2real(&psi(ib, 0), wfcr.data(), ik);
255+
// FFT on device and copy result back to host
256+
if (std::is_same<Device, base_device::DEVICE_CPU>::value)
257+
{
258+
pelec->basis->recip2real(&psi(ib, 0), wfcr.data(), ik);
259+
}
260+
else
261+
{
262+
pelec->basis->recip_to_real(ctx, &psi(ib, 0), wfcr_dev, ik);
263+
264+
base_device::memory::synchronize_memory_op<std::complex<double>, base_device::DEVICE_CPU, Device>()(
265+
wfcr.data(), wfcr_dev, pelec->basis->nrxx);
266+
}
217267
const double weight = pelec->klist->wk[ik] / ucell.omega;
218268

219269
for (int ir = 0; ir < pelec->basis->nrxx; ir++)
@@ -240,6 +290,12 @@ void ldos_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
240290
}
241291
}
242292

293+
// Free device memory
294+
if (!std::is_same<Device, base_device::DEVICE_CPU>::value)
295+
{
296+
base_device::memory::delete_memory_op<std::complex<double>, Device>()(wfcr_dev);
297+
}
298+
243299
std::ofstream ofs_ldos;
244300
std::stringstream fn;
245301
fn << PARAM.globalv.global_out_dir << "LDOS.txt";
@@ -368,4 +424,17 @@ void trilinear_interpolate(const std::vector<std::vector<int>>& points,
368424
#endif
369425
}
370426

427+
template void cal_ldos_pw<base_device::DEVICE_CPU>(const elecstate::ElecStatePW<std::complex<double>>* pelec,
428+
const psi::Psi<std::complex<double>, base_device::DEVICE_CPU>& psi,
429+
const base_device::DEVICE_CPU* ctx,
430+
const Parallel_Grid& pgrid,
431+
const UnitCell& ucell);
432+
#if defined(__CUDA) || defined(__ROCM)
433+
template void cal_ldos_pw<base_device::DEVICE_GPU>(const elecstate::ElecStatePW<std::complex<double>>* pelec,
434+
const psi::Psi<std::complex<double>, base_device::DEVICE_GPU>& psi,
435+
const base_device::DEVICE_GPU* ctx,
436+
const Parallel_Grid& pgrid,
437+
const UnitCell& ucell);
438+
#endif
439+
371440
} // namespace ModuleIO

‎source/source_io/module_dos/cal_ldos.h‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,18 +32,24 @@ class Cal_ldos
3232

3333
}; // namespace Cal_ldos
3434

35+
template <typename Device>
3536
void cal_ldos_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
36-
const psi::Psi<std::complex<double>>& psi,
37+
const psi::Psi<std::complex<double>, Device>& psi,
38+
const Device* ctx,
3739
const Parallel_Grid& pgrid,
3840
const UnitCell& ucell);
3941

42+
template <typename Device>
4043
void stm_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
41-
const psi::Psi<std::complex<double>>& psi,
44+
const psi::Psi<std::complex<double>, Device>& psi,
45+
const Device* ctx,
4246
const Parallel_Grid& pgrid,
4347
const UnitCell& ucell);
4448

49+
template <typename Device>
4550
void ldos_mode_pw(const elecstate::ElecStatePW<std::complex<double>>* pelec,
46-
const psi::Psi<std::complex<double>>& psi,
51+
const psi::Psi<std::complex<double>, Device>& psi,
52+
const Device* ctx,
4753
const Parallel_Grid& pgrid,
4854
const UnitCell& ucell);
4955

0 commit comments

Comments
 (0)