Skip to content

Perf: openmp for cal_force_stress #5956

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
merged 8 commits into from
Mar 5, 2025
Merged

Perf: openmp for cal_force_stress #5956

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
a50d09c
remove wrong timer
dzzz2001 Mar 4, 2025
415c33b
omp for cal_force_stress
dzzz2001 Mar 4, 2025
59a8303
openmp for cal_force_stress in dftu
dzzz2001 Mar 4, 2025
7316e09
openmp for cal_force_stress in dspin
dzzz2001 Mar 4, 2025
db76d77
little change
dzzz2001 Mar 4, 2025
1ed0342
fix bug
dzzz2001 Mar 4, 2025
873a0c2
fix a bug
dzzz2001 Mar 5, 2025
f4c6cc8
Merge branch 'develop' into opt
dzzz2001 Mar 5, 2025
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
4 changes: 0 additions & 4 deletions source/module_elecstate/module_charge/charge_mixing.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,7 +193,6 @@ void Charge_Mixing::mix_reset()
bool Charge_Mixing::if_scf_oscillate(const int iteration, const double drho, const int iternum_used, const double threshold)
{
ModuleBase::TITLE("Charge_Mixing", "if_scf_oscillate");
ModuleBase::timer::tick("Charge_Mixing", "if_scf_oscillate");

if(this->_drho_history.size() == 0)
{
Expand Down Expand Up @@ -241,7 +240,4 @@ bool Charge_Mixing::if_scf_oscillate(const int iteration, const double drho, con
}

return false;

ModuleBase::timer::tick("Charge_Mixing", "if_scf_oscillate");

}
30 changes: 24 additions & 6 deletions source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/dftu_force_stress.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,11 @@ void DFTU<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
}
// 1. calculate <psi|beta> for each pair of atoms
// loop over all on-site atoms
int atom_index = 0;
#pragma omp parallel
{
std::vector<double> stress_local(6, 0);
ModuleBase::matrix force_local(force.nr, force.nc);
#pragma omp for schedule(dynamic)
for (int iat0 = 0; iat0 < this->ucell->nat; iat0++)
{
// skip the atoms without plus-U
Expand All @@ -51,7 +55,7 @@ void DFTU<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
continue;
}
const int tlp1 = 2 * target_L + 1;
AdjacentAtomInfo& adjs = this->adjs_all[atom_index++];
AdjacentAtomInfo& adjs = this->adjs_all[iat0];

std::vector<std::unordered_map<int, std::vector<double>>> nlm_tot;
nlm_tot.resize(adjs.adj_num + 1);
Expand Down Expand Up @@ -156,8 +160,8 @@ void DFTU<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
const int T1 = adjs.ntype[ad1];
const int I1 = adjs.natom[ad1];
const int iat1 = ucell->itia2iat(T1, I1);
double* force_tmp1 = (cal_force) ? &force(iat1, 0) : nullptr;
double* force_tmp2 = (cal_force) ? &force(iat0, 0) : nullptr;
double* force_tmp1 = (cal_force) ? &force_local(iat1, 0) : nullptr;
double* force_tmp2 = (cal_force) ? &force_local(iat0, 0) : nullptr;
ModuleBase::Vector3<int>& R_index1 = adjs.box[ad1];
ModuleBase::Vector3<double> dis1 = adjs.adjacent_tau[ad1] - tau0;
for (int ad2 = 0; ad2 < adjs.adj_num + 1; ++ad2)
Expand Down Expand Up @@ -205,13 +209,27 @@ void DFTU<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
this->nspin,
dis1,
dis2,
stress_tmp.data());
stress_local.data());
}
}
}
}
}

#pragma omp critical
{
if(cal_force)
{
force += force_local;
}
if(cal_stress)
{
for(int i = 0; i < 6; i++)
{
stress_tmp[i] += stress_local[i];
}
}
}
}
if (cal_force)
{
#ifdef __MPI
Expand Down
27 changes: 23 additions & 4 deletions source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/dspin_force_stress.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,11 @@ void DeltaSpin<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
}
// 1. calculate <psi|beta> for each pair of atoms
// loop over all on-site atoms
int atom_index = 0;
#pragma omp parallel
{
std::vector<double> stress_local(6, 0);
ModuleBase::matrix force_local(force.nr, force.nc);
#pragma omp for schedule(dynamic)
for (int iat0 = 0; iat0 < this->ucell->nat; iat0++)
{
if(!this->constraint_atom_list[iat0])
Expand Down Expand Up @@ -133,8 +137,8 @@ void DeltaSpin<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
const int T1 = adjs.ntype[ad1];
const int I1 = adjs.natom[ad1];
const int iat1 = ucell->itia2iat(T1, I1);
double* force_tmp1 = (cal_force) ? &force(iat1, 0) : nullptr;
double* force_tmp2 = (cal_force) ? &force(iat0, 0) : nullptr;
double* force_tmp1 = (cal_force) ? &force_local(iat1, 0) : nullptr;
double* force_tmp2 = (cal_force) ? &force_local(iat0, 0) : nullptr;
ModuleBase::Vector3<int>& R_index1 = adjs.box[ad1];
ModuleBase::Vector3<double> dis1 = adjs.adjacent_tau[ad1] - tau0;
for (int ad2 = 0; ad2 < adjs.adj_num + 1; ++ad2)
Expand Down Expand Up @@ -183,12 +187,27 @@ void DeltaSpin<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
this->nspin,
dis1,
dis2,
stress_tmp.data());
stress_local.data());
}
}
}
}
}
#pragma omp critical
{
if(cal_force)
{
force += force_local;
}
if(cal_stress)
{
for(int i = 0; i < 6; i++)
{
stress_tmp[i] += stress_local[i];
}
}
}
}

if (cal_force)
{
Expand Down
75 changes: 51 additions & 24 deletions source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/nonlocal_force_stress.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,17 +27,22 @@ void NonlocalNew<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
}
// 1. calculate <psi|beta> for each pair of atoms
// loop over all on-site atoms
int atom_index = 0;
#pragma omp parallel
{
std::vector<double> stress_local(6, 0);
ModuleBase::matrix force_local(force.nr, force.nc);
#pragma omp for schedule(dynamic)
for (int iat0 = 0; iat0 < this->ucell->nat; iat0++)
{
// skip the atoms without plus-U
auto tau0 = ucell->get_tau(iat0);
int I0 = 0;
ucell->iat2iait(iat0, &I0, &this->current_type);
int T0 = 0;
ucell->iat2iait(iat0, &I0, &T0);

// first step: find the adjacent atoms and filter the real adjacent atoms
AdjacentAtomInfo adjs;
this->gridD->Find_atom(*ucell, tau0, this->current_type, I0, &adjs);
this->gridD->Find_atom(*ucell, tau0, T0, I0, &adjs);

std::vector<bool> is_adj(adjs.adj_num + 1, false);
for (int ad = 0; ad < adjs.adj_num + 1; ++ad)
Expand All @@ -51,7 +56,7 @@ void NonlocalNew<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
// When equal, the theoretical value of matrix element is zero,
// but the calculated value is not zero due to the numerical error, which would lead to result changes.
if (this->ucell->cal_dtau(iat0, iat1, R_index1).norm() * this->ucell->lat0
< orb_cutoff_[T1] + this->ucell->infoNL.Beta[this->current_type].get_rcut_max())
< orb_cutoff_[T1] + this->ucell->infoNL.Beta[T0].get_rcut_max())
{
is_adj[ad] = true;
}
Expand Down Expand Up @@ -89,7 +94,7 @@ void NonlocalNew<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
int M1 = (m1 % 2 == 0) ? -m1 / 2 : (m1 + 1) / 2;

ModuleBase::Vector3<double> dtau = tau0 - tau1;
intor_->snap(T1, L1, N1, M1, this->current_type, dtau * this->ucell->lat0, true /*cal_deri*/, nlm);
intor_->snap(T1, L1, N1, M1, T0, dtau * this->ucell->lat0, true /*cal_deri*/, nlm);
// select the elements of nlm with target_L
const int length = nlm[0].size();
std::vector<double> nlm_target(length * 4);
Expand All @@ -111,8 +116,8 @@ void NonlocalNew<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
const int T1 = adjs.ntype[ad1];
const int I1 = adjs.natom[ad1];
const int iat1 = ucell->itia2iat(T1, I1);
double* force_tmp1 = (cal_force) ? &force(iat1, 0) : nullptr;
double* force_tmp2 = (cal_force) ? &force(iat0, 0) : nullptr;
double* force_tmp1 = (cal_force) ? &force_local(iat1, 0) : nullptr;
double* force_tmp2 = (cal_force) ? &force_local(iat0, 0) : nullptr;
ModuleBase::Vector3<int>& R_index1 = adjs.box[ad1];
ModuleBase::Vector3<double> dis1 = adjs.adjacent_tau[ad1] - tau0;
for (int ad2 = 0; ad2 < adjs.adj_num + 1; ++ad2)
Expand All @@ -139,6 +144,7 @@ void NonlocalNew<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
if (cal_force) {
this->cal_force_IJR(iat1,
iat2,
T0,
paraV,
nlm_iat0[ad1],
nlm_iat0[ad2],
Expand All @@ -151,18 +157,35 @@ void NonlocalNew<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
if (cal_stress) {
this->cal_stress_IJR(iat1,
iat2,
T0,
paraV,
nlm_iat0[ad1],
nlm_iat0[ad2],
tmp,
dis1,
dis2,
stress_tmp.data());
stress_local.data());
}
}
}
}
}
#pragma omp critical
{
if(cal_force)
{
force += force_local;
}
if(cal_stress)
{
for(int i = 0; i < 6; i++)
{
stress_tmp[i] += stress_local[i];
}
}

}
}

if (cal_force)
{
Expand Down Expand Up @@ -202,6 +225,7 @@ void NonlocalNew<OperatorLCAO<TK, TR>>::cal_force_stress(const bool cal_force,
template <>
void NonlocalNew<OperatorLCAO<std::complex<double>, std::complex<double>>>::cal_force_IJR(const int& iat1,
const int& iat2,
const int& T0,
const Parallel_Orbitals* paraV,
const std::unordered_map<int, std::vector<double>>& nlm1_all,
const std::unordered_map<int, std::vector<double>>& nlm2_all,
Expand Down Expand Up @@ -241,11 +265,11 @@ void NonlocalNew<OperatorLCAO<std::complex<double>, std::complex<double>>>::cal_
std::vector<std::complex<double>> nlm_tmp(12, ModuleBase::ZERO);
for (int is = 0; is < 4; ++is)
{
for (int no = 0; no < this->ucell->atoms[this->current_type].ncpp.non_zero_count_soc[is]; no++)
for (int no = 0; no < this->ucell->atoms[T0].ncpp.non_zero_count_soc[is]; no++)
{
const int p1 = this->ucell->atoms[this->current_type].ncpp.index1_soc[is][no];
const int p2 = this->ucell->atoms[this->current_type].ncpp.index2_soc[is][no];
this->ucell->atoms[this->current_type].ncpp.get_d(is, p1, p2, tmp_d);
const int p1 = this->ucell->atoms[T0].ncpp.index1_soc[is][no];
const int p2 = this->ucell->atoms[T0].ncpp.index2_soc[is][no];
this->ucell->atoms[T0].ncpp.get_d(is, p1, p2, tmp_d);
nlm_tmp[is*3] += nlm1[p1 + length] * nlm2[p2] * (*tmp_d);
nlm_tmp[is*3+1] += nlm1[p1 + length * 2] * nlm2[p2] * (*tmp_d);
nlm_tmp[is*3+2] += nlm1[p1 + length * 3] * nlm2[p2] * (*tmp_d);
Expand All @@ -270,6 +294,7 @@ void NonlocalNew<OperatorLCAO<std::complex<double>, std::complex<double>>>::cal_
template <>
void NonlocalNew<OperatorLCAO<std::complex<double>, std::complex<double>>>::cal_stress_IJR(const int& iat1,
const int& iat2,
const int& T0,
const Parallel_Orbitals* paraV,
const std::unordered_map<int, std::vector<double>>& nlm1_all,
const std::unordered_map<int, std::vector<double>>& nlm2_all,
Expand Down Expand Up @@ -311,11 +336,11 @@ void NonlocalNew<OperatorLCAO<std::complex<double>, std::complex<double>>>::cal_
std::vector<std::complex<double>> nlm_tmp(npol2 * 6, ModuleBase::ZERO);
for (int is = 0; is < 4; ++is)
{
for (int no = 0; no < this->ucell->atoms[this->current_type].ncpp.non_zero_count_soc[is]; no++)
for (int no = 0; no < this->ucell->atoms[T0].ncpp.non_zero_count_soc[is]; no++)
{
const int p1 = this->ucell->atoms[this->current_type].ncpp.index1_soc[is][no];
const int p2 = this->ucell->atoms[this->current_type].ncpp.index2_soc[is][no];
this->ucell->atoms[this->current_type].ncpp.get_d(is, p1, p2, tmp_d);
const int p1 = this->ucell->atoms[T0].ncpp.index1_soc[is][no];
const int p2 = this->ucell->atoms[T0].ncpp.index2_soc[is][no];
this->ucell->atoms[T0].ncpp.get_d(is, p1, p2, tmp_d);
nlm_tmp[is*6] += (nlm1[p1 + length] * dis1.x * nlm2[p2] + nlm1[p1] * nlm2[p2 + length] * dis2.x) * (*tmp_d);
nlm_tmp[is*6+1] += (nlm1[p1 + length] * dis1.y * nlm2[p2] + nlm1[p1] * nlm2[p2 + length] * dis2.y) * (*tmp_d);
nlm_tmp[is*6+2] += (nlm1[p1 + length] * dis1.z * nlm2[p2] + nlm1[p1] * nlm2[p2 + length] * dis2.z) * (*tmp_d);
Expand All @@ -341,6 +366,7 @@ void NonlocalNew<OperatorLCAO<std::complex<double>, std::complex<double>>>::cal_
template <typename TK, typename TR>
void NonlocalNew<OperatorLCAO<TK, TR>>::cal_force_IJR(const int& iat1,
const int& iat2,
const int& T0,
const Parallel_Orbitals* paraV,
const std::unordered_map<int, std::vector<double>>& nlm1_all,
const std::unordered_map<int, std::vector<double>>& nlm2_all,
Expand All @@ -367,11 +393,11 @@ void NonlocalNew<OperatorLCAO<TK, TR>>::cal_force_IJR(const int& iat1,
assert(nlm1.size() == nlm2.size());
#endif
std::vector<double> nlm_tmp(3, 0.0);
for (int no = 0; no < this->ucell->atoms[this->current_type].ncpp.non_zero_count_soc[0]; no++)
for (int no = 0; no < this->ucell->atoms[T0].ncpp.non_zero_count_soc[0]; no++)
{
const int p1 = this->ucell->atoms[this->current_type].ncpp.index1_soc[0][no];
const int p2 = this->ucell->atoms[this->current_type].ncpp.index2_soc[0][no];
this->ucell->atoms[this->current_type].ncpp.get_d(0, p1, p2, tmp_d);
const int p1 = this->ucell->atoms[T0].ncpp.index1_soc[0][no];
const int p2 = this->ucell->atoms[T0].ncpp.index2_soc[0][no];
this->ucell->atoms[T0].ncpp.get_d(0, p1, p2, tmp_d);
nlm_tmp[0] += nlm1[p1 + length] * nlm2[p2] * (*tmp_d);
nlm_tmp[1] += nlm1[p1 + length * 2] * nlm2[p2] * (*tmp_d);
nlm_tmp[2] += nlm1[p1 + length * 3] * nlm2[p2] * (*tmp_d);
Expand All @@ -390,6 +416,7 @@ void NonlocalNew<OperatorLCAO<TK, TR>>::cal_force_IJR(const int& iat1,
template <typename TK, typename TR>
void NonlocalNew<OperatorLCAO<TK, TR>>::cal_stress_IJR(const int& iat1,
const int& iat2,
const int& T0,
const Parallel_Orbitals* paraV,
const std::unordered_map<int, std::vector<double>>& nlm1_all,
const std::unordered_map<int, std::vector<double>>& nlm2_all,
Expand Down Expand Up @@ -417,11 +444,11 @@ void NonlocalNew<OperatorLCAO<TK, TR>>::cal_stress_IJR(const int& iat1,
assert(nlm1.size() == nlm2.size());
#endif
std::vector<double> nlm_tmp(6, 0.0);
for (int no = 0; no < this->ucell->atoms[this->current_type].ncpp.non_zero_count_soc[0]; no++)
for (int no = 0; no < this->ucell->atoms[T0].ncpp.non_zero_count_soc[0]; no++)
{
const int p1 = this->ucell->atoms[this->current_type].ncpp.index1_soc[0][no];
const int p2 = this->ucell->atoms[this->current_type].ncpp.index2_soc[0][no];
this->ucell->atoms[this->current_type].ncpp.get_d(0, p1, p2, tmp_d);
const int p1 = this->ucell->atoms[T0].ncpp.index1_soc[0][no];
const int p2 = this->ucell->atoms[T0].ncpp.index2_soc[0][no];
this->ucell->atoms[T0].ncpp.get_d(0, p1, p2, tmp_d);
nlm_tmp[0] += (nlm1[p1 + length] * dis1.x * nlm2[p2] + nlm1[p1] * nlm2[p2 + length] * dis2.x) * (*tmp_d);
nlm_tmp[1] += (nlm1[p1 + length] * dis1.y * nlm2[p2] + nlm1[p1] * nlm2[p2 + length] * dis2.y) * (*tmp_d);
nlm_tmp[2] += (nlm1[p1 + length] * dis1.z * nlm2[p2] + nlm1[p1] * nlm2[p2 + length] * dis2.z) * (*tmp_d);
Expand Down
4 changes: 3 additions & 1 deletion source/module_hamilt_lcao/hamilt_lcaodft/operator_lcao/nonlocal_new.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,12 +102,13 @@ class NonlocalNew<OperatorLCAO<TK, TR>> : public OperatorLCAO<TK, TR>
TR* data_pointer);

const Grid_Driver* gridD = nullptr;
int current_type = 0;

/**
* @brief calculate the atomic Force of <I,J,R> atom pair
*/
void cal_force_IJR(const int& iat1,
const int& iat2,
const int& T0,
const Parallel_Orbitals* paraV,
const std::unordered_map<int, std::vector<double>>& nlm1_all,
const std::unordered_map<int, std::vector<double>>& nlm2_all,
Expand All @@ -119,6 +120,7 @@ class NonlocalNew<OperatorLCAO<TK, TR>> : public OperatorLCAO<TK, TR>
*/
void cal_stress_IJR(const int& iat1,
const int& iat2,
const int& T0,
const Parallel_Orbitals* paraV,
const std::unordered_map<int, std::vector<double>>& nlm1_all,
const std::unordered_map<int, std::vector<double>>& nlm2_all,
Expand Down
Loading