Skip to content

Refactor/merge openmp #7446

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
lijianing-sudo wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

Refactor/merge openmp #7446

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
76 changes: 47 additions & 29 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 @@ -44,6 +48,20 @@ void ESolver_DP::before_all_runners(UnitCell& ucell, const Input_para& inp)

atype.resize(ucell.nat);

// Build flat atom index for OpenMP coordinate fill in runner()
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++;
}
}

rescaling = inp.mdp.dp_rescaling;
fparam = inp.mdp.dp_fparam;
aparam = inp.mdp.dp_aparam;
Expand All @@ -59,38 +77,36 @@ 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);
int iat = 0;
for (int it = 0; it < ucell.ntype; ++it)
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);
const int nat = ucell.nat;
#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[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;
iat++;
}
const int it = atom_type_index[iat];
const int ia = atom_local_index[iat];
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;
}
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 @@ -101,18 +117,20 @@ void ESolver_DP::runner(UnitCell& ucell, const int istep)
GlobalV::ofs_running << " #TOTAL ENERGY# " << std::setprecision(11) << dp_potential * ModuleBase::Ry_to_eV << " eV"
<< std::endl;

for (int i = 0; i < ucell.nat; ++i)
const int nat_f = ucell.nat;
#pragma omp parallel for schedule(static) if (nat_f >= 256)
for (int i = 0; i < nat_f; ++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
6 changes: 6 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 @@ -109,12 +109,18 @@ class ESolver_DP : public ESolver

std::string dp_file; ///< directory of DP model file
std::vector<int> atype = {}; ///< atom type corresponding to DP model
std::vector<int> atom_type_index; ///< type index (it) for each global atom iat
std::vector<int> atom_local_index; ///< local index (ia) within type for each global atom iat
std::vector<double> fparam = {}; ///< frame parameter for dp potential: dim_fparam
std::vector<double> aparam = {}; ///< atomic parameter for dp potential: natoms x dim_aparam
double rescaling = 1.0; ///< rescaling factor for DP model
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
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
Loading
Loading