Optimize MD and NEP hot paths with OpenMP-ready refactors by Audrey-777 · Pull Request #7419 · deepmodeling/abacus-develop

Audrey-777 · 2026-06-02T04:26:32Z

Summary

This PR refactors selected MD and NEP-related hot paths and adds OpenMP parallelization for loops with clear performance potential.

Main changes:

Extract MD statistic state handling into reusable helpers.
Refactor MD runner and ML buffer handling to reduce duplicated logic.
Add OpenMP parallel loops for selected MD base/function loops.
Improve NEP virial accumulation using scalar OpenMP reductions instead of a critical-section based accumulation.
Adjust PW Hamiltonian construction interfaces needed by the refactored MD/ESolver flow.
Refactor MD unit test setup into a shared fixture to reduce duplicated initialization code.

This PR intentionally excludes local planning notes, benchmark artifacts, and process documentation.

Tests

The following checks were run locally before preparing this clean PR branch:

cmake --build build-prod -j2
- Passed.
- Confirmed source/source_esolver/esolver_nep.cpp is compiled in the regular build configuration.
cmake --build build-mpi-test --target MODULE_MD_LJ_pot MODULE_MD_func MODULE_MD_fire MODULE_MD_verlet MODULE_MD_nhc MODULE_MD_msst MODULE_MD_lgv -j2
- Passed.
ctest --test-dir build-mpi-test -R 'MODULE_MD' --output-on-failure
- Passed, 7/7 tests.

Note: full NEP linked runtime tests were not run because the local environment does not provide the external NEP library/header required for __NEP execution.

Copilot

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

This PR refactors MD unit tests to share common setup via fixtures, adds new “state-returning” MD helper APIs, and introduces performance-oriented refactors (OpenMP parallel loops and buffer reuse) across MD and ML potential e-solvers.

Changes:

Introduced shared MD test fixtures and migrated multiple MD tests to use them (RAII via std::unique_ptr).
Added MDKineticState/MDStressState and corresponding calc_*_state APIs; parallelized several MD loops with OpenMP.
Reduced per-step allocations in NEP/DP solvers by reusing buffers; refactored MD runner creation; added a templated HamiltPW copy constructor.

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
source/source_pw/module_pwdft/hamilt_pw.h	Reordered includes; declares new templated `HamiltPW` copy ctor; adjusts access specifier formatting.
source/source_pw/module_pwdft/hamilt_pw.cpp	Adds templated `HamiltPW` copy ctor; minor formatting; include reordering.
source/source_md/test/verlet_test.cpp	Migrates test to shared MD integrator fixture.
source/source_md/test/nhchain_test.cpp	Migrates test to shared base fixture; switches to `unique_ptr` ownership.
source/source_md/test/msst_test.cpp	Migrates test to shared MD integrator fixture; adjusts downcasts for `unique_ptr`.
source/source_md/test/md_test_fixture.h	Adds reusable gtest fixtures for MD integrators and MD utility tests.
source/source_md/test/md_func_test.cpp	Migrates test to shared MD function fixture; removes manual allocations.
source/source_md/test/lj_pot_test.cpp	Migrates test to shared LJ fixture; switches to `unique_ptr` e-solver usage.
source/source_md/test/langevin_test.cpp	Migrates test to shared MD integrator fixture.
source/source_md/test/fire_test.cpp	Migrates test to shared MD integrator fixture; adjusts downcasts for `unique_ptr`.
source/source_md/test/CMakeLists.txt	Removes `global_variable.cpp` from MD test dependency list.
source/source_md/run_md.cpp	Refactors MD runner selection into helper returning `std::unique_ptr`.
source/source_md/md_statistics.h	Adds structs for returning kinetic/stress results without mutating caller-owned outputs.
source/source_md/md_func.h	Declares new `calc_kinetic_state` / `calc_stress_state` APIs.
source/source_md/md_func.cpp	Implements new state APIs; adds OpenMP reductions and parallel loops; refactors `compute_stress` / `current_temp`.
source/source_md/md_base.cpp	Adds OpenMP parallel loops in position/velocity updates.
source/source_esolver/esolver_nep.h	Adds persistent buffers and atom index mapping storage for NEP.
source/source_esolver/esolver_nep.cpp	Reuses NEP buffers; parallelizes energy/virial reductions and coordinate packing.
source/source_esolver/esolver_ks_pw.cpp	Changes ELF tau computation call to use `stp.psi_cpu` directly.
source/source_esolver/esolver_dp.h	Adds persistent buffers for DP cell/coords and model outputs.
source/source_esolver/esolver_dp.cpp	Reuses DP buffers and stores model force/virial outputs as members.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Audrey-777 · 2026-06-26T14:22:38Z

+template<typename T, typename Device>
+template<typename T_in, typename Device_in>
+HamiltPW<T, Device>::HamiltPW(const HamiltPW<T_in, Device_in> *hamilt)
+{
+    this->classname = hamilt->classname;
+    this->ppcell = hamilt->ppcell;
+    this->qq_nt = hamilt->qq_nt;
+    this->qq_so = hamilt->qq_so;
+    this->vkb = hamilt->vkb;
+    OperatorPW<std::complex<T_in>, Device_in> * node =
+            reinterpret_cast<OperatorPW<std::complex<T_in>, Device_in> *>(hamilt->ops);
+
+    while(node != nullptr) {
+        if (node->classname == "Ekinetic") {
+            Operator<T, Device>* ekinetic =
+                    new Ekinetic<OperatorPW<T, Device>>(
+                            reinterpret_cast<const Ekinetic<OperatorPW<T_in, Device_in>>*>(node));
+            if(this->ops == nullptr) {
+                this->ops = ekinetic;
+            }
+            else {
+                this->ops->add(ekinetic);
+            }
+            // this->ops = reinterpret_cast<Operator<T, Device>*>(node);
+        }
+        else if (node->classname == "Nonlocal") {
+            Operator<T, Device>* nonlocal =
+                    new Nonlocal<OperatorPW<T, Device>>(
+                            reinterpret_cast<const Nonlocal<OperatorPW<T_in, Device_in>>*>(node));
+            if(this->ops == nullptr) {
+                this->ops = nonlocal;
+            }
+            else {
+                this->ops->add(nonlocal);
+            }
+        }
+        else if (node->classname == "Veff") {
+            Operator<T, Device>* veff =
+                    new Veff<OperatorPW<T, Device>>(
+                            reinterpret_cast<const Veff<OperatorPW<T_in, Device_in>>*>(node));
+            if(this->ops == nullptr) {
+                this->ops = veff;
+            }
+            else {
+                this->ops->add(veff);
+            }
+        }
+        else if (node->classname == "Meta") {
+            Operator<T, Device>* meta =
+                    new Meta<OperatorPW<T, Device>>(
+                            reinterpret_cast<const Meta<OperatorPW<T_in, Device_in>>*>(node));
+            if(this->ops == nullptr) {
+                this->ops = meta;
+            }
+            else {
+                this->ops->add(meta);
+            }
+        }
+        else if (node->classname == "OnsiteProj") {
+            Operator<T, Device>* onsite_proj =
+                    new OnsiteProj<OperatorPW<T, Device>>(
+                            reinterpret_cast<const OnsiteProj<OperatorPW<T_in, Device_in>>*>(node));
+            if(this->ops == nullptr) {
+                this->ops = onsite_proj;
+            }
+            else {
+                this->ops->add(onsite_proj);
+            }
+        }
+        else {
+            ModuleBase::WARNING_QUIT("HamiltPW", "Unrecognized Operator type!");
+        }
+        node = reinterpret_cast<OperatorPW<std::complex<T_in>, Device_in> *>(node->next_op);
+    }
 }


Fixed. The unsafe HamiltPW copy-constructor change was removed from this PR, so there is no longer any reinterpret_cast-based operator-chain copy or missing ucell initialization in the diff.

Audrey-777 · 2026-06-26T14:23:08Z

+		HamiltPW(elecstate::Potential* pot_in,
+				ModulePW::PW_Basis_K* wfc_basis,
+				K_Vectors* p_kv,
+			pseudopot_cell_vnl* nlpp,
+			Plus_U *p_dftu, // mohan add 2025-11-06
+			const UnitCell* ucell);


Fixed together with comment above. The new constructor declaration was removed from this PR, so the indentation and pointer-style issue no longer exists in the changed files.

Audrey-777 · 2026-06-26T14:23:33Z

-        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));


Fixed. The ELF cal_tau change was removed/ restored, and this PR no longer hardwires the calculation to stp.psi_cpu. The device-aware path is no longer modified in this PR.

Audrey-777 · 2026-06-26T14:23:52Z

+MDKineticState calc_kinetic_state(const int& natom,
+                                  const int& frozen_freedom,
+                                  const double* allmass,
+                                  const ModuleBase::Vector3<double>* vel)
+{
+    MDKineticState state;
+    if (3 * natom == frozen_freedom)
+    {
+        return state;
+    }
+
+    state.kinetic = kinetic_energy(natom, vel, allmass);
+    state.temperature = 2 * state.kinetic / (3 * natom - frozen_freedom);
+    return state;
+}
+
+MDStressState calc_stress_state(const int& natom,
+                                const double& omega,
+                                const ModuleBase::Vector3<double>* vel,
+                                const double* allmass,
+                                const ModuleBase::matrix& virial)
+{
+    MDStressState state;
+    temp_vector(natom, vel, allmass, state.temperature_tensor);
+    state.stress.create(3, 3);
+
+    for (int i = 0; i < 3; ++i)
+    {
+        for (int j = 0; j < 3; ++j)
+        {
+            state.stress(i, j) = virial(i, j) + state.temperature_tensor(i, j) / omega;
+        }
+    }
+
+    return state;
+}


Fixed. Direct unit coverage was added for the new MD helper APIs, including deterministic kinetic/stress checks, the fully frozen frozen_freedom == 3 * natom case, and equivalence with the existing current_temp / compute_stress behavior.

Audrey-777 · 2026-06-26T14:24:24Z

+std::unique_ptr<MD_base> create_md_runner(const Parameter& param_in, UnitCell& unit_in)
 {
-    ModuleBase::TITLE("Run_MD", "md_line");
-    ModuleBase::timer::start("Run_MD", "md_line");
-
-    /// determine the md_type
-    MD_base* mdrun = nullptr;
    if (param_in.mdp.md_type == "fire")
    {
-        mdrun = new FIRE(param_in, unit_in);
+        return std::unique_ptr<MD_base>(new FIRE(param_in, unit_in));
    }
-    else if ((param_in.mdp.md_type == "nvt" && param_in.mdp.md_thermostat == "nhc") || param_in.mdp.md_type == "npt")
+    if ((param_in.mdp.md_type == "nvt" && param_in.mdp.md_thermostat == "nhc") || param_in.mdp.md_type == "npt")
    {
-        mdrun = new Nose_Hoover(param_in, unit_in);
+        return std::unique_ptr<MD_base>(new Nose_Hoover(param_in, unit_in));
    }
-    else if (param_in.mdp.md_type == "nve" || param_in.mdp.md_type == "nvt")
+    if (param_in.mdp.md_type == "nve" || param_in.mdp.md_type == "nvt")
    {
-        mdrun = new Verlet(param_in, unit_in);
+        return std::unique_ptr<MD_base>(new Verlet(param_in, unit_in));
    }
-    else if (param_in.mdp.md_type == "langevin")
+    if (param_in.mdp.md_type == "langevin")
    {
-        mdrun = new Langevin(param_in, unit_in);
+        return std::unique_ptr<MD_base>(new Langevin(param_in, unit_in));
    }
-    else if (param_in.mdp.md_type == "msst")
+    if (param_in.mdp.md_type == "msst")
    {
-        mdrun = new MSST(param_in, unit_in);
-    }
-    else
-    {
-        ModuleBase::WARNING_QUIT("md_line", "no such md_type!");
+        return std::unique_ptr<MD_base>(new MSST(param_in, unit_in));
    }

+    ModuleBase::WARNING_QUIT("md_line", "no such md_type!");
+    return nullptr;
+}


Fixed. The MD runner selection logic is now an else if ladder. std::make_unique was not used because this codebase still builds in a C++11-compatible mode; the existing std::unique_ptr(new T(...)) style is consistent with that constraint. The return nullptr after WARNING_QUIT is kept as a conservative fallback.

refactor: optimize MD and NEP hot paths

cbdfa40

Copilot AI review requested due to automatic review settings June 2, 2026 04:26

Copilot AI reviewed Jun 2, 2026

View reviewed changes

mohanchen added project_learning Refactor Refactor ABACUS codes MD & LAM MD and Larege Atomic Models labels Jun 2, 2026

Audrey-777 added 2 commits June 25, 2026 16:51

fix: address OpenMP PR review readiness issues

78e7758

Merge branch 'deepmodeling:develop' into refactor/md-openmp-pr

b3ae935

lijianing-sudo force-pushed the refactor/md-openmp-pr branch 3 times, most recently from b0fccfa to b3ae935 Compare June 26, 2026 11:45

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

Optimize MD and NEP hot paths with OpenMP-ready refactors#7419
Audrey-777 wants to merge 3 commits into
deepmodeling:developfrom
Audrey-777:refactor/md-openmp-pr

Audrey-777 commented Jun 2, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Audrey-777 Jun 26, 2026

Uh oh!

Audrey-777 Jun 26, 2026

Uh oh!

Audrey-777 Jun 26, 2026

Uh oh!

Audrey-777 Jun 26, 2026

Uh oh!

Audrey-777 Jun 26, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

Audrey-777 commented Jun 2, 2026

Summary

Tests

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Audrey-777 Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

Audrey-777 Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

Audrey-777 Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

Audrey-777 Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

Audrey-777 Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants