Skip to content

Commit aaab6dd

Browse files
author
dyzheng
committed
Fix(paged): fix PAGED_GPU buffer overflow, invalid D2H copy, and psi_norm abort
- psi_prepare.cpp: use nbands_l instead of nbands_start in memcpy to prevent buffer overflow when nbands_start > nbands - psi_prepare.cpp: skip invalid D2H copy to GPU pointer in dav solver path for PAGED_GPU mode (data already in kspw_psi CPU buffer) - diago_dav_subspace.cpp: replace WARNING_QUIT with convergence handling when psi_norm <= 1e-12 for both GPU and CPU paths; when all new directions have near-zero norms, mark subspace as converged instead of aborting
1 parent c996775 commit aaab6dd

2 files changed

Lines changed: 35 additions & 16 deletions

File tree

source/source_hsolver/diago_dav_subspace.cpp

Lines changed: 34 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -440,13 +440,23 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
440440
{
441441
if (psi_norm_host[i] <= 1.0e-12)
442442
{
443-
std::cout << "Diago_DavSubspace::cal_grad: psi_norm <= 0 for band " << i << std::endl;
444-
std::cout << "This may be due to npwx < nbands: the number of plane waves is less than" << std::endl;
445-
std::cout << "the number of bands, leading to a rank-deficient problem." << std::endl;
446-
std::cout << "Please increase ecutwfc or reduce nbands." << std::endl;
447-
delmem_real_h_op()(psi_norm_host);
448-
delmem_real_op()(psi_norm);
449-
ModuleBase::WARNING_QUIT("cal_grad", "psi_norm <= 0");
443+
std::cout << "Diago_DavSubspace::cal_grad: norm of new direction for band " << i
444+
<< " is too small, treating as converged." << std::endl;
445+
}
446+
}
447+
{
448+
bool all_zero = true;
449+
for (int i = 0; i < notconv; i++)
450+
{
451+
if (psi_norm_host[i] > 1.0e-12)
452+
{
453+
all_zero = false;
454+
break;
455+
}
456+
}
457+
if (all_zero && notconv > 0)
458+
{
459+
this->notconv = 0;
450460
}
451461
}
452462
delmem_real_h_op()(psi_norm_host);
@@ -470,12 +480,23 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
470480
{
471481
if (psi_norm[i] <= 1.0e-12)
472482
{
473-
std::cout << "Diago_DavSubspace::cal_grad: psi_norm <= 0 for band " << i << std::endl;
474-
std::cout << "This may be due to npwx < nbands: the number of plane waves is less than" << std::endl;
475-
std::cout << "the number of bands, leading to a rank-deficient problem." << std::endl;
476-
std::cout << "Please increase ecutwfc or reduce nbands." << std::endl;
477-
delmem_real_h_op()(psi_norm);
478-
ModuleBase::WARNING_QUIT("cal_grad", "psi_norm <= 0");
483+
std::cout << "Diago_DavSubspace::cal_grad: norm of new direction for band " << i
484+
<< " is too small, treating as converged." << std::endl;
485+
}
486+
}
487+
{
488+
bool all_zero = true;
489+
for (int i = 0; i < notconv; i++)
490+
{
491+
if (psi_norm[i] > 1.0e-12)
492+
{
493+
all_zero = false;
494+
break;
495+
}
496+
}
497+
if (all_zero && notconv > 0)
498+
{
499+
this->notconv = 0;
479500
}
480501
}
481502
delmem_real_h_op()(psi_norm);

source/source_psi/psi_prepare.cpp

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,7 @@ void PSIPrepare<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
179179
if (kspw_psi->get_storage_mode() == psi::PsiStorageMode::PAGED_GPU)
180180
{
181181
T* dst_cpu = kspw_psi->get_cpu_pointer(ik);
182-
std::memcpy(dst_cpu, psi_cpu->get_pointer(), sizeof(T) * nbands_start * nbasis);
182+
std::memcpy(dst_cpu, psi_cpu->get_pointer(), sizeof(T) * nbands_l * nbasis);
183183
kspw_psi->load_k_to_gpu(ik);
184184
}
185185
else
@@ -219,8 +219,6 @@ void PSIPrepare<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
219219
{
220220
if (kspw_psi->get_storage_mode() == psi::PsiStorageMode::PAGED_GPU)
221221
{
222-
base_device::memory::synchronize_memory_op<T, base_device::DEVICE_CPU, Device>()(
223-
kspw_psi->get_pointer(), psi_device->get_pointer(), nbands_l * nbasis);
224222
}
225223
else
226224
{

0 commit comments

Comments
 (0)