Skip to content

Fix: a mpi bug in DeePKS. #7485

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mohanchen merged 5 commits into from
Jun 19, 2026
Merged

Fix: a mpi bug in DeePKS. #7485

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c67bf68
Fix a mpi bug in DeePKS.
ErjieWu Jun 18, 2026
f4db38d
Merge branch 'develop' into fix
ErjieWu Jun 18, 2026
2c843fb
Fix the mpi bug and update the reference value in tests.
ErjieWu Jun 18, 2026
18cd3d7
Update DeePKS UT.
ErjieWu Jun 18, 2026
af377b7
Merge branch 'develop' into fix
ErjieWu Jun 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion source/source_lcao/module_deepks/LCAO_deepks_interface.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -629,7 +629,7 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
int R_size = DeePKS_domain::get_R_size(*h_deltaR);
torch::Tensor overlap_out;
torch::Tensor iRmat;
DeePKS_domain::prepare_phialpha_iRmat(nlocal, R_size, deepks_param, phialpha, ucell, orb, GridD, overlap_out, iRmat);
DeePKS_domain::prepare_phialpha_iRmat(nlocal, R_size, deepks_param, phialpha, ucell, orb, *ParaV, GridD, overlap_out, iRmat);
const std::string file_overlap = PARAM.globalv.global_out_dir + "deepks_phialpha_r.npy";
LCAO_deepks_io::save_tensor2npy<double>(file_overlap, overlap_out, rank);
const std::string file_iRmat = PARAM.globalv.global_out_dir + "deepks_iRmat.npy";
Expand Down
8 changes: 8 additions & 0 deletions source/source_lcao/module_deepks/deepks_vdrpre.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal,
const std::vector<hamilt::HContainer<double>*> phialpha,
const UnitCell& ucell,
const LCAO_Orbitals& orb,
const Parallel_Orbitals& pv,
const Grid_Driver& GridD,
torch::Tensor& overlap,
torch::Tensor& iRmat)
Expand Down Expand Up @@ -81,6 +82,10 @@ void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal,

for (int ix = 0; ix < nw_tot; ix++)
{
if (pv.global2local_row(start + ix) < 0 || pv.global2local_col(start + ix) < 0)
{
continue;
}
for (int iy = 0; iy < deepks_param.des_per_atom; iy++)
{
overlap_accessor[iat][nnmax_vec[iat]][start + ix][iy] = overlap_mat->get_value(ix, iy);
Expand All @@ -89,6 +94,9 @@ void DeePKS_domain::prepare_phialpha_iRmat(const int nlocal,
nnmax_vec[iat]++;
}
);
#ifdef __MPI
Parallel_Reduce::reduce_all(overlap.data_ptr<double>(), overlap.numel());
#endif
iRmat = mapping_R(dRmat_tmp.unsqueeze(1) - dRmat_tmp.unsqueeze(2));
ModuleBase::timer::end("DeePKS_domain", "prepare_phialpha_iRmat");
return;
Expand Down
1 change: 1 addition & 0 deletions source/source_lcao/module_deepks/deepks_vdrpre.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ void prepare_phialpha_iRmat(const int nlocal,
const std::vector<hamilt::HContainer<double>*> phialpha,
const UnitCell& ucell,
const LCAO_Orbitals& orb,
const Parallel_Orbitals& pv,
const Grid_Driver& GridD,
torch::Tensor& overlap,
torch::Tensor& iRmat);
Expand Down
1 change: 1 addition & 0 deletions source/source_lcao/module_deepks/test/LCAO_deepks_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -308,6 +308,7 @@ void test_deepks<T>::check_vdrpre()
this->ld.phialpha,
ucell,
ORB,
ParaO,
Test_Deepks::GridD,
overlap_out,
iRmat);
Expand Down
2 changes: 1 addition & 1 deletion tests/09_DeePKS/22_NO_GO_deepks_vdelta_r_2/result.ref
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,6 @@ deepks_e_label 17.126764505645642
deepks_edelta 0.09815855485768665
deepks_hr_label_pass 0
deepks_vdelta_r_pass 0
deepks_phialpha_r 40.10770566607966
deepks_phialpha_r 73.40481582340394
deepks_gevdm 54.0
totaltimeref 1.81
2 changes: 1 addition & 1 deletion tests/09_DeePKS/24_NO_KP_deepks_vdelta_r_2/result.ref
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,6 @@ deepks_e_label 17.12676450564565
deepks_edelta 0.09815855485768665
deepks_hr_label_pass 0
deepks_vdelta_r_pass 0
deepks_phialpha_r 40.10770566607966
deepks_phialpha_r 73.40481582340394
deepks_gevdm 54.0
totaltimeref 1.81
Loading