Skip to content

Commit a4be870

Browse files
authored
Feature: Support ultrasoft pseudopotentials (USPP) in PW BPCG (#7904)
* Feature: Enable USPP generalized eigenproblems in PW BPCG * Reject unsupported USPP SDFT and nspin=4 calculations * Fix: Synchronize PW projector table caches after cell rescaling * Fix: Correct float input handling in radial_fft_q * Fix: Use double becsum for single-precision USPP force and stress * Reject unsupported USPP calculations before they start * Fix: Use a dense FFT grid when only the z dimension is increased * Test: Add USPP BPCG regression coverage with BNDPAR * Resolve CI/CD failure regarding Makefile * Update 005_PW_SDFT_MALL_BPCG_GPU reference * Update BPCG output references after solver correction
1 parent cc06269 commit a4be870

34 files changed

Lines changed: 1645 additions & 1262 deletions

File tree

source/Makefile.Objects

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -819,6 +819,7 @@ OBJS_SRCPW=h_ewald_pw.o\
819819
setup_pot.o\
820820
setup_pwrho.o\
821821
setup_pwwfc.o\
822+
uspp_support.o\
822823
update_cell_pw.o\
823824
dftu_base.o\
824825
dftu_base_io.o\

source/source_esolver/esolver_fp.cpp

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include "esolver_fp.h"
22

3+
#include "source_base/tool_quit.h"
34
#include "source_cell/cal_ux.h"
45
#include "source_estate/module_charge/symm_rho.h"
56
#include "source_cell/read_pp_ucell.h"
@@ -12,6 +13,7 @@
1213
#include "source_io/module_parameter/parameter.h"
1314

1415
#include "source_pw/module_pwdft/setup_pwrho.h" // mohan 20251005
16+
#include "source_pw/module_pwdft/uspp_support.h"
1517
#include "source_hamilt/module_xc/xc_functional.h" // mohan 20251005
1618
#include "source_io/module_ctrl/ctrl_output_fp.h"
1719
#include "source_io/module_chgpot/write_init.h" // write_chg_init, write_pot_init
@@ -69,8 +71,20 @@ void ESolver_FP::before_all_runners(BaseCell& basecell, const Input_para& inp)
6971
this->inp_->bndpar,
7072
this->inp_->nelec,
7173
this->inp_->nupdown);
74+
7275
elecstate::ParamUpdater::update_from_atoms_info(atoms_info);
7376

77+
XC_Functional::set_xc_type(ucell.atoms[0].ncpp.xc_func);
78+
pw::validate_uspp_support(atoms_info.use_uspp,
79+
inp.basis_type,
80+
inp.esolver_type,
81+
inp.nspin,
82+
XC_Functional::get_func_type(),
83+
inp.berry_phase,
84+
inp.towannier90,
85+
inp.cal_cond);
86+
GlobalV::ofs_running << XC_Functional::output_info() << std::endl;
87+
7488
//! 2) setup pw_rho, pw_rhod, pw_big, sf, and read_pseudopotentials
7589
pw::setup_pwrho(ucell, PARAM.globalv.double_grid, this->pw_rho_flag,
7690
this->pw_rho, this->pw_rhod, this->pw_big, this->classname, inp);
@@ -112,10 +126,6 @@ void ESolver_FP::before_all_runners(BaseCell& basecell, const Input_para& inp)
112126
//! 10) calculate the structure factor
113127
this->sf.setup(&ucell, Pgrid, this->pw_rhod);
114128

115-
//! 11) setup the xc functional
116-
XC_Functional::set_xc_type(ucell.atoms[0].ncpp.xc_func);
117-
GlobalV::ofs_running<<XC_Functional::output_info()<<std::endl;
118-
119129
//! 11) initialize the charge density, we need to first set xc_type,
120130
// then we can call chr.allocate()
121131
this->chr.set_rhopw(this->pw_rhod); // mohan add 20251130

source/source_estate/elecstate_pw.cpp

Lines changed: 31 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -49,10 +49,6 @@ ElecStatePW<T, Device>::~ElecStatePW()
4949
delete[] this->kin_r;
5050
}
5151
}
52-
if (PARAM.globalv.use_uspp)
53-
{
54-
delmem_var_h_op()(this->becsum);
55-
}
5652
delmem_complex_op()(this->wfcr);
5753
delmem_complex_op()(this->wfcr_another_spin);
5854
}
@@ -291,9 +287,8 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
291287
const int nkb = this->ppcell->nkb;
292288
this->vkb = this->ppcell->template get_vkb_data<Real>();
293289
const int nh_tot = this->ppcell->nhm * (this->ppcell->nhm + 1) / 2;
294-
// becsum on CPU (forces_us / stress_us use CPU dgemm)
295-
resmem_var_h_op()(becsum, nh_tot * ucell->nat * PARAM.inp.nspin, "ElecState<PW>::becsum");
296-
setmem_var_h_op()(becsum, 0, nh_tot * ucell->nat * PARAM.inp.nspin);
290+
const int becsum_size = nh_tot * ucell->nat * PARAM.inp.nspin;
291+
this->becsum_.assign(becsum_size, 0.0);
297292

298293
// becp: device buffer for gemm, then D2H for host loops
299294
T* becp = nullptr;
@@ -429,11 +424,11 @@ void ElecStatePW<T, Device>::cal_becsum(const psi::Psi<T, Device>& psi)
429424
{
430425
if (ih == jh)
431426
{
432-
becsum[index + ijh] += std::real(aux_gk_host[ih * nh_atom + jh]);
427+
this->becsum_[index + ijh] += static_cast<double>(std::real(aux_gk_host[ih * nh_atom + jh]));
433428
}
434429
else
435430
{
436-
becsum[index + ijh] += 2.0 * std::real(aux_gk_host[ih * nh_atom + jh]);
431+
this->becsum_[index + ijh] += 2.0 * static_cast<double>(std::real(aux_gk_host[ih * nh_atom + jh]));
437432
}
438433
ijh++;
439434
}
@@ -470,7 +465,7 @@ void ElecStatePW<T, Device>::add_usrho(const psi::Psi<T, Device>& psi)
470465
// add to the charge density in reciprocal space the part which is due to the US augmentation.
471466
if (PARAM.globalv.use_uspp)
472467
{
473-
this->addusdens_g(becsum, this->charge->rhog);
468+
this->addusdens_g(this->charge->rhog);
474469
}
475470
// transform back to real space using dense grids
476471
if (PARAM.globalv.double_grid || PARAM.globalv.use_uspp)
@@ -483,13 +478,14 @@ void ElecStatePW<T, Device>::add_usrho(const psi::Psi<T, Device>& psi)
483478
}
484479

485480
template <typename T, typename Device>
486-
void ElecStatePW<T, Device>::addusdens_g(const Real* becsum, std::complex<double>** rhog)
481+
void ElecStatePW<T, Device>::addusdens_g(std::complex<double>** rhog)
487482
{
488483
const T one{1, 0};
489484
const T zero{0, 0};
490485
const int npw = this->charge->rhopw->npw;
491486
const int lmaxq = this->ppcell->lmaxq;
492487
const int nh_tot = this->ppcell->nhm * (this->ppcell->nhm + 1) / 2;
488+
const double* becsum = this->becsum_.data();
493489
Structure_Factor* psf = this->ppcell->psf;
494490
const std::complex<double> ci_tpi = ModuleBase::NEG_IMAG_UNIT * ModuleBase::TWO_PI;
495491

@@ -577,11 +573,35 @@ void ElecStatePW<T, Device>::addusdens_g(const Real* becsum, std::complex<double
577573
delmem_var_op()(ylmk0);
578574
}
579575

576+
// Taoni add 2026-09-02
577+
// Added to fix USPP single force/stress reading the former float becsum as double.
578+
// The double-only drivers receive a base ElecState, while becsum belongs to the precision-templated ElecStatePW.
579+
// Refactor this bridge for true float force/stress.
580+
template <typename Device>
581+
const std::vector<double>* get_becsum(const ElecState& elec)
582+
{
583+
const ElecStatePW<std::complex<double>, Device>* double_elec = dynamic_cast<const ElecStatePW<std::complex<double>, Device>*>(&elec);
584+
if (double_elec != nullptr)
585+
{
586+
return &double_elec->get_becsum();
587+
}
588+
589+
const ElecStatePW<std::complex<float>, Device>* single_elec = dynamic_cast<const ElecStatePW<std::complex<float>, Device>*>(&elec);
590+
if (single_elec != nullptr)
591+
{
592+
return &single_elec->get_becsum();
593+
}
594+
595+
return nullptr;
596+
}
597+
580598
template class ElecStatePW<std::complex<float>, base_device::DEVICE_CPU>;
581599
template class ElecStatePW<std::complex<double>, base_device::DEVICE_CPU>;
600+
template const std::vector<double>* get_becsum<base_device::DEVICE_CPU>(const ElecState& elec);
582601
#if ((defined __CUDA) || (defined __ROCM))
583602
template class ElecStatePW<std::complex<float>, base_device::DEVICE_GPU>;
584603
template class ElecStatePW<std::complex<double>, base_device::DEVICE_GPU>;
604+
template const std::vector<double>* get_becsum<base_device::DEVICE_GPU>(const ElecState& elec);
585605
#endif
586606

587607
} // namespace elecstate

source/source_estate/elecstate_pw.h

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
#ifndef ELECSTATEPW_H
22
#define ELECSTATEPW_H
33

4+
#include <vector>
5+
46
#include <source_base/macros.h>
57

68
#include "elecstate.h"
@@ -41,7 +43,15 @@ class ElecStatePW : public ElecState
4143
//! calculate becsum for uspp
4244
void cal_becsum(const psi::Psi<T, Device>& psi);
4345

44-
Real* becsum = nullptr;
46+
/**
47+
* @brief Return the USPP projector occupancy coefficients.
48+
*
49+
* @return Read-only canonical double-precision coefficients.
50+
*/
51+
const std::vector<double>& get_becsum() const
52+
{
53+
return becsum_;
54+
}
4555

4656
//! init rho_data and kin_r_data
4757
void init_rho_data();
@@ -74,7 +84,7 @@ class ElecStatePW : public ElecState
7484

7585
//! Non-local pseudopotentials
7686
//! \sum_lm Q_lm(r) \sum_i <psi_i|beta_l><beta_m|psi_i> w_i
77-
void addusdens_g(const Real* becsum, std::complex<double>** rhog);
87+
void addusdens_g(std::complex<double>** rhog);
7888

7989
Device * ctx = {};
8090

@@ -89,6 +99,8 @@ class ElecStatePW : public ElecState
8999
T* wfcr_another_spin = nullptr;
90100

91101
private:
102+
std::vector<double> becsum_;
103+
92104
using meta_op = hamilt::meta_pw_op<Real, Device>;
93105
using elecstate_pw_op = elecstate::elecstate_pw_op<Real, Device>;
94106

@@ -106,16 +118,22 @@ class ElecStatePW : public ElecState
106118
using syncmem_complex_d2h_op = base_device::memory::synchronize_memory_op<T, base_device::DEVICE_CPU, Device>;
107119
using syncmem_complex_h2d_op = base_device::memory::synchronize_memory_op<T, Device, base_device::DEVICE_CPU>;
108120

109-
using resmem_var_h_op = base_device::memory::resize_memory_op<Real, base_device::DEVICE_CPU>;
110-
using delmem_var_h_op = base_device::memory::delete_memory_op<Real, base_device::DEVICE_CPU>;
111-
using setmem_var_h_op = base_device::memory::set_memory_op<Real, base_device::DEVICE_CPU>;
112121
using syncmem_var_h2d_op = base_device::memory::synchronize_memory_op<Real, Device, base_device::DEVICE_CPU>;
113122
using syncmem_var_d2h_op = base_device::memory::synchronize_memory_op<Real, base_device::DEVICE_CPU, Device>;
114123

115124
using gemv_op = ModuleBase::gemv_op<T, Device>;
116125
using gemm_op = ModuleBase::gemm_op<T, Device>;
117126
};
118127

128+
/**
129+
* @brief Return the USPP projector occupancy coefficients of a PW electronic state.
130+
*
131+
* @param elec Electronic state to inspect.
132+
* @return Read-only coefficients, or nullptr when the state is not a supported PW state for Device.
133+
*/
134+
template <typename Device>
135+
const std::vector<double>* get_becsum(const ElecState& elec);
136+
119137
} // namespace elecstate
120138

121139
#endif

source/source_estate/test/elecstate_pw_test.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,7 @@ TEST_F(ElecStatePWTest, ConstructorDouble)
242242
EXPECT_EQ(elecstate_pw_d->charge, chg);
243243
EXPECT_EQ(elecstate_pw_d->klist, klist);
244244
EXPECT_EQ(elecstate_pw_d->bigpw, bigpw);
245+
EXPECT_TRUE(elecstate_pw_d->get_becsum().empty());
245246
}
246247

247248
TEST_F(ElecStatePWTest, ConstructorSingle)
@@ -257,6 +258,7 @@ TEST_F(ElecStatePWTest, ConstructorSingle)
257258
EXPECT_EQ(elecstate_pw_s->charge, chg);
258259
EXPECT_EQ(elecstate_pw_s->klist, klist);
259260
EXPECT_EQ(elecstate_pw_s->bigpw, bigpw);
261+
EXPECT_TRUE(elecstate_pw_s->get_becsum().empty());
260262
}
261263

262264
TEST_F(ElecStatePWTest, InitRhoDataDouble)

0 commit comments

Comments
 (0)