Optimize MD and NEP hot paths with OpenMP-ready refactors#7419
Open
Audrey-777 wants to merge 1 commit into
Open
Optimize MD and NEP hot paths with OpenMP-ready refactors#7419Audrey-777 wants to merge 1 commit into
Audrey-777 wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
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/MDStressStateand correspondingcalc_*_stateAPIs; parallelized several MD loops with OpenMP. - Reduced per-step allocations in NEP/DP solvers by reusing buffers; refactored MD runner creation; added a templated
HamiltPWcopy 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.
Comment on lines
+164
to
238
| 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); | ||
| } | ||
| } |
Comment on lines
+28
to
+33
| 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); |
| 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)); |
Comment on lines
+56
to
+91
| 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; | ||
| } |
Comment on lines
+19
to
+44
| 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; | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR refactors selected MD and NEP-related hot paths and adds OpenMP parallelization for loops with clear performance potential.
Main changes:
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 -j2source/source_esolver/esolver_nep.cppis 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 -j2ctest --test-dir build-mpi-test -R 'MODULE_MD' --output-on-failureNote: full NEP linked runtime tests were not run because the local environment does not provide the external NEP library/header required for
__NEPexecution.