Skip to content

Optimize MD and NEP hot paths with OpenMP-ready refactors #7419

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

Open
Audrey-777 wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

Optimize MD and NEP hot paths with OpenMP-ready refactors #7419

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
46 changes: 25 additions & 21 deletions source/source_esolver/esolver_dp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,10 @@ void ESolver_DP::before_all_runners(UnitCell& ucell, const Input_para& inp)
dp_potential = 0;
dp_force.create(ucell.nat, 3);
dp_virial.create(3, 3);
dp_cell.resize(9);
dp_coord.resize(3 * ucell.nat);
dp_model_force.clear();
dp_model_virial.clear();

ModuleIO::CifParser::write(PARAM.globalv.global_out_dir + "STRU.cif",
ucell,
Expand All @@ -59,38 +63,38 @@ void ESolver_DP::runner(UnitCell& ucell, const int istep)
ModuleBase::TITLE("ESolver_DP", "runner");
ModuleBase::timer::start("ESolver_DP", "runner");

std::vector<double> cell(9, 0.0);
cell[0] = ucell.latvec.e11 * ucell.lat0_angstrom;
cell[1] = ucell.latvec.e12 * ucell.lat0_angstrom;
cell[2] = ucell.latvec.e13 * ucell.lat0_angstrom;
cell[3] = ucell.latvec.e21 * ucell.lat0_angstrom;
cell[4] = ucell.latvec.e22 * ucell.lat0_angstrom;
cell[5] = ucell.latvec.e23 * ucell.lat0_angstrom;
cell[6] = ucell.latvec.e31 * ucell.lat0_angstrom;
cell[7] = ucell.latvec.e32 * ucell.lat0_angstrom;
cell[8] = ucell.latvec.e33 * ucell.lat0_angstrom;

std::vector<double> coord(3 * ucell.nat, 0.0);
dp_cell[0] = ucell.latvec.e11 * ucell.lat0_angstrom;
dp_cell[1] = ucell.latvec.e12 * ucell.lat0_angstrom;
dp_cell[2] = ucell.latvec.e13 * ucell.lat0_angstrom;
dp_cell[3] = ucell.latvec.e21 * ucell.lat0_angstrom;
dp_cell[4] = ucell.latvec.e22 * ucell.lat0_angstrom;
dp_cell[5] = ucell.latvec.e23 * ucell.lat0_angstrom;
dp_cell[6] = ucell.latvec.e31 * ucell.lat0_angstrom;
dp_cell[7] = ucell.latvec.e32 * ucell.lat0_angstrom;
dp_cell[8] = ucell.latvec.e33 * ucell.lat0_angstrom;

dp_coord.resize(3 * ucell.nat);
int iat = 0;
for (int it = 0; it < ucell.ntype; ++it)
{
for (int ia = 0; ia < ucell.atoms[it].na; ++ia)
{
coord[3 * iat] = ucell.atoms[it].tau[ia].x * ucell.lat0_angstrom;
coord[3 * iat + 1] = ucell.atoms[it].tau[ia].y * ucell.lat0_angstrom;
coord[3 * iat + 2] = ucell.atoms[it].tau[ia].z * ucell.lat0_angstrom;
dp_coord[3 * iat] = ucell.atoms[it].tau[ia].x * ucell.lat0_angstrom;
dp_coord[3 * iat + 1] = ucell.atoms[it].tau[ia].y * ucell.lat0_angstrom;
dp_coord[3 * iat + 2] = ucell.atoms[it].tau[ia].z * ucell.lat0_angstrom;
iat++;
}
}
assert(ucell.nat == iat);

#ifdef __DPMD
std::vector<double> f, v;
dp_potential = 0;
dp_force.zero_out();
dp_virial.zero_out();
dp_model_force.clear();
dp_model_virial.clear();

dp.compute(dp_potential, f, v, coord, atype, cell, fparam, aparam);
dp.compute(dp_potential, dp_model_force, dp_model_virial, dp_coord, atype, dp_cell, fparam, aparam);

// rescale the energy, force, and stress
const double fact_e = rescaling / ModuleBase::Ry_to_eV;
Expand All @@ -103,16 +107,16 @@ void ESolver_DP::runner(UnitCell& ucell, const int istep)

for (int i = 0; i < ucell.nat; ++i)
{
dp_force(i, 0) = f[3 * i] * fact_f;
dp_force(i, 1) = f[3 * i + 1] * fact_f;
dp_force(i, 2) = f[3 * i + 2] * fact_f;
dp_force(i, 0) = dp_model_force[3 * i] * fact_f;
dp_force(i, 1) = dp_model_force[3 * i + 1] * fact_f;
dp_force(i, 2) = dp_model_force[3 * i + 2] * fact_f;
}

for (int i = 0; i < 3; ++i)
{
for (int j = 0; j < 3; ++j)
{
dp_virial(i, j) = v[3 * i + j] * fact_v;
dp_virial(i, j) = dp_model_virial[3 * i + j] * fact_v;
}
}
#else
Expand Down
4 changes: 4 additions & 0 deletions source/source_esolver/esolver_dp.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,10 @@ class ESolver_DP : public ESolver
double dp_potential = 0.0; ///< computed potential energy
ModuleBase::matrix dp_force; ///< computed atomic forces
ModuleBase::matrix dp_virial; ///< computed lattice virials
std::vector<double> dp_cell; ///< DP cell buffer in Angstrom
std::vector<double> dp_coord; ///< DP coordinate buffer in Angstrom
std::vector<double> dp_model_force; ///< raw force buffer returned by DP
std::vector<double> dp_model_virial; ///< raw virial buffer returned by DP
};

} // namespace ModuleESolver
Expand Down
4 changes: 1 addition & 3 deletions source/source_esolver/esolver_ks_pw.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -279,9 +279,7 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
// Calculate kinetic energy density tau for ELF if needed
if (PARAM.inp.out_elf[0] > 0)
{
auto* elec_pw = static_cast<elecstate::ElecStatePW<T, Device>*>(this->pelec);
auto& psi = *this->stp.template get_psi_t<T, Device>();
elec_pw->cal_tau(psi);
this->pelec->cal_tau(*(this->stp.psi_cpu));
}

ESolver_KS::after_scf(ucell, istep, conv_esolver);
Expand Down
106 changes: 74 additions & 32 deletions source/source_esolver/esolver_nep.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
#include "source_io/module_output/output_log.h"
#include "source_io/module_output/cif_io.h"

#include <numeric>
#include <algorithm>
#include <unordered_map>

using namespace ModuleESolver;
Expand All @@ -34,9 +34,26 @@ void ESolver_NEP::before_all_runners(UnitCell& ucell, const Input_para& inp)
nep_force.create(ucell.nat, 3);
nep_virial.create(3, 3);
atype.resize(ucell.nat);
nep_cell.resize(9);
nep_coord.resize(3 * ucell.nat);
nep_virial_sum.resize(9);
_e.resize(ucell.nat);
_f.resize(3 * ucell.nat);
_v.resize(9 * ucell.nat);
atom_type_index.resize(ucell.nat);
atom_local_index.resize(ucell.nat);

int iat = 0;
for (int it = 0; it < ucell.ntype; ++it)
{
for (int ia = 0; ia < ucell.atoms[it].na; ++ia)
{
atom_type_index[iat] = it;
atom_local_index[iat] = ia;
++iat;
}
}
assert(ucell.nat == iat);

ModuleIO::CifParser::write(PARAM.globalv.global_out_dir + "STRU.cif",
ucell,
Expand All @@ -56,39 +73,35 @@ void ESolver_NEP::runner(UnitCell& ucell, const int istep)

// note that NEP are column major, thus a transpose is needed
// cell
std::vector<double> cell(9, 0.0);
cell[0] = ucell.latvec.e11 * ucell.lat0_angstrom;
cell[1] = ucell.latvec.e21 * ucell.lat0_angstrom;
cell[2] = ucell.latvec.e31 * ucell.lat0_angstrom;
cell[3] = ucell.latvec.e12 * ucell.lat0_angstrom;
cell[4] = ucell.latvec.e22 * ucell.lat0_angstrom;
cell[5] = ucell.latvec.e32 * ucell.lat0_angstrom;
cell[6] = ucell.latvec.e13 * ucell.lat0_angstrom;
cell[7] = ucell.latvec.e23 * ucell.lat0_angstrom;
cell[8] = ucell.latvec.e33 * ucell.lat0_angstrom;
nep_cell[0] = ucell.latvec.e11 * ucell.lat0_angstrom;
nep_cell[1] = ucell.latvec.e21 * ucell.lat0_angstrom;
nep_cell[2] = ucell.latvec.e31 * ucell.lat0_angstrom;
nep_cell[3] = ucell.latvec.e12 * ucell.lat0_angstrom;
nep_cell[4] = ucell.latvec.e22 * ucell.lat0_angstrom;
nep_cell[5] = ucell.latvec.e32 * ucell.lat0_angstrom;
nep_cell[6] = ucell.latvec.e13 * ucell.lat0_angstrom;
nep_cell[7] = ucell.latvec.e23 * ucell.lat0_angstrom;
nep_cell[8] = ucell.latvec.e33 * ucell.lat0_angstrom;

// coord
std::vector<double> coord(3 * ucell.nat, 0.0);
int iat = 0;
nep_coord.resize(3 * ucell.nat);
const int nat = ucell.nat;
for (int it = 0; it < ucell.ntype; ++it)
#pragma omp parallel for schedule(static) if (nat >= 256)
for (int iat = 0; iat < nat; ++iat)
{
for (int ia = 0; ia < ucell.atoms[it].na; ++ia)
{
coord[iat] = ucell.atoms[it].tau[ia].x * ucell.lat0_angstrom;
coord[iat + nat] = ucell.atoms[it].tau[ia].y * ucell.lat0_angstrom;
coord[iat + 2 * nat] = ucell.atoms[it].tau[ia].z * ucell.lat0_angstrom;
iat++;
}
const int it = atom_type_index[iat];
const int ia = atom_local_index[iat];
nep_coord[iat] = ucell.atoms[it].tau[ia].x * ucell.lat0_angstrom;
nep_coord[iat + nat] = ucell.atoms[it].tau[ia].y * ucell.lat0_angstrom;
nep_coord[iat + 2 * nat] = ucell.atoms[it].tau[ia].z * ucell.lat0_angstrom;
}
assert(ucell.nat == iat);

#ifdef __NEP
nep_potential = 0.0;
nep_force.zero_out();
nep_virial.zero_out();

nep.compute(atype, cell, coord, _e, _f, _v);
nep.compute(atype, nep_cell, nep_coord, _e, _f, _v);

// unit conversion
const double fact_e = 1.0 / ModuleBase::Ry_to_eV;
Expand All @@ -97,11 +110,18 @@ void ESolver_NEP::runner(UnitCell& ucell, const int istep)


// potential energy
nep_potential = fact_e * std::accumulate(_e.begin(), _e.end(), 0.0) ;
double energy_sum = 0.0;
#pragma omp parallel for reduction(+:energy_sum) schedule(static) if (nat >= 256)
for (int i = 0; i < nat; ++i)
{
energy_sum += _e[i];
}
nep_potential = fact_e * energy_sum;
GlobalV::ofs_running << " #TOTAL ENERGY# " << std::setprecision(11) << nep_potential * ModuleBase::Ry_to_eV << " eV"
<< std::endl;

// forces
#pragma omp parallel for schedule(static) if (nat >= 256)
for (int i = 0; i < nat; ++i)
{
nep_force(i, 0) = _f[i] * fact_f;
Expand All @@ -110,22 +130,44 @@ void ESolver_NEP::runner(UnitCell& ucell, const int istep)
}

// virial
std::vector<double> v_sum(9, 0.0);
for (int j = 0; j < 9; ++j)
double v0 = 0.0;
double v1 = 0.0;
double v2 = 0.0;
double v3 = 0.0;
double v4 = 0.0;
double v5 = 0.0;
double v6 = 0.0;
double v7 = 0.0;
double v8 = 0.0;
#pragma omp parallel for reduction(+:v0, v1, v2, v3, v4, v5, v6, v7, v8) schedule(static) if (nat >= 256)
for (int i = 0; i < nat; ++i)
{
for (int i = 0; i < nat; ++i)
{
int index = j * nat + i;
v_sum[j] += _v[index];
}
v0 += _v[i];
v1 += _v[nat + i];
v2 += _v[2 * nat + i];
v3 += _v[3 * nat + i];
v4 += _v[4 * nat + i];
v5 += _v[5 * nat + i];
v6 += _v[6 * nat + i];
v7 += _v[7 * nat + i];
v8 += _v[8 * nat + i];
}
nep_virial_sum[0] = v0;
nep_virial_sum[1] = v1;
nep_virial_sum[2] = v2;
nep_virial_sum[3] = v3;
nep_virial_sum[4] = v4;
nep_virial_sum[5] = v5;
nep_virial_sum[6] = v6;
nep_virial_sum[7] = v7;
nep_virial_sum[8] = v8;

// virial -> stress
for (int i = 0; i < 3; ++i)
{
for (int j = 0; j < 3; ++j)
{
nep_virial(i, j) = v_sum[3 * i + j] * fact_v;
nep_virial(i, j) = nep_virial_sum[3 * i + j] * fact_v;
}
}
#else
Expand Down
7 changes: 6 additions & 1 deletion source/source_esolver/esolver_nep.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,14 +95,19 @@ class ESolver_NEP : public ESolver

std::string nep_file; ///< directory of NEP model file
std::vector<int> atype = {}; ///< atom type mapping for NEP model
std::vector<int> atom_type_index; ///< global atom index to UnitCell atom type
std::vector<int> atom_local_index; ///< global atom index to local index inside atom type
double nep_potential; ///< computed potential energy
ModuleBase::matrix nep_force; ///< computed atomic forces
ModuleBase::matrix nep_virial; ///< computed lattice virials
std::vector<double> nep_cell; ///< NEP cell buffer in Angstrom, column-major
std::vector<double> nep_coord; ///< NEP coordinate buffer in Angstrom, column-major
std::vector<double> nep_virial_sum; ///< summed per-atom virial components
std::vector<double> _e; ///< temporary storage for energy computation
std::vector<double> _f; ///< temporary storage for force computation
std::vector<double> _v; ///< temporary storage for virial computation
};

} // namespace ModuleESolver

#endif
#endif
8 changes: 6 additions & 2 deletions source/source_md/md_base.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,9 @@ void MD_base::update_pos()
{
if (my_rank == 0)
{
for (int i = 0; i < ucell.nat; ++i)
const int natom = ucell.nat;
#pragma omp parallel for schedule(static) if (natom >= 256)
for (int i = 0; i < natom; ++i)
{
for (int k = 0; k < 3; ++k)
{
Expand Down Expand Up @@ -127,7 +129,9 @@ void MD_base::update_vel(const ModuleBase::Vector3<double>* force)
{
if (my_rank == 0)
{
for (int i = 0; i < ucell.nat; ++i)
const int natom = ucell.nat;
#pragma omp parallel for schedule(static) if (natom >= 256)
for (int i = 0; i < natom; ++i)
{
for (int k = 0; k < 3; ++k)
{
Expand Down
Loading
Loading