From ffd66c43401f42053ed2b62eca397cb406aff541 Mon Sep 17 00:00:00 2001 From: Fei Yang <2501213217@stu.pku.edu.cn> Date: Wed, 14 Oct 2026 20:26:45 +0800 Subject: [PATCH 1/5] fix: pass MDCell ghost communication plan to DeePMD --- cmake/modules/FindDeePMD.cmake | 43 +++++++ source/CMakeLists.txt | 3 + source/source_cell/domain_decomposition.cpp | 49 ++++++++ source/source_cell/domain_decomposition.h | 2 + source/source_cell/ghost_exchange_plan.h | 19 +++ source/source_cell/mdcell.cpp | 22 ++++ source/source_cell/mdcell.h | 6 + .../test/test_domain_decomposition.cpp | 34 ++++++ source/source_esolver/esolver_dp.cpp | 115 ++++++++++++++++-- 9 files changed, 284 insertions(+), 9 deletions(-) create mode 100644 source/source_cell/ghost_exchange_plan.h diff --git a/cmake/modules/FindDeePMD.cmake b/cmake/modules/FindDeePMD.cmake index 50563799216..a0ab7dfdb42 100644 --- a/cmake/modules/FindDeePMD.cmake +++ b/cmake/modules/FindDeePMD.cmake @@ -64,4 +64,47 @@ endif() set(CMAKE_REQUIRED_INCLUDES ${CMAKE_REQUIRED_INCLUDES} ${DeePMD_INCLUDE_DIR}) +include(CheckCXXSourceCompiles) +if(DeePMD_FOUND) + # The communication constructor is inline and calls into the DeePMD + # library, so the feature probe must link against the selected API library. + set(_abacus_saved_required_libraries ${CMAKE_REQUIRED_LIBRARIES}) + if(DeePMDC_FOUND) + set(CMAKE_REQUIRED_LIBRARIES ${CMAKE_REQUIRED_LIBRARIES} ${deepmd_c}) + else() + set(CMAKE_REQUIRED_LIBRARIES ${CMAKE_REQUIRED_LIBRARIES} ${deepmd_cc}) + endif() + if(DeePMDC_FOUND) + check_cxx_source_compiles( + "#include + int main() + { + deepmd::hpp::InputNlist nlist(0, nullptr, nullptr, nullptr, + 0, nullptr, nullptr, nullptr, + nullptr, nullptr, nullptr, nullptr, 1); + return 0; + }" + DeePMD_HAS_COMM_NLIST) + else() + check_cxx_source_compiles( + "#include + int main() + { + deepmd::InputNlist nlist(0, nullptr, nullptr, nullptr, + 0, nullptr, nullptr, nullptr, + nullptr, nullptr, nullptr, nullptr, 1); + return 0; + }" + DeePMD_HAS_COMM_NLIST) + endif() + set(CMAKE_REQUIRED_LIBRARIES ${_abacus_saved_required_libraries}) + unset(_abacus_saved_required_libraries) +endif() + +if(DeePMD_FOUND AND NOT DeePMD_HAS_COMM_NLIST) + message(WARNING + "The installed DeePMD API has no communication-capable InputNlist constructor; " + "message-passing MDCell support requires a newer DeePMD-kit API.") +endif() + mark_as_advanced(DeePMD_INCLUDE_DIR deepmd_c deepmd_cc) diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt index 05ee4118ee5..f5f8d02c911 100644 --- a/source/CMakeLists.txt +++ b/source/CMakeLists.txt @@ -74,6 +74,9 @@ if(ABACUS_COMMIT_INFO_INCLUDE_DIR) endif() if(DEFINED DeePMD_DIR) list(APPEND _abacus_feature_definitions __DPMD HIGH_PREC) + if(DeePMD_HAS_COMM_NLIST) + list(APPEND _abacus_feature_definitions __DPMD_COMM_NLIST) + endif() if(DeePMDC_FOUND) list(APPEND _abacus_feature_definitions __DPMDC) endif() diff --git a/source/source_cell/domain_decomposition.cpp b/source/source_cell/domain_decomposition.cpp index b48a75dd9a6..6f89b716c87 100644 --- a/source/source_cell/domain_decomposition.cpp +++ b/source/source_cell/domain_decomposition.cpp @@ -176,6 +176,7 @@ void DomainDecomposition::exchange_ghost_atoms(MDCell& cell) std::vector& ghost_atoms = cell.ghost_atoms_; ghost_atoms.clear(); ghost_layout_valid_ = false; + cell.clear_ghost_exchange_plan_(); ghost_slots_.clear(); build_ghost_exchange_slots(ghost_slots_); std::vector& slots = ghost_slots_; @@ -311,10 +312,58 @@ void DomainDecomposition::exchange_ghost_atoms(MDCell& cell) slot.ghost_count = recv_count; } ghost_layout_valid_ = true; + cell.set_ghost_exchange_plan_(build_ghost_exchange_plan(cell)); cell.clear_forces_(ghost_atoms); cell.neighbor_layout_valid_ = false; } +GhostExchangePlan DomainDecomposition::build_ghost_exchange_plan(const MDCell& cell) const +{ + if (!initialized_ || !ghost_layout_valid_) + { + throw std::runtime_error("Ghost exchange plan is not available."); + } + + GhostExchangePlan plan; + plan.nprocs = size_; + plan.sendnum.reserve(ghost_slots_.size()); + plan.recvnum.reserve(ghost_slots_.size()); + plan.firstrecv.reserve(ghost_slots_.size()); + plan.send_atom_indices.reserve(ghost_slots_.size()); + plan.sendproc.reserve(ghost_slots_.size()); + plan.recvproc.reserve(ghost_slots_.size()); + + for (const GhostExchangeSlot& slot : ghost_slots_) + { + if (slot.send_atom_indices.size() > static_cast(std::numeric_limits::max())) + { + throw std::overflow_error("Ghost send list exceeds int range."); + } + if (slot.ghost_begin > static_cast(std::numeric_limits::max())) + { + throw std::overflow_error("Ghost receive offset exceeds int range."); + } + if (slot.ghost_count < 0) + { + throw std::runtime_error("Invalid ghost receive count."); + } + + plan.sendnum.push_back(static_cast(slot.send_atom_indices.size())); + plan.recvnum.push_back(slot.ghost_count); + const std::size_t firstrecv = cell.owned_atoms_.size() + slot.ghost_begin; + if (firstrecv > static_cast(std::numeric_limits::max())) + { + throw std::overflow_error("Ghost receive index exceeds int range."); + } + plan.firstrecv.push_back(static_cast(firstrecv)); + plan.send_atom_indices.push_back(slot.send_atom_indices); + plan.sendproc.push_back(slot.send_rank); + plan.recvproc.push_back(slot.recv_rank); + } + + return plan; +} + void DomainDecomposition::update_ghost_atom_positions(MDCell& cell) { synchronize_geometry_(cell); diff --git a/source/source_cell/domain_decomposition.h b/source/source_cell/domain_decomposition.h index 8c503525bad..d5099320fbf 100644 --- a/source/source_cell/domain_decomposition.h +++ b/source/source_cell/domain_decomposition.h @@ -3,6 +3,7 @@ #include "source_base/matrix3.h" #include "source_base/vector3.h" +#include "source_cell/ghost_exchange_plan.h" #include "source_cell/module_neighlist/local_atom.h" #include @@ -47,6 +48,7 @@ class DomainDecomposition void update_ghost_atom_positions(MDCell& cell); void accumulate_ghost_forces(MDCell& cell); void prepare_neighbors(MDCell& cell); + GhostExchangePlan build_ghost_exchange_plan(const MDCell& cell) const; int owner_rank_from_frac(const ModuleBase::Vector3& frac) const; diff --git a/source/source_cell/ghost_exchange_plan.h b/source/source_cell/ghost_exchange_plan.h new file mode 100644 index 00000000000..06b8eeb1e2a --- /dev/null +++ b/source/source_cell/ghost_exchange_plan.h @@ -0,0 +1,19 @@ +#ifndef GHOST_EXCHANGE_PLAN_H +#define GHOST_EXCHANGE_PLAN_H + +#include + +struct GhostExchangePlan +{ + std::vector sendnum; + std::vector recvnum; + std::vector firstrecv; + std::vector > send_atom_indices; + std::vector sendproc; + std::vector recvproc; + int nprocs; + + GhostExchangePlan() : nprocs(1) {} +}; + +#endif // GHOST_EXCHANGE_PLAN_H diff --git a/source/source_cell/mdcell.cpp b/source/source_cell/mdcell.cpp index b82d119fde3..dcc64a09e4f 100644 --- a/source/source_cell/mdcell.cpp +++ b/source/source_cell/mdcell.cpp @@ -43,6 +43,27 @@ const ModuleBase::Matrix3& MDCell::get_GT() const return gt_; } +const GhostExchangePlan& MDCell::get_ghost_exchange_plan() const +{ + if (!ghost_exchange_plan_valid_) + { + throw std::runtime_error("MDCell ghost exchange plan has not been prepared."); + } + return ghost_exchange_plan_; +} + +void MDCell::clear_ghost_exchange_plan_() +{ + ghost_exchange_plan_ = GhostExchangePlan(); + ghost_exchange_plan_valid_ = false; +} + +void MDCell::set_ghost_exchange_plan_(const GhostExchangePlan& plan) +{ + ghost_exchange_plan_ = plan; + ghost_exchange_plan_valid_ = true; +} + void MDCell::set_backing_unitcell(UnitCell& ucell) { backing_unitcell_ = &ucell; @@ -85,6 +106,7 @@ void MDCell::initialize_from_owned_atoms(const ModuleBase::Matrix3& latvec, neighbor_search_.reset(); neighbor_layout_valid_ = false; ghost_atoms_.clear(); + clear_ghost_exchange_plan_(); #ifdef __MPI comm_ = comm_domain.communicator(); rank_ = comm_domain.rank(); diff --git a/source/source_cell/mdcell.h b/source/source_cell/mdcell.h index 2bfe9c8fd65..ebf891311ec 100644 --- a/source/source_cell/mdcell.h +++ b/source/source_cell/mdcell.h @@ -2,6 +2,7 @@ #define MDCELL_H #include "source_cell/basecell.h" +#include "source_cell/ghost_exchange_plan.h" #include "source_cell/strumeta.h" #include "source_cell/module_neighlist/local_atom.h" #include "source_base/matrix3.h" @@ -65,6 +66,7 @@ class MDCell : public BaseCell const std::vector& type_labels() const { return type_labels_; } const std::vector& type_masses() const { return type_masses_; } const std::vector& type_atom_counts() const { return type_atom_counts_; } + const GhostExchangePlan& get_ghost_exchange_plan() const; StruMeta& mutable_stru_meta() { return stru_meta_; } const StruMeta& stru_meta() const { return stru_meta_; } int nowned_atoms() const { return static_cast(owned_atoms_.size()); } @@ -91,6 +93,8 @@ class MDCell : public BaseCell const ModuleBase::Matrix3& get_GT() const override; friend class DomainDecomposition; + void clear_ghost_exchange_plan_(); + void set_ghost_exchange_plan_(const GhostExchangePlan& plan); void sync_backing_unitcell_geometry_(); void clear_forces_(std::vector& atoms); static double wrap_fractional_(double value); @@ -113,6 +117,8 @@ class MDCell : public BaseCell std::unique_ptr neighbor_search_; std::vector > neighbor_reference_frac_; bool neighbor_layout_valid_ = false; + GhostExchangePlan ghost_exchange_plan_; + bool ghost_exchange_plan_valid_ = false; #ifdef __MPI MPI_Comm comm_ = MPI_COMM_NULL; diff --git a/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp b/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp index e7ba6e1ec49..54fafdcbdd0 100644 --- a/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp +++ b/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp @@ -80,6 +80,40 @@ TEST_F(DomainDecompositionTest, CutoffSetterDoesNotExchangeAtoms) } } +TEST_F(DomainDecompositionTest, GhostExchangePlanMatchesGhostLayout) +{ + decomp.prepare_neighbors(cell); + const GhostExchangePlan plan = decomp.build_ghost_exchange_plan(cell); + const std::size_t exchange_count = plan.sendnum.size(); + ASSERT_EQ(plan.recvnum.size(), exchange_count); + ASSERT_EQ(plan.firstrecv.size(), exchange_count); + ASSERT_EQ(plan.send_atom_indices.size(), exchange_count); + ASSERT_EQ(plan.sendproc.size(), exchange_count); + ASSERT_EQ(plan.recvproc.size(), exchange_count); + EXPECT_EQ(plan.nprocs, domain.size()); + + std::size_t ghost_offset = cell.owned_atoms().size(); + std::size_t ghost_count = 0; + for (std::size_t iswap = 0; iswap < exchange_count; ++iswap) + { + EXPECT_EQ(plan.sendnum[iswap], static_cast(plan.send_atom_indices[iswap].size())); + EXPECT_EQ(plan.firstrecv[iswap], static_cast(ghost_offset)); + EXPECT_GE(plan.recvnum[iswap], 0); + EXPECT_GE(plan.sendproc[iswap], 0); + EXPECT_LT(plan.sendproc[iswap], domain.size()); + EXPECT_GE(plan.recvproc[iswap], 0); + EXPECT_LT(plan.recvproc[iswap], domain.size()); + for (const int atom_index : plan.send_atom_indices[iswap]) + { + EXPECT_GE(atom_index, 0); + EXPECT_LT(atom_index, static_cast(cell.owned_atoms().size())); + } + ghost_offset += static_cast(plan.recvnum[iswap]); + ghost_count += static_cast(plan.recvnum[iswap]); + } + EXPECT_EQ(ghost_count, cell.ghost_atoms().size()); +} + TEST_F(DomainDecompositionTest, GhostForcesAreReturnedAndConsumed) { decomp.prepare_neighbors(cell); diff --git a/source/source_esolver/esolver_dp.cpp b/source/source_esolver/esolver_dp.cpp index 3de1522e053..ff7eb7a4693 100644 --- a/source/source_esolver/esolver_dp.cpp +++ b/source/source_esolver/esolver_dp.cpp @@ -20,6 +20,7 @@ #include "esolver_dp.h" #include "source_base/parallel_common.h" #include "source_base/timer.h" +#include "source_cell/domain_decomposition.h" #include "source_cell/mdcell.h" #include "source_cell/module_neighlist/neighbor_search.h" #include "source_cell/cif_io.h" @@ -28,11 +29,21 @@ #include #include +#include +#include #include #include using namespace ModuleESolver; +#ifdef __DPMD +#ifdef __DPMDC +using DPInputNlist = deepmd::hpp::InputNlist; +#else +using DPInputNlist = deepmd::InputNlist; +#endif +#endif + void ESolver_DP::before_all_runners(BaseCell& basecell, const Input_para& inp) { this->inp_ = &inp; @@ -143,16 +154,102 @@ void ESolver_DP::runner(BaseCell& basecell, const int istep) ? NULL : sorted_neighbors[static_cast(iat)].data(); } -#ifdef __DPMDC - deepmd::hpp::InputNlist nlist(nowned_atoms, - nowned_atoms > 0 ? &ilist[0] : NULL, - nowned_atoms > 0 ? &numneigh[0] : NULL, - nowned_atoms > 0 ? &firstneigh[0] : NULL); + + const GhostExchangePlan& ghost_plan = mdcell.get_ghost_exchange_plan(); + const std::size_t nswap_size = ghost_plan.sendnum.size(); + if (ghost_plan.recvnum.size() != nswap_size || ghost_plan.firstrecv.size() != nswap_size + || ghost_plan.send_atom_indices.size() != nswap_size + || ghost_plan.sendproc.size() != nswap_size || ghost_plan.recvproc.size() != nswap_size) + { + ModuleBase::WARNING_QUIT("ESolver_DP", "Invalid MDCell ghost exchange plan."); + } + if (nswap_size > static_cast(std::numeric_limits::max())) + { + ModuleBase::WARNING_QUIT("ESolver_DP", "MDCell ghost exchange plan exceeds int range."); + } + const bool use_mapping = ghost_plan.nprocs == 1; + std::vector mapping; + if (use_mapping) + { + std::map, int> owned_atom_map; + for (int iat = 0; iat < nowned_atoms; ++iat) + { + const LocalAtom& atom = owned_atoms[static_cast(iat)]; + owned_atom_map[std::make_pair(atom.type, atom.type_index)] = iat; + } + mapping.resize(static_cast(natom), 0); + for (int iat = 0; iat < natom; ++iat) + { + if (iat < nowned_atoms) + { + mapping[static_cast(iat)] = iat; + continue; + } + const LocalAtom& atom = ghost_atoms[static_cast(iat - nowned_atoms)]; + const std::map, int>::const_iterator owner + = owned_atom_map.find(std::make_pair(atom.type, atom.type_index)); + if (owner == owned_atom_map.end()) + { + ModuleBase::WARNING_QUIT("ESolver_DP", "Cannot match a DeePMD ghost atom to its local owner."); + } + mapping[static_cast(iat)] = owner->second; + } + } + const int nswap = static_cast(nswap_size); + std::vector > sendlist_storage = ghost_plan.send_atom_indices; + std::vector sendlist(nswap_size, NULL); + for (std::size_t iswap = 0; iswap < nswap_size; ++iswap) + { + if (!sendlist_storage[iswap].empty()) + { + sendlist[iswap] = sendlist_storage[iswap].data(); + } + } + std::vector sendnum_data = ghost_plan.sendnum; + std::vector recvnum_data = ghost_plan.recvnum; + std::vector firstrecv_data = ghost_plan.firstrecv; + std::vector sendproc_data = ghost_plan.sendproc; + std::vector recvproc_data = ghost_plan.recvproc; + int* sendnum = sendnum_data.empty() ? NULL : &sendnum_data[0]; + int* recvnum = recvnum_data.empty() ? NULL : &recvnum_data[0]; + int* firstrecv = firstrecv_data.empty() ? NULL : &firstrecv_data[0]; + int** sendlist_data = sendlist.empty() ? NULL : &sendlist[0]; + int* sendproc = sendproc_data.empty() ? NULL : &sendproc_data[0]; + int* recvproc = recvproc_data.empty() ? NULL : &recvproc_data[0]; + int* ilist_data = nowned_atoms > 0 ? &ilist[0] : NULL; + int* numneigh_data = nowned_atoms > 0 ? &numneigh[0] : NULL; + int** firstneigh_data = nowned_atoms > 0 ? &firstneigh[0] : NULL; + int* mapping_data = use_mapping && natom > 0 ? &mapping[0] : NULL; + void* dp_world = NULL; +#ifdef __MPI + MPI_Comm dp_world_comm = mdcell.communicator(); + dp_world = &dp_world_comm; +#endif +#ifdef __DPMD_COMM_NLIST + DPInputNlist nlist(nowned_atoms, + ilist_data, + numneigh_data, + firstneigh_data, + nswap, + sendnum, + recvnum, + firstrecv, + sendlist_data, + sendproc, + recvproc, + dp_world, + ghost_plan.nprocs); #else - deepmd::InputNlist nlist(nowned_atoms, - nowned_atoms > 0 ? &ilist[0] : NULL, - nowned_atoms > 0 ? &numneigh[0] : NULL, - nowned_atoms > 0 ? &firstneigh[0] : NULL); + DPInputNlist nlist(nowned_atoms, + ilist_data, + numneigh_data, + firstneigh_data); +#endif +#if !defined(__DPMDC) || defined(__DPMD_COMM_NLIST) + if (use_mapping) + { + nlist.set_mapping(mapping_data); + } #endif double local_energy = 0.0; std::vector force, virial; From efc47261577f3c260952b40b4276da424e676b55 Mon Sep 17 00:00:00 2001 From: Fei Yang <2501213217@stu.pku.edu.cn> Date: Thu, 15 Oct 2026 17:20:51 +0800 Subject: [PATCH 2/5] fix: reuse ghost ownership for DeepMD neighbor lists --- cmake/modules/FindDeePMD.cmake | 56 +++++++++++----- source/CMakeLists.txt | 7 +- source/Makefile | 18 ++++++ source/source_cell/domain_decomposition.cpp | 25 ++++++++ source/source_cell/domain_decomposition.h | 4 +- source/source_cell/ghost_exchange_plan.h | 4 ++ .../test/test_domain_decomposition.cpp | 64 ++++++++++++++++++- source/source_esolver/esolver_dp.cpp | 47 ++++++-------- tools/deepmd_input_nlist_probe.cpp | 45 +++++++++++++ 9 files changed, 220 insertions(+), 50 deletions(-) create mode 100644 tools/deepmd_input_nlist_probe.cpp diff --git a/cmake/modules/FindDeePMD.cmake b/cmake/modules/FindDeePMD.cmake index a0ab7dfdb42..d8aab27efda 100644 --- a/cmake/modules/FindDeePMD.cmake +++ b/cmake/modules/FindDeePMD.cmake @@ -75,33 +75,57 @@ if(DeePMD_FOUND) set(CMAKE_REQUIRED_LIBRARIES ${CMAKE_REQUIRED_LIBRARIES} ${deepmd_cc}) endif() if(DeePMDC_FOUND) - check_cxx_source_compiles( - "#include - int main() - { - deepmd::hpp::InputNlist nlist(0, nullptr, nullptr, nullptr, - 0, nullptr, nullptr, nullptr, - nullptr, nullptr, nullptr, nullptr, 1); - return 0; - }" - DeePMD_HAS_COMM_NLIST) + set(_abacus_dp_nlist_type "deepmd::hpp::InputNlist") + set(_abacus_dp_nlist_header "#include ") + else() + set(_abacus_dp_nlist_type "deepmd::InputNlist") + set(_abacus_dp_nlist_header "#include ") + endif() + check_cxx_source_compiles( + "${_abacus_dp_nlist_header} + int main() + { + ${_abacus_dp_nlist_type} nlist(0, nullptr, nullptr, nullptr, + 0, nullptr, nullptr, nullptr, + nullptr, nullptr, nullptr, nullptr, 1); + return 0; + }" + DeePMD_HAS_COMM_NLIST_13) + if(DeePMD_HAS_COMM_NLIST_13) + set(DeePMD_INPUT_NLIST_COMM_ARGS 13) else() check_cxx_source_compiles( - "#include + "${_abacus_dp_nlist_header} int main() { - deepmd::InputNlist nlist(0, nullptr, nullptr, nullptr, - 0, nullptr, nullptr, nullptr, - nullptr, nullptr, nullptr, nullptr, 1); + ${_abacus_dp_nlist_type} nlist(0, nullptr, nullptr, nullptr, + 0, nullptr, nullptr, nullptr, + nullptr, nullptr, nullptr, nullptr); return 0; }" - DeePMD_HAS_COMM_NLIST) + DeePMD_HAS_COMM_NLIST_12) + if(DeePMD_HAS_COMM_NLIST_12) + set(DeePMD_INPUT_NLIST_COMM_ARGS 12) + else() + set(DeePMD_INPUT_NLIST_COMM_ARGS 0) + endif() endif() + check_cxx_source_compiles( + "${_abacus_dp_nlist_header} + int main() + { + ${_abacus_dp_nlist_type} nlist(0, nullptr, nullptr, nullptr); + nlist.set_mapping(nullptr); + return 0; + }" + DeePMD_HAS_NLIST_MAPPING) set(CMAKE_REQUIRED_LIBRARIES ${_abacus_saved_required_libraries}) unset(_abacus_saved_required_libraries) + unset(_abacus_dp_nlist_type) + unset(_abacus_dp_nlist_header) endif() -if(DeePMD_FOUND AND NOT DeePMD_HAS_COMM_NLIST) +if(DeePMD_FOUND AND DEFINED DeePMD_INPUT_NLIST_COMM_ARGS AND DeePMD_INPUT_NLIST_COMM_ARGS EQUAL 0) message(WARNING "The installed DeePMD API has no communication-capable InputNlist constructor; " "message-passing MDCell support requires a newer DeePMD-kit API.") diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt index f5f8d02c911..6555f0bc809 100644 --- a/source/CMakeLists.txt +++ b/source/CMakeLists.txt @@ -74,8 +74,11 @@ if(ABACUS_COMMIT_INFO_INCLUDE_DIR) endif() if(DEFINED DeePMD_DIR) list(APPEND _abacus_feature_definitions __DPMD HIGH_PREC) - if(DeePMD_HAS_COMM_NLIST) - list(APPEND _abacus_feature_definitions __DPMD_COMM_NLIST) + if(DEFINED DeePMD_INPUT_NLIST_COMM_ARGS AND DeePMD_INPUT_NLIST_COMM_ARGS GREATER 0) + list(APPEND _abacus_feature_definitions __DPMD_COMM_NLIST_ARGS=${DeePMD_INPUT_NLIST_COMM_ARGS}) + endif() + if(DeePMD_HAS_NLIST_MAPPING) + list(APPEND _abacus_feature_definitions __DPMD_HAS_NLIST_MAPPING) endif() if(DeePMDC_FOUND) list(APPEND _abacus_feature_definitions __DPMDC) diff --git a/source/Makefile b/source/Makefile index 92dfe1aa8c9..40d53bfed42 100644 --- a/source/Makefile +++ b/source/Makefile @@ -182,6 +182,24 @@ ifdef DeePMD_DIR endif endif +ifdef DeePMD_DIR + # Keep the legacy Make build on the same DeePMD API path as CMake. These + # are syntax-only probes, so no temporary executable or extra link step is + # required during the normal build. + DPMD_PROBE_SOURCE = ${ABACUS_ROOT}../tools/deepmd_input_nlist_probe.cpp + DPMD_COMM_NLIST_13 := $(shell $(CXX) $(OPTS) $(INCLUDES) $(HONG) -DDPMD_PROBE_COMM_13 -fsyntax-only $(DPMD_PROBE_SOURCE) >/dev/null 2>&1 && echo 1) + DPMD_COMM_NLIST_12 := $(shell $(CXX) $(OPTS) $(INCLUDES) $(HONG) -DDPMD_PROBE_COMM_12 -fsyntax-only $(DPMD_PROBE_SOURCE) >/dev/null 2>&1 && echo 1) + DPMD_NLIST_MAPPING := $(shell $(CXX) $(OPTS) $(INCLUDES) $(HONG) -DDPMD_PROBE_MAPPING -fsyntax-only $(DPMD_PROBE_SOURCE) >/dev/null 2>&1 && echo 1) + ifeq ($(DPMD_COMM_NLIST_13),1) + HONG += -D__DPMD_COMM_NLIST_ARGS=13 + else ifeq ($(DPMD_COMM_NLIST_12),1) + HONG += -D__DPMD_COMM_NLIST_ARGS=12 + endif + ifeq ($(DPMD_NLIST_MAPPING),1) + HONG += -D__DPMD_HAS_NLIST_MAPPING + endif +endif + ifdef TensorFlow_DIR #========================== # TensorFlow package diff --git a/source/source_cell/domain_decomposition.cpp b/source/source_cell/domain_decomposition.cpp index 6f89b716c87..d2628977a4f 100644 --- a/source/source_cell/domain_decomposition.cpp +++ b/source/source_cell/domain_decomposition.cpp @@ -1,5 +1,6 @@ #include "source_cell/domain_decomposition.h" #include "source_base/parallel_cell.h" +#include "source_cell/ghost_exchange_plan.h" #include "source_cell/unitcell.h" #include "source_cell/mdcell.h" #include "source_cell/module_neighlist/neighbor_search.h" @@ -361,6 +362,30 @@ GhostExchangePlan DomainDecomposition::build_ghost_exchange_plan(const MDCell& c plan.recvproc.push_back(slot.recv_rank); } + if (plan.nprocs == 1) + { + const std::size_t natom = cell.owned_atoms_.size() + cell.ghost_atoms_.size(); + plan.mapping.resize(natom, 0); + std::map, int> owned_lookup; + for (std::size_t iat = 0; iat < cell.owned_atoms_.size(); ++iat) + { + const LocalAtom& atom = cell.owned_atoms_[iat]; + owned_lookup[std::make_pair(atom.type, atom.type_index)] = static_cast(iat); + plan.mapping[iat] = static_cast(iat); + } + for (std::size_t iat = 0; iat < cell.ghost_atoms_.size(); ++iat) + { + const LocalAtom& atom = cell.ghost_atoms_[iat]; + const std::map, int>::const_iterator owner + = owned_lookup.find(std::make_pair(atom.type, atom.type_index)); + if (owner == owned_lookup.end()) + { + throw std::runtime_error("Cannot match a single-rank ghost atom to an owned atom."); + } + plan.mapping[cell.owned_atoms_.size() + iat] = owner->second; + } + } + return plan; } diff --git a/source/source_cell/domain_decomposition.h b/source/source_cell/domain_decomposition.h index d5099320fbf..c6a57829f23 100644 --- a/source/source_cell/domain_decomposition.h +++ b/source/source_cell/domain_decomposition.h @@ -3,7 +3,6 @@ #include "source_base/matrix3.h" #include "source_base/vector3.h" -#include "source_cell/ghost_exchange_plan.h" #include "source_cell/module_neighlist/local_atom.h" #include @@ -16,6 +15,7 @@ class UnitCell; class MDCell; +struct GhostExchangePlan; namespace ModuleBase { class CommunicationDomain; @@ -48,7 +48,6 @@ class DomainDecomposition void update_ghost_atom_positions(MDCell& cell); void accumulate_ghost_forces(MDCell& cell); void prepare_neighbors(MDCell& cell); - GhostExchangePlan build_ghost_exchange_plan(const MDCell& cell) const; int owner_rank_from_frac(const ModuleBase::Vector3& frac) const; @@ -119,6 +118,7 @@ class DomainDecomposition const std::array& previous_image_shift, const GhostExchangeSlot& slot) const; void build_ghost_exchange_slots(std::vector& slots) const; + GhostExchangePlan build_ghost_exchange_plan(const MDCell& cell) const; PackedAtom pack_atom(const LocalAtom& atom, const std::array& image_shift) const; LocalAtom unpack_ghost_atom(const PackedAtom& packed) const; LocalAtom unpack_owned_atom(const PackedAtom& packed) const; diff --git a/source/source_cell/ghost_exchange_plan.h b/source/source_cell/ghost_exchange_plan.h index 06b8eeb1e2a..8d9fe0efd2a 100644 --- a/source/source_cell/ghost_exchange_plan.h +++ b/source/source_cell/ghost_exchange_plan.h @@ -11,6 +11,10 @@ struct GhostExchangePlan std::vector > send_atom_indices; std::vector sendproc; std::vector recvproc; + // For a single rank, map every local atom (owned and ghost) to its owned + // atom index. Distributed DeePMD obtains this information from exchange + // metadata instead. + std::vector mapping; int nprocs; GhostExchangePlan() : nprocs(1) {} diff --git a/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp b/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp index 54fafdcbdd0..b1c0f8d83fd 100644 --- a/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp +++ b/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp @@ -83,7 +83,7 @@ TEST_F(DomainDecompositionTest, CutoffSetterDoesNotExchangeAtoms) TEST_F(DomainDecompositionTest, GhostExchangePlanMatchesGhostLayout) { decomp.prepare_neighbors(cell); - const GhostExchangePlan plan = decomp.build_ghost_exchange_plan(cell); + const GhostExchangePlan& plan = cell.get_ghost_exchange_plan(); const std::size_t exchange_count = plan.sendnum.size(); ASSERT_EQ(plan.recvnum.size(), exchange_count); ASSERT_EQ(plan.firstrecv.size(), exchange_count); @@ -91,6 +91,26 @@ TEST_F(DomainDecompositionTest, GhostExchangePlanMatchesGhostLayout) ASSERT_EQ(plan.sendproc.size(), exchange_count); ASSERT_EQ(plan.recvproc.size(), exchange_count); EXPECT_EQ(plan.nprocs, domain.size()); + if (domain.size() == 1) + { + ASSERT_EQ(plan.mapping.size(), cell.owned_atoms().size() + cell.ghost_atoms().size()); + for (std::size_t iat = 0; iat < cell.owned_atoms().size(); ++iat) + { + EXPECT_EQ(plan.mapping[iat], static_cast(iat)); + } + for (std::size_t iat = 0; iat < cell.ghost_atoms().size(); ++iat) + { + const int owner = plan.mapping[cell.owned_atoms().size() + iat]; + ASSERT_GE(owner, 0); + ASSERT_LT(owner, static_cast(cell.owned_atoms().size())); + EXPECT_EQ(cell.ghost_atoms()[iat].type_index, + cell.owned_atoms()[static_cast(owner)].type_index); + } + } + else + { + EXPECT_TRUE(plan.mapping.empty()); + } std::size_t ghost_offset = cell.owned_atoms().size(); std::size_t ghost_count = 0; @@ -112,6 +132,48 @@ TEST_F(DomainDecompositionTest, GhostExchangePlanMatchesGhostLayout) ghost_count += static_cast(plan.recvnum[iswap]); } EXPECT_EQ(ghost_count, cell.ghost_atoms().size()); + + if (domain.size() == 1) + { + for (std::size_t iswap = 0; iswap < exchange_count; ++iswap) + { + for (std::size_t iatom = 0; iatom < plan.send_atom_indices[iswap].size(); ++iatom) + { + const int owned_index = plan.send_atom_indices[iswap][iatom]; + ASSERT_GE(plan.firstrecv[iswap], static_cast(cell.owned_atoms().size())); + const std::size_t ghost_index = static_cast(plan.firstrecv[iswap] + - cell.owned_atoms().size()) + + iatom; + ASSERT_LT(ghost_index, cell.ghost_atoms().size()); + EXPECT_EQ(cell.ghost_atoms()[ghost_index].type_index, + cell.owned_atoms()[static_cast(owned_index)].type_index); + } + } + } + + std::vector first_identity; + first_identity.reserve(cell.ghost_atoms().size()); + for (const LocalAtom& atom : cell.ghost_atoms()) + { + first_identity.push_back(atom.type_index); + } + for (LocalAtom& atom : cell.owned_atoms()) + { + atom.frac.x += 0.01; + atom.cart = atom.frac * cell.latvec(); + } + decomp.prepare_neighbors(cell); + const GhostExchangePlan& rebuilt_plan = cell.get_ghost_exchange_plan(); + ASSERT_EQ(rebuilt_plan.firstrecv.size(), rebuilt_plan.recvnum.size()); + if (domain.size() == 1) + { + EXPECT_EQ(rebuilt_plan.mapping.size(), cell.owned_atoms().size() + cell.ghost_atoms().size()); + } + ASSERT_EQ(cell.ghost_atoms().size(), first_identity.size()); + for (std::size_t i = 0; i < first_identity.size(); ++i) + { + EXPECT_EQ(cell.ghost_atoms()[i].type_index, first_identity[i]); + } } TEST_F(DomainDecompositionTest, GhostForcesAreReturnedAndConsumed) diff --git a/source/source_esolver/esolver_dp.cpp b/source/source_esolver/esolver_dp.cpp index ff7eb7a4693..92fbe51841a 100644 --- a/source/source_esolver/esolver_dp.cpp +++ b/source/source_esolver/esolver_dp.cpp @@ -20,7 +20,6 @@ #include "esolver_dp.h" #include "source_base/parallel_common.h" #include "source_base/timer.h" -#include "source_cell/domain_decomposition.h" #include "source_cell/mdcell.h" #include "source_cell/module_neighlist/neighbor_search.h" #include "source_cell/cif_io.h" @@ -30,7 +29,6 @@ #include #include #include -#include #include #include @@ -168,32 +166,9 @@ void ESolver_DP::runner(BaseCell& basecell, const int istep) ModuleBase::WARNING_QUIT("ESolver_DP", "MDCell ghost exchange plan exceeds int range."); } const bool use_mapping = ghost_plan.nprocs == 1; - std::vector mapping; - if (use_mapping) + if (use_mapping && ghost_plan.mapping.size() != static_cast(natom)) { - std::map, int> owned_atom_map; - for (int iat = 0; iat < nowned_atoms; ++iat) - { - const LocalAtom& atom = owned_atoms[static_cast(iat)]; - owned_atom_map[std::make_pair(atom.type, atom.type_index)] = iat; - } - mapping.resize(static_cast(natom), 0); - for (int iat = 0; iat < natom; ++iat) - { - if (iat < nowned_atoms) - { - mapping[static_cast(iat)] = iat; - continue; - } - const LocalAtom& atom = ghost_atoms[static_cast(iat - nowned_atoms)]; - const std::map, int>::const_iterator owner - = owned_atom_map.find(std::make_pair(atom.type, atom.type_index)); - if (owner == owned_atom_map.end()) - { - ModuleBase::WARNING_QUIT("ESolver_DP", "Cannot match a DeePMD ghost atom to its local owner."); - } - mapping[static_cast(iat)] = owner->second; - } + ModuleBase::WARNING_QUIT("ESolver_DP", "Invalid single-rank DeePMD atom mapping."); } const int nswap = static_cast(nswap_size); std::vector > sendlist_storage = ghost_plan.send_atom_indices; @@ -219,13 +194,14 @@ void ESolver_DP::runner(BaseCell& basecell, const int istep) int* ilist_data = nowned_atoms > 0 ? &ilist[0] : NULL; int* numneigh_data = nowned_atoms > 0 ? &numneigh[0] : NULL; int** firstneigh_data = nowned_atoms > 0 ? &firstneigh[0] : NULL; + std::vector mapping = ghost_plan.mapping; int* mapping_data = use_mapping && natom > 0 ? &mapping[0] : NULL; void* dp_world = NULL; #ifdef __MPI MPI_Comm dp_world_comm = mdcell.communicator(); dp_world = &dp_world_comm; #endif -#ifdef __DPMD_COMM_NLIST +#if defined(__DPMD_COMM_NLIST_ARGS) && __DPMD_COMM_NLIST_ARGS == 13 DPInputNlist nlist(nowned_atoms, ilist_data, numneigh_data, @@ -239,13 +215,26 @@ void ESolver_DP::runner(BaseCell& basecell, const int istep) recvproc, dp_world, ghost_plan.nprocs); +#elif defined(__DPMD_COMM_NLIST_ARGS) && __DPMD_COMM_NLIST_ARGS == 12 + DPInputNlist nlist(nowned_atoms, + ilist_data, + numneigh_data, + firstneigh_data, + nswap, + sendnum, + recvnum, + firstrecv, + sendlist_data, + sendproc, + recvproc, + dp_world); #else DPInputNlist nlist(nowned_atoms, ilist_data, numneigh_data, firstneigh_data); #endif -#if !defined(__DPMDC) || defined(__DPMD_COMM_NLIST) +#ifdef __DPMD_HAS_NLIST_MAPPING if (use_mapping) { nlist.set_mapping(mapping_data); diff --git a/tools/deepmd_input_nlist_probe.cpp b/tools/deepmd_input_nlist_probe.cpp new file mode 100644 index 00000000000..ab163119d70 --- /dev/null +++ b/tools/deepmd_input_nlist_probe.cpp @@ -0,0 +1,45 @@ +#if defined(__DPMDC) +#include +using DPProbeInputNlist = deepmd::hpp::InputNlist; +#else +#include +using DPProbeInputNlist = deepmd::InputNlist; +#endif + +int main() +{ +#if defined(DPMD_PROBE_COMM_13) + DPProbeInputNlist nlist(0, + nullptr, + nullptr, + nullptr, + 0, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + 1); +#elif defined(DPMD_PROBE_COMM_12) + DPProbeInputNlist nlist(0, + nullptr, + nullptr, + nullptr, + 0, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr); +#elif defined(DPMD_PROBE_MAPPING) + DPProbeInputNlist nlist(0, nullptr, nullptr, nullptr); + nlist.set_mapping(nullptr); +#else +#error "Select one DeePMD InputNlist probe" +#endif + return 0; +} From 79649902fcf329a97d1cea790cff9af298e8efde Mon Sep 17 00:00:00 2001 From: Fei Yang <2501213217@stu.pku.edu.cn> Date: Thu, 15 Oct 2026 17:30:15 +0800 Subject: [PATCH 3/5] refactor: validate DeepMD ghost exchange plans --- source/source_cell/ghost_exchange_plan.h | 13 ++++++++ source/source_esolver/esolver_dp.cpp | 42 +++--------------------- 2 files changed, 17 insertions(+), 38 deletions(-) diff --git a/source/source_cell/ghost_exchange_plan.h b/source/source_cell/ghost_exchange_plan.h index 8d9fe0efd2a..17f69ecaf53 100644 --- a/source/source_cell/ghost_exchange_plan.h +++ b/source/source_cell/ghost_exchange_plan.h @@ -1,6 +1,8 @@ #ifndef GHOST_EXCHANGE_PLAN_H #define GHOST_EXCHANGE_PLAN_H +#include +#include #include struct GhostExchangePlan @@ -18,6 +20,17 @@ struct GhostExchangePlan int nprocs; GhostExchangePlan() : nprocs(1) {} + + bool is_valid(const std::size_t natom) const + { + const std::size_t nswap = sendnum.size(); + const bool communication_sizes_match = recvnum.size() == nswap && firstrecv.size() == nswap + && send_atom_indices.size() == nswap + && sendproc.size() == nswap && recvproc.size() == nswap; + const bool nswap_fits_int = nswap <= static_cast(std::numeric_limits::max()); + const bool mapping_matches_atoms = nprocs != 1 || mapping.size() == natom; + return communication_sizes_match && nswap_fits_int && mapping_matches_atoms; + } }; #endif // GHOST_EXCHANGE_PLAN_H diff --git a/source/source_esolver/esolver_dp.cpp b/source/source_esolver/esolver_dp.cpp index 92fbe51841a..07363ceeaf3 100644 --- a/source/source_esolver/esolver_dp.cpp +++ b/source/source_esolver/esolver_dp.cpp @@ -28,7 +28,6 @@ #include #include -#include #include #include @@ -155,20 +154,10 @@ void ESolver_DP::runner(BaseCell& basecell, const int istep) const GhostExchangePlan& ghost_plan = mdcell.get_ghost_exchange_plan(); const std::size_t nswap_size = ghost_plan.sendnum.size(); - if (ghost_plan.recvnum.size() != nswap_size || ghost_plan.firstrecv.size() != nswap_size - || ghost_plan.send_atom_indices.size() != nswap_size - || ghost_plan.sendproc.size() != nswap_size || ghost_plan.recvproc.size() != nswap_size) - { - ModuleBase::WARNING_QUIT("ESolver_DP", "Invalid MDCell ghost exchange plan."); - } - if (nswap_size > static_cast(std::numeric_limits::max())) - { - ModuleBase::WARNING_QUIT("ESolver_DP", "MDCell ghost exchange plan exceeds int range."); - } const bool use_mapping = ghost_plan.nprocs == 1; - if (use_mapping && ghost_plan.mapping.size() != static_cast(natom)) + if (!ghost_plan.is_valid(static_cast(natom))) { - ModuleBase::WARNING_QUIT("ESolver_DP", "Invalid single-rank DeePMD atom mapping."); + ModuleBase::WARNING_QUIT("ESolver_DP", "Invalid MDCell ghost exchange plan."); } const int nswap = static_cast(nswap_size); std::vector > sendlist_storage = ghost_plan.send_atom_indices; @@ -202,32 +191,9 @@ void ESolver_DP::runner(BaseCell& basecell, const int istep) dp_world = &dp_world_comm; #endif #if defined(__DPMD_COMM_NLIST_ARGS) && __DPMD_COMM_NLIST_ARGS == 13 - DPInputNlist nlist(nowned_atoms, - ilist_data, - numneigh_data, - firstneigh_data, - nswap, - sendnum, - recvnum, - firstrecv, - sendlist_data, - sendproc, - recvproc, - dp_world, - ghost_plan.nprocs); + DPInputNlist nlist(nowned_atoms, ilist_data, numneigh_data, firstneigh_data, nswap, sendnum, recvnum, firstrecv, sendlist_data, sendproc, recvproc, dp_world, ghost_plan.nprocs); #elif defined(__DPMD_COMM_NLIST_ARGS) && __DPMD_COMM_NLIST_ARGS == 12 - DPInputNlist nlist(nowned_atoms, - ilist_data, - numneigh_data, - firstneigh_data, - nswap, - sendnum, - recvnum, - firstrecv, - sendlist_data, - sendproc, - recvproc, - dp_world); + DPInputNlist nlist(nowned_atoms, ilist_data, numneigh_data, firstneigh_data, nswap, sendnum, recvnum, firstrecv, sendlist_data, sendproc, recvproc, dp_world); #else DPInputNlist nlist(nowned_atoms, ilist_data, From 039322e5dfc79862e10173ac77c219e637a5a425 Mon Sep 17 00:00:00 2001 From: Fei Yang <2501213217@stu.pku.edu.cn> Date: Thu, 15 Oct 2026 18:13:32 +0800 Subject: [PATCH 4/5] test: sync evolve psi test with output info API --- source/source_lcao/module_rt/test/CMakeLists.txt | 1 + source/source_lcao/module_rt/test/test_evolve_psi.cpp | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/source/source_lcao/module_rt/test/CMakeLists.txt b/source/source_lcao/module_rt/test/CMakeLists.txt index 57114b79d83..86a5a955d35 100644 --- a/source/source_lcao/module_rt/test/CMakeLists.txt +++ b/source/source_lcao/module_rt/test/CMakeLists.txt @@ -19,6 +19,7 @@ if(NOT ENABLE_CUBLASMP) ../band_energy.cpp ../upsi.cpp ../solve_propagation.cpp ../propagator.cpp ../propagator_cn2.cpp ../propagator_taylor.cpp ../propagator_etrs.cpp ../../../source_basis/module_ao/parallel_orbitals.cpp + ../../../source_io/module_energy/td_eij_io.cpp ) endif() diff --git a/source/source_lcao/module_rt/test/test_evolve_psi.cpp b/source/source_lcao/module_rt/test/test_evolve_psi.cpp index 942a48c70c4..41afa734cb4 100644 --- a/source/source_lcao/module_rt/test/test_evolve_psi.cpp +++ b/source/source_lcao/module_rt/test/test_evolve_psi.cpp @@ -35,6 +35,7 @@ class EvolvePsiTest : public testing::Test EndpointHamiltonian hamiltonian; std::array previous_s = {{1.0, 0.0, 0.0, 1.0}}; std::array previous_psi = {{1.0, Complex(0.0, 0.5)}}; + const module_rt::EijOutputInfo output_info{0.0, 0, 0, 0, 0, 0}; std::ofstream log; void SetUp() override @@ -83,7 +84,7 @@ class EvolvePsiTest : public testing::Test const bool use_lapack = true; module_rt::evolve_psi_tensor( 1, 2, &pv, &hamiltonian, psi, last_psi, last_h, last_s, energy, 0, log, 0, - use_lapack, resources); + use_lapack, resources, output_info); expect_endpoint_norm(psi.data()); EXPECT_NEAR(energy.data()[0], 0.0, 1e-12); } @@ -95,7 +96,7 @@ TEST_F(EvolvePsiTest, CpuUsesEndpointOverlap) std::array last_h = {{0.0, 0.0, 0.0, 0.0}}; double energy = 0.0; module_rt::evolve_psi(1, 2, &pv, &hamiltonian, psi.data(), previous_psi.data(), - last_h.data(), previous_s.data(), nullptr, false, &energy, 0, log, 0); + last_h.data(), previous_s.data(), nullptr, false, &energy, 0, log, 0, output_info); expect_endpoint_norm(psi.data()); EXPECT_NEAR(energy, 0.0, 1e-12); } From 447ed27150f5c814f3e6d4b079d6395d225a75a1 Mon Sep 17 00:00:00 2001 From: Fei Yang <2501213217@stu.pku.edu.cn> Date: Thu, 15 Oct 2026 19:12:40 +0800 Subject: [PATCH 5/5] refactor: reduce DeePMD plan handling in solver --- source/source_cell/domain_decomposition.cpp | 25 ++++++++++++++ source/source_cell/ghost_exchange_plan.h | 14 ++++++++ source/source_cell/mdcell.cpp | 9 +++++ source/source_cell/mdcell.h | 1 + .../test/test_domain_decomposition.cpp | 16 ++++++++- source/source_esolver/esolver_dp.cpp | 34 ++++--------------- 6 files changed, 71 insertions(+), 28 deletions(-) diff --git a/source/source_cell/domain_decomposition.cpp b/source/source_cell/domain_decomposition.cpp index d2628977a4f..8565e8c6a10 100644 --- a/source/source_cell/domain_decomposition.cpp +++ b/source/source_cell/domain_decomposition.cpp @@ -43,6 +43,31 @@ DomainDecomposition::~DomainDecomposition() #endif } +GhostExchangePlanView GhostExchangePlan::make_view() +{ + GhostExchangePlanView view; + view.nswap = static_cast(sendnum.size()); + view.sendnum = sendnum.empty() ? NULL : sendnum.data(); + view.recvnum = recvnum.empty() ? NULL : recvnum.data(); + view.firstrecv = firstrecv.empty() ? NULL : firstrecv.data(); + view.sendproc = sendproc.empty() ? NULL : sendproc.data(); + view.recvproc = recvproc.empty() ? NULL : recvproc.data(); + view.mapping = mapping.empty() ? NULL : mapping.data(); + view.sendlist_storage.resize(send_atom_indices.size()); + for (std::size_t iswap = 0; iswap < send_atom_indices.size(); ++iswap) + { + if (!send_atom_indices[iswap].empty()) + { + view.sendlist_storage[iswap] = send_atom_indices[iswap].data(); + } + else + { + view.sendlist_storage[iswap] = NULL; + } + } + return view; +} + void DomainDecomposition::init(const ModuleBase::CommunicationDomain& comm_domain, const ModuleBase::Matrix3& latvec, double lat0, diff --git a/source/source_cell/ghost_exchange_plan.h b/source/source_cell/ghost_exchange_plan.h index 17f69ecaf53..422e9f91622 100644 --- a/source/source_cell/ghost_exchange_plan.h +++ b/source/source_cell/ghost_exchange_plan.h @@ -5,6 +5,18 @@ #include #include +struct GhostExchangePlanView +{ + int nswap; + int* sendnum; + int* recvnum; + int* firstrecv; + int* sendproc; + int* recvproc; + int* mapping; + std::vector sendlist_storage; +}; + struct GhostExchangePlan { std::vector sendnum; @@ -21,6 +33,8 @@ struct GhostExchangePlan GhostExchangePlan() : nprocs(1) {} + GhostExchangePlanView make_view(); + bool is_valid(const std::size_t natom) const { const std::size_t nswap = sendnum.size(); diff --git a/source/source_cell/mdcell.cpp b/source/source_cell/mdcell.cpp index dcc64a09e4f..486f539f9c7 100644 --- a/source/source_cell/mdcell.cpp +++ b/source/source_cell/mdcell.cpp @@ -52,6 +52,15 @@ const GhostExchangePlan& MDCell::get_ghost_exchange_plan() const return ghost_exchange_plan_; } +GhostExchangePlan& MDCell::get_ghost_exchange_plan() +{ + if (!ghost_exchange_plan_valid_) + { + ModuleBase::WARNING_QUIT("MDCell", "Ghost exchange plan is not available."); + } + return ghost_exchange_plan_; +} + void MDCell::clear_ghost_exchange_plan_() { ghost_exchange_plan_ = GhostExchangePlan(); diff --git a/source/source_cell/mdcell.h b/source/source_cell/mdcell.h index ebf891311ec..04fa383e26d 100644 --- a/source/source_cell/mdcell.h +++ b/source/source_cell/mdcell.h @@ -67,6 +67,7 @@ class MDCell : public BaseCell const std::vector& type_masses() const { return type_masses_; } const std::vector& type_atom_counts() const { return type_atom_counts_; } const GhostExchangePlan& get_ghost_exchange_plan() const; + GhostExchangePlan& get_ghost_exchange_plan(); StruMeta& mutable_stru_meta() { return stru_meta_; } const StruMeta& stru_meta() const { return stru_meta_; } int nowned_atoms() const { return static_cast(owned_atoms_.size()); } diff --git a/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp b/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp index b1c0f8d83fd..d615f3856c7 100644 --- a/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp +++ b/source/source_cell/module_neighlist/test/test_domain_decomposition.cpp @@ -83,13 +83,27 @@ TEST_F(DomainDecompositionTest, CutoffSetterDoesNotExchangeAtoms) TEST_F(DomainDecompositionTest, GhostExchangePlanMatchesGhostLayout) { decomp.prepare_neighbors(cell); - const GhostExchangePlan& plan = cell.get_ghost_exchange_plan(); + GhostExchangePlan& plan = cell.get_ghost_exchange_plan(); const std::size_t exchange_count = plan.sendnum.size(); ASSERT_EQ(plan.recvnum.size(), exchange_count); ASSERT_EQ(plan.firstrecv.size(), exchange_count); ASSERT_EQ(plan.send_atom_indices.size(), exchange_count); ASSERT_EQ(plan.sendproc.size(), exchange_count); ASSERT_EQ(plan.recvproc.size(), exchange_count); + GhostExchangePlanView view = plan.make_view(); + EXPECT_EQ(view.nswap, static_cast(exchange_count)); + ASSERT_EQ(view.sendlist_storage.size(), exchange_count); + for (std::size_t iswap = 0; iswap < exchange_count; ++iswap) + { + if (plan.send_atom_indices[iswap].empty()) + { + EXPECT_EQ(view.sendlist_storage[iswap], static_cast(NULL)); + } + else + { + EXPECT_EQ(view.sendlist_storage[iswap], plan.send_atom_indices[iswap].data()); + } + } EXPECT_EQ(plan.nprocs, domain.size()); if (domain.size() == 1) { diff --git a/source/source_esolver/esolver_dp.cpp b/source/source_esolver/esolver_dp.cpp index 07363ceeaf3..98708ee051e 100644 --- a/source/source_esolver/esolver_dp.cpp +++ b/source/source_esolver/esolver_dp.cpp @@ -20,6 +20,7 @@ #include "esolver_dp.h" #include "source_base/parallel_common.h" #include "source_base/timer.h" +#include "source_cell/ghost_exchange_plan.h" #include "source_cell/mdcell.h" #include "source_cell/module_neighlist/neighbor_search.h" #include "source_cell/cif_io.h" @@ -152,48 +153,27 @@ void ESolver_DP::runner(BaseCell& basecell, const int istep) : sorted_neighbors[static_cast(iat)].data(); } - const GhostExchangePlan& ghost_plan = mdcell.get_ghost_exchange_plan(); - const std::size_t nswap_size = ghost_plan.sendnum.size(); + GhostExchangePlan& ghost_plan = mdcell.get_ghost_exchange_plan(); const bool use_mapping = ghost_plan.nprocs == 1; if (!ghost_plan.is_valid(static_cast(natom))) { ModuleBase::WARNING_QUIT("ESolver_DP", "Invalid MDCell ghost exchange plan."); } - const int nswap = static_cast(nswap_size); - std::vector > sendlist_storage = ghost_plan.send_atom_indices; - std::vector sendlist(nswap_size, NULL); - for (std::size_t iswap = 0; iswap < nswap_size; ++iswap) - { - if (!sendlist_storage[iswap].empty()) - { - sendlist[iswap] = sendlist_storage[iswap].data(); - } - } - std::vector sendnum_data = ghost_plan.sendnum; - std::vector recvnum_data = ghost_plan.recvnum; - std::vector firstrecv_data = ghost_plan.firstrecv; - std::vector sendproc_data = ghost_plan.sendproc; - std::vector recvproc_data = ghost_plan.recvproc; - int* sendnum = sendnum_data.empty() ? NULL : &sendnum_data[0]; - int* recvnum = recvnum_data.empty() ? NULL : &recvnum_data[0]; - int* firstrecv = firstrecv_data.empty() ? NULL : &firstrecv_data[0]; - int** sendlist_data = sendlist.empty() ? NULL : &sendlist[0]; - int* sendproc = sendproc_data.empty() ? NULL : &sendproc_data[0]; - int* recvproc = recvproc_data.empty() ? NULL : &recvproc_data[0]; + GhostExchangePlanView view = ghost_plan.make_view(); + int** sendlist_data = view.sendlist_storage.empty() ? NULL : view.sendlist_storage.data(); int* ilist_data = nowned_atoms > 0 ? &ilist[0] : NULL; int* numneigh_data = nowned_atoms > 0 ? &numneigh[0] : NULL; int** firstneigh_data = nowned_atoms > 0 ? &firstneigh[0] : NULL; - std::vector mapping = ghost_plan.mapping; - int* mapping_data = use_mapping && natom > 0 ? &mapping[0] : NULL; + int* mapping_data = use_mapping ? view.mapping : NULL; void* dp_world = NULL; #ifdef __MPI MPI_Comm dp_world_comm = mdcell.communicator(); dp_world = &dp_world_comm; #endif #if defined(__DPMD_COMM_NLIST_ARGS) && __DPMD_COMM_NLIST_ARGS == 13 - DPInputNlist nlist(nowned_atoms, ilist_data, numneigh_data, firstneigh_data, nswap, sendnum, recvnum, firstrecv, sendlist_data, sendproc, recvproc, dp_world, ghost_plan.nprocs); + DPInputNlist nlist(nowned_atoms, ilist_data, numneigh_data, firstneigh_data, view.nswap, view.sendnum, view.recvnum, view.firstrecv, sendlist_data, view.sendproc, view.recvproc, dp_world, ghost_plan.nprocs); #elif defined(__DPMD_COMM_NLIST_ARGS) && __DPMD_COMM_NLIST_ARGS == 12 - DPInputNlist nlist(nowned_atoms, ilist_data, numneigh_data, firstneigh_data, nswap, sendnum, recvnum, firstrecv, sendlist_data, sendproc, recvproc, dp_world); + DPInputNlist nlist(nowned_atoms, ilist_data, numneigh_data, firstneigh_data, view.nswap, view.sendnum, view.recvnum, view.firstrecv, sendlist_data, view.sendproc, view.recvproc, dp_world); #else DPInputNlist nlist(nowned_atoms, ilist_data,