fix: pass MDCell ghost communication plan to DeePMD

cmake/modules/FindDeePMD.cmake

@@ -64,4 +64,71 @@ 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)
+        set(_abacus_dp_nlist_type "deepmd::hpp::InputNlist")
+        set(_abacus_dp_nlist_header "#include <deepmd/deepmd.hpp>")
+    else()
+        set(_abacus_dp_nlist_type "deepmd::InputNlist")
+        set(_abacus_dp_nlist_header "#include <deepmd/DeepPot.h>")
+    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(
+            "${_abacus_dp_nlist_header}
+             int main()
+             {
+                 ${_abacus_dp_nlist_type} nlist(0, nullptr, nullptr, nullptr,
+                                                0, nullptr, nullptr, nullptr,
+                                                nullptr, nullptr, nullptr, nullptr);
+                 return 0;
+             }"
+            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 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.")
+endif()
+
 mark_as_advanced(DeePMD_INCLUDE_DIR deepmd_c deepmd_cc)

source/CMakeLists.txt

@@ -74,6 +74,12 @@ if(ABACUS_COMMIT_INFO_INCLUDE_DIR)
 endif()
 if(DEFINED DeePMD_DIR)
   list(APPEND _abacus_feature_definitions __DPMD HIGH_PREC)
+  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)
   endif()

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

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"
@@ -42,6 +43,31 @@ DomainDecomposition::~DomainDecomposition()
 #endif
 }
 
+GhostExchangePlanView GhostExchangePlan::make_view()
+{
+    GhostExchangePlanView view;
+    view.nswap = static_cast<int>(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,
@@ -176,6 +202,7 @@ void DomainDecomposition::exchange_ghost_atoms(MDCell& cell)
     std::vector<LocalAtom>& 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<GhostExchangeSlot>& slots = ghost_slots_;
@@ -311,10 +338,82 @@ 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::size_t>(std::numeric_limits<int>::max()))
+        {
+            throw std::overflow_error("Ghost send list exceeds int range.");
+        }
+        if (slot.ghost_begin > static_cast<std::size_t>(std::numeric_limits<int>::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<int>(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::size_t>(std::numeric_limits<int>::max()))
+        {
+            throw std::overflow_error("Ghost receive index exceeds int range.");
+        }
+        plan.firstrecv.push_back(static_cast<int>(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);
+    }
+
+    if (plan.nprocs == 1)
+    {
+        const std::size_t natom = cell.owned_atoms_.size() + cell.ghost_atoms_.size();
+        plan.mapping.resize(natom, 0);
+        std::map<std::pair<int, std::int64_t>, 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<int>(iat);
+            plan.mapping[iat] = static_cast<int>(iat);
+        }
+        for (std::size_t iat = 0; iat < cell.ghost_atoms_.size(); ++iat)
+        {
+            const LocalAtom& atom = cell.ghost_atoms_[iat];
+            const std::map<std::pair<int, std::int64_t>, 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;
+}
+
 void DomainDecomposition::update_ghost_atom_positions(MDCell& cell)
 {
     synchronize_geometry_(cell);

source/source_cell/domain_decomposition.h

@@ -15,6 +15,7 @@
 
 class UnitCell;
 class MDCell;
+struct GhostExchangePlan;
 namespace ModuleBase
 {
 class CommunicationDomain;
@@ -117,6 +118,7 @@ class DomainDecomposition
                                               const std::array<int, 3>& previous_image_shift,
                                               const GhostExchangeSlot& slot) const;
     void build_ghost_exchange_slots(std::vector<GhostExchangeSlot>& slots) const;
+    GhostExchangePlan build_ghost_exchange_plan(const MDCell& cell) const;
     PackedAtom pack_atom(const LocalAtom& atom, const std::array<int, 3>& image_shift) const;
     LocalAtom unpack_ghost_atom(const PackedAtom& packed) const;
     LocalAtom unpack_owned_atom(const PackedAtom& packed) const;

source/source_cell/ghost_exchange_plan.h

@@ -0,0 +1,50 @@
+#ifndef GHOST_EXCHANGE_PLAN_H
+#define GHOST_EXCHANGE_PLAN_H
+
+#include <cstddef>
+#include <limits>
+#include <vector>
+
+struct GhostExchangePlanView
+{
+    int nswap;
+    int* sendnum;
+    int* recvnum;
+    int* firstrecv;
+    int* sendproc;
+    int* recvproc;
+    int* mapping;
+    std::vector<int*> sendlist_storage;
+};
+
+struct GhostExchangePlan
+{
+    std::vector<int> sendnum;
+    std::vector<int> recvnum;
+    std::vector<int> firstrecv;
+    std::vector<std::vector<int> > send_atom_indices;
+    std::vector<int> sendproc;
+    std::vector<int> 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<int> mapping;
+    int nprocs;
+
+    GhostExchangePlan() : nprocs(1) {}
+
+    GhostExchangePlanView make_view();
+
+    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::size_t>(std::numeric_limits<int>::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

source/source_cell/mdcell.cpp

@@ -43,6 +43,36 @@ 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_;
+}
+
+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();
+    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 +115,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();

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,8 @@ class MDCell : public BaseCell
     const std::vector<std::string>& type_labels() const { return type_labels_; }
     const std::vector<double>& type_masses() const { return type_masses_; }
     const std::vector<std::int64_t>& 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<int>(owned_atoms_.size()); }
@@ -91,6 +94,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<LocalAtom>& atoms);
     static double wrap_fractional_(double value);
@@ -113,6 +118,8 @@ class MDCell : public BaseCell
     std::unique_ptr<NeighborSearch> neighbor_search_;
     std::vector<ModuleBase::Vector3<double> > 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;