Refactor/merge openmp (#7446)

* optimize: add OpenMP parallelization to MD per-atom loops

Add #pragma omp parallel for to major per-atom loops in MD module,
enabling multi-threaded execution for NEP/DPMD potentials and
thermostat/integrator operations.

Scope (23 files):
- source/source_md/: md_base, md_func, fire, msst, nhchain, verlet,
  run_md, md_statistics.h
- source/source_esolver/: esolver_nep, esolver_dp
- source/source_md/test/: 7 unit tests + md_test_fixture.h

Strategy: schedule(static) with if(nat>=256), reduction clauses,
atomic/critical for shared accumulators.
LJ esolver excluded (upstream refactored to UnitCellLite API).

Rebased onto deepmodeling/develop.

Co-Authored-By: Claude <noreply@anthropic.com>

* Resolve merge conflicts and fix Verlet CSVR test# Please enter the commit message for your changes. Lines starting
Resolve merge conflicts and fix Verlet CSVR test#

* fix(md): restore CSVR thermostat after conflict resolution

---------

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Mohan Chen <mohanchen@pku.edu.cn>

Author: 3 people

source/source_esolver/esolver_dp.cpp

@@ -37,9 +37,27 @@ void ESolver_DP::before_all_runners(BaseCell& basecell, 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();
 
     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;
@@ -58,38 +76,36 @@ void ESolver_DP::runner(BaseCell& basecell, 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;
@@ -100,18 +116,20 @@ void ESolver_DP::runner(BaseCell& basecell, 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

source/source_esolver/esolver_dp.h

@@ -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

source/source_esolver/esolver_nep.cpp

@@ -21,7 +21,7 @@
 #include "source_io/module_output/output_log.h"
 #include "source_io/module_parameter/parameter.h"
 
-#include <numeric>
+#include <algorithm>
 #include <unordered_map>
 
 using namespace ModuleESolver;
@@ -35,6 +35,9 @@ void ESolver_NEP::before_all_runners(BaseCell& basecell, 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);
@@ -55,51 +58,54 @@ void ESolver_NEP::runner(BaseCell& basecell, 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;
     const double fact_f = 1.0 / (ModuleBase::Ry_to_eV * ModuleBase::ANGSTROM_AU);
     const double fact_v = 1.0 / (ucell.omega * ModuleBase::Ry_to_eV);
 
     // 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;
@@ -108,22 +114,44 @@ void ESolver_NEP::runner(BaseCell& basecell, 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

source/source_esolver/esolver_nep.h

@@ -93,16 +93,21 @@ class ESolver_NEP : public ESolver
     NEP nep; ///< NEP object for NEP calculations
 #endif
 
-    std::string nep_file;          ///< directory of NEP model file
-    std::vector<int> atype = {};   ///< atom type mapping for NEP model
-    double nep_potential;          ///< computed potential energy
-    ModuleBase::matrix nep_force;  ///< computed atomic forces
-    ModuleBase::matrix nep_virial; ///< computed lattice virials
-    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
+    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