update fix, ref to deepmodeling#5268

Author: iProzd

source/api_cc/include/DeepPotPT.h

@@ -340,6 +340,10 @@ class DeepPotPT : public DeepPotBackend {
   at::Tensor firstneigh_tensor;
   c10::optional<torch::Tensor> mapping_tensor;
   torch::Dict<std::string, torch::Tensor> comm_dict;
+  std::vector<std::vector<int>> remapped_sendlist;
+  std::vector<int*> remapped_sendlist_ptrs;
+  std::vector<int> remapped_sendnum;
+  std::vector<int> remapped_recvnum;
   bool profiler_enabled{false};
   std::string profiler_file;
   /**

source/api_cc/include/DeepSpinPT.h

@@ -262,6 +262,10 @@ class DeepSpinPT : public DeepSpinBackend {
   at::Tensor firstneigh_tensor;
   c10::optional<torch::Tensor> mapping_tensor;
   torch::Dict<std::string, torch::Tensor> comm_dict;
+  std::vector<std::vector<int>> remapped_sendlist;
+  std::vector<int*> remapped_sendlist_ptrs;
+  std::vector<int> remapped_sendnum;
+  std::vector<int> remapped_recvnum;
   /**
    * @brief Translate PyTorch exceptions to the DeePMD-kit exception.
    * @param[in] f The function to run.

source/api_cc/include/common.h

@@ -153,6 +153,27 @@ void select_map_inv(typename std::vector<VT>::iterator out,
                     const std::vector<int>& fwd_map,
                     const int& stride);
 
+/**
+ * @brief Remap communication sendlist for message passing with NULL-type atoms.
+ *
+ * When NULL-type (virtual) atoms are present, the original LAMMPS sendlist
+ * contains indices referring to virtual atoms that have been filtered out by
+ * select_real_atoms_coord. This function remaps those indices through fwd_map
+ * and independently recomputes recvnum using firstrecv.
+ *
+ * @param[out] new_sendlist Remapped send lists per swap (vector of vectors).
+ * @param[out] new_sendnum Number of atoms to send per swap after remapping.
+ * @param[out] new_recvnum Number of atoms to receive per swap after remapping.
+ * @param[in] lmp_list The LAMMPS neighbor list containing communication info.
+ * @param[in] fwd_map Forward map from original atom index to real-atom index
+ *            (-1 for virtual/NULL atoms).
+ */
+void remap_comm_sendlist(std::vector<std::vector<int>>& new_sendlist,
+                         std::vector<int>& new_sendnum,
+                         std::vector<int>& new_recvnum,
+                         const InputNlist& lmp_list,
+                         const std::vector<int>& fwd_map);
+
 /**
  * @brief Get the number of threads from the environment variable.
  * @details A warning will be thrown if environment variables are not set.

source/api_cc/include/commonPT.h

@@ -0,0 +1,106 @@
+// SPDX-License-Identifier: LGPL-3.0-or-later
+#pragma once
+
+#ifdef BUILD_PYTORCH
+#include <torch/torch.h>
+
+#include <cstdint>
+#include <vector>
+
+#include "common.h"
+#include "neighbor_list.h"
+
+namespace deepmd {
+
+/**
+ * @brief Build comm_dict tensors from sendlist/sendnum/recvnum buffers.
+ *
+ * This is the shared tensor-building logic for all PyTorch backends
+ * (DeepPotPT, DeepSpinPT). Backend-specific entries (e.g. has_spin)
+ * should be added by the caller after this function returns.
+ *
+ * @param[out] comm_dict The communication dictionary to populate.
+ * @param[in] lmp_list The LAMMPS neighbor list (for sendproc/recvproc/world).
+ * @param[in] sendlist Pointer array (int**) for each swap's send list.
+ * @param[in] sendnum Number of send atoms per swap.
+ * @param[in] recvnum Number of recv atoms per swap.
+ */
+inline void build_comm_dict(torch::Dict<std::string, torch::Tensor>& comm_dict,
+                            const InputNlist& lmp_list,
+                            int** sendlist,
+                            int* sendnum,
+                            int* recvnum) {
+  int nswap = lmp_list.nswap;
+  auto int32_option =
+      torch::TensorOptions().device(torch::kCPU).dtype(torch::kInt32);
+  auto int64_option =
+      torch::TensorOptions().device(torch::kCPU).dtype(torch::kInt64);
+
+  torch::Tensor sendlist_tensor =
+      torch::from_blob(static_cast<void*>(sendlist), {nswap}, int64_option);
+  torch::Tensor sendnum_tensor =
+      torch::from_blob(sendnum, {nswap}, int32_option);
+  torch::Tensor recvnum_tensor =
+      torch::from_blob(recvnum, {nswap}, int32_option);
+  torch::Tensor sendproc_tensor =
+      torch::from_blob(lmp_list.sendproc, {nswap}, int32_option);
+  torch::Tensor recvproc_tensor =
+      torch::from_blob(lmp_list.recvproc, {nswap}, int32_option);
+
+  torch::Tensor communicator_tensor;
+  static std::int64_t null_communicator = 0;
+  if (lmp_list.world == nullptr) {
+    communicator_tensor =
+        torch::from_blob(&null_communicator, {1}, torch::kInt64);
+  } else {
+    communicator_tensor =
+        torch::from_blob(const_cast<void*>(lmp_list.world), {1}, torch::kInt64);
+  }
+
+  comm_dict.insert_or_assign("send_list", sendlist_tensor);
+  comm_dict.insert_or_assign("send_proc", sendproc_tensor);
+  comm_dict.insert_or_assign("recv_proc", recvproc_tensor);
+  comm_dict.insert_or_assign("send_num", sendnum_tensor);
+  comm_dict.insert_or_assign("recv_num", recvnum_tensor);
+  comm_dict.insert_or_assign("communicator", communicator_tensor);
+}
+
+/**
+ * @brief Build comm_dict with sendlist remapping for virtual (NULL-type) atoms.
+ *
+ * Calls remap_comm_sendlist() to remap indices through fwd_map, then
+ * build_comm_dict() to create tensors. Backend-specific entries (e.g.
+ * has_spin) should be added by the caller after this function returns.
+ *
+ * @param[out] comm_dict The communication dictionary to populate.
+ * @param[in] lmp_list The LAMMPS neighbor list containing communication info.
+ * @param[in] fwd_map Map from original atom index to real-atom index (-1 for
+ *            virtual atoms).
+ * @param[out] remapped_sendlist Storage for remapped send lists (kept alive for
+ *             tensor lifetime).
+ * @param[out] remapped_sendlist_ptrs Pointer array into remapped_sendlist.
+ * @param[out] remapped_sendnum Remapped send counts per swap.
+ * @param[out] remapped_recvnum Remapped recv counts per swap.
+ */
+inline void build_comm_dict_with_virtual_atoms(
+    torch::Dict<std::string, torch::Tensor>& comm_dict,
+    const InputNlist& lmp_list,
+    const std::vector<int>& fwd_map,
+    std::vector<std::vector<int>>& remapped_sendlist,
+    std::vector<int*>& remapped_sendlist_ptrs,
+    std::vector<int>& remapped_sendnum,
+    std::vector<int>& remapped_recvnum) {
+  remap_comm_sendlist(remapped_sendlist, remapped_sendnum, remapped_recvnum,
+                      lmp_list, fwd_map);
+  int nswap = lmp_list.nswap;
+  remapped_sendlist_ptrs.resize(nswap);
+  for (int s = 0; s < nswap; ++s) {
+    remapped_sendlist_ptrs[s] = remapped_sendlist[s].data();
+  }
+  build_comm_dict(comm_dict, lmp_list, remapped_sendlist_ptrs.data(),
+                  remapped_sendnum.data(), remapped_recvnum.data());
+}
+
+}  // namespace deepmd
+
+#endif  // BUILD_PYTORCH

source/api_cc/src/DeepPotPT.cc

@@ -8,6 +8,7 @@
 #include <cstdint>
 
 #include "common.h"
+#include "commonPT.h"
 #include "device.h"
 #include "errors.h"
 
@@ -198,6 +199,8 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
                           bkw_map, nall_real, nloc_real, coord, atype, aparam,
                           nghost, ntypes, 1, daparam, nall, aparam_nall);
   int nloc = nall_real - nghost_real;
+  // Detect whether any NULL-type atoms were filtered out.
+  bool has_null_atoms = (nall_real < nall);
   int nframes = 1;
   std::vector<VALUETYPE> coord_wrapped = dcoord;
   at::Tensor coord_wrapped_Tensor =
@@ -211,36 +214,14 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
     nlist_data.shuffle_exclude_empty(fwd_map);
     nlist_data.padding();
     if (do_message_passing) {
-      int nswap = lmp_list.nswap;
-      torch::Tensor sendproc_tensor =
-          torch::from_blob(lmp_list.sendproc, {nswap}, int32_option);
-      torch::Tensor recvproc_tensor =
-          torch::from_blob(lmp_list.recvproc, {nswap}, int32_option);
-      torch::Tensor firstrecv_tensor =
-          torch::from_blob(lmp_list.firstrecv, {nswap}, int32_option);
-      torch::Tensor recvnum_tensor =
-          torch::from_blob(lmp_list.recvnum, {nswap}, int32_option);
-      torch::Tensor sendnum_tensor =
-          torch::from_blob(lmp_list.sendnum, {nswap}, int32_option);
-      torch::Tensor communicator_tensor;
-      if (lmp_list.world == 0) {
-        communicator_tensor = torch::empty({1}, torch::kInt64);
+      if (has_null_atoms) {
+        build_comm_dict_with_virtual_atoms(
+            comm_dict, lmp_list, fwd_map, remapped_sendlist,
+            remapped_sendlist_ptrs, remapped_sendnum, remapped_recvnum);
       } else {
-        communicator_tensor = torch::from_blob(
-            const_cast<void*>(lmp_list.world), {1}, torch::kInt64);
+        build_comm_dict(comm_dict, lmp_list, lmp_list.sendlist,
+                        lmp_list.sendnum, lmp_list.recvnum);
       }
-
-      torch::Tensor nswap_tensor = torch::tensor(nswap, int32_option);
-      int total_send =
-          std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0);
-      torch::Tensor sendlist_tensor =
-          torch::from_blob(lmp_list.sendlist, {total_send}, int32_option);
-      comm_dict.insert_or_assign("send_list", sendlist_tensor);
-      comm_dict.insert_or_assign("send_proc", sendproc_tensor);
-      comm_dict.insert_or_assign("recv_proc", recvproc_tensor);
-      comm_dict.insert_or_assign("send_num", sendnum_tensor);
-      comm_dict.insert_or_assign("recv_num", recvnum_tensor);
-      comm_dict.insert_or_assign("communicator", communicator_tensor);
     }
     if (lmp_list.mapping) {
       std::vector<std::int64_t> mapping(nall_real);
@@ -272,7 +253,7 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
             options)
             .to(device);
   }
-  c10::Dict<c10::IValue, c10::IValue> outputs =
+  auto outputs =
       (do_message_passing)
           ? module
                 .run_method("forward_lower", coord_wrapped_Tensor, atype_Tensor,

source/api_cc/src/DeepSpinPT.cc

@@ -7,6 +7,7 @@
 #include <cstdint>
 
 #include "common.h"
+#include "commonPT.h"
 #include "device.h"
 #include "errors.h"
 
@@ -163,6 +164,8 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener,
                           bkw_map, nall_real, nloc_real, coord, atype, aparam,
                           nghost, ntypes, 1, daparam, nall, aparam_nall);
   int nloc = nall_real - nghost_real;
+  // Detect whether any NULL-type atoms were filtered out.
+  bool has_null_atoms = (nall_real < nall);
   int nframes = 1;
   std::vector<VALUETYPE> coord_wrapped = dcoord;
   at::Tensor coord_wrapped_Tensor =
@@ -175,43 +178,32 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener,
   std::vector<std::int64_t> atype_64(datype.begin(), datype.end());
   at::Tensor atype_Tensor =
       torch::from_blob(atype_64.data(), {1, nall_real}, int_option).to(device);
-  c10::optional<torch::Tensor> mapping_tensor;
   if (ago == 0) {
     nlist_data.copy_from_nlist(lmp_list, nall - nghost);
     nlist_data.shuffle_exclude_empty(fwd_map);
     nlist_data.padding();
     if (do_message_passing) {
-      int nswap = lmp_list.nswap;
-      torch::Tensor sendproc_tensor =
-          torch::from_blob(lmp_list.sendproc, {nswap}, int32_option);
-      torch::Tensor recvproc_tensor =
-          torch::from_blob(lmp_list.recvproc, {nswap}, int32_option);
-      torch::Tensor firstrecv_tensor =
-          torch::from_blob(lmp_list.firstrecv, {nswap}, int32_option);
-      torch::Tensor recvnum_tensor =
-          torch::from_blob(lmp_list.recvnum, {nswap}, int32_option);
-      torch::Tensor sendnum_tensor =
-          torch::from_blob(lmp_list.sendnum, {nswap}, int32_option);
-      torch::Tensor communicator_tensor;
-      if (lmp_list.world == 0) {
-        communicator_tensor = torch::empty({1}, torch::kInt64);
+      if (has_null_atoms) {
+        build_comm_dict_with_virtual_atoms(
+            comm_dict, lmp_list, fwd_map, remapped_sendlist,
+            remapped_sendlist_ptrs, remapped_sendnum, remapped_recvnum);
       } else {
-        communicator_tensor = torch::from_blob(
-            const_cast<void*>(lmp_list.world), {1}, torch::kInt64);
+        build_comm_dict(comm_dict, lmp_list, lmp_list.sendlist,
+                        lmp_list.sendnum, lmp_list.recvnum);
       }
-      torch::Tensor nswap_tensor = torch::tensor(nswap, int32_option);
-      int total_send =
-          std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0);
-      torch::Tensor sendlist_tensor =
-          torch::from_blob(lmp_list.sendlist, {total_send}, int32_option);
-      torch::Tensor has_spin = torch::tensor({1}, int32_option);
-      comm_dict.insert_or_assign("send_list", sendlist_tensor);
-      comm_dict.insert_or_assign("send_proc", sendproc_tensor);
-      comm_dict.insert_or_assign("recv_proc", recvproc_tensor);
-      comm_dict.insert_or_assign("send_num", sendnum_tensor);
-      comm_dict.insert_or_assign("recv_num", recvnum_tensor);
-      comm_dict.insert_or_assign("communicator", communicator_tensor);
-      comm_dict.insert_or_assign("has_spin", has_spin);
+      // DeepSpin-specific: signal spin model to the Python side
+      auto int32_option =
+          torch::TensorOptions().device(torch::kCPU).dtype(torch::kInt32);
+      comm_dict.insert_or_assign("has_spin", torch::tensor({1}, int32_option));
+    }
+    if (lmp_list.mapping) {
+      std::vector<std::int64_t> mapping(nall_real);
+      for (size_t ii = 0; ii < nall_real; ii++) {
+        mapping[ii] = lmp_list.mapping[fwd_map[ii]];
+      }
+      mapping_tensor =
+          torch::from_blob(mapping.data(), {1, nall_real}, int_option)
+              .to(device);
     }
   }
   at::Tensor firstneigh = createNlistTensor2(nlist_data.jlist);
@@ -234,7 +226,7 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener,
             options)
             .to(device);
   }
-  c10::Dict<c10::IValue, c10::IValue> outputs =
+  auto outputs =
       (do_message_passing)
           ? module
                 .run_method("forward_lower", coord_wrapped_Tensor, atype_Tensor,

source/api_cc/src/common.cc

@@ -232,6 +232,45 @@ template void deepmd::select_real_atoms_coord<float>(
     const int& nall,
     const bool aparam_nall);
 
+void deepmd::remap_comm_sendlist(std::vector<std::vector<int>>& new_sendlist,
+                                 std::vector<int>& new_sendnum,
+                                 std::vector<int>& new_recvnum,
+                                 const InputNlist& lmp_list,
+                                 const std::vector<int>& fwd_map) {
+  int nswap = lmp_list.nswap;
+  new_sendlist.resize(nswap);
+  new_sendnum.resize(nswap);
+  new_recvnum.resize(nswap);
+
+  for (int s = 0; s < nswap; ++s) {
+    int orig_sendnum = lmp_list.sendnum[s];
+    new_sendlist[s].clear();
+    new_sendlist[s].reserve(orig_sendnum);
+
+    for (int k = 0; k < orig_sendnum; ++k) {
+      int orig_idx = lmp_list.sendlist[s][k];
+      int real_idx = fwd_map[orig_idx];
+      if (real_idx >= 0) {
+        new_sendlist[s].push_back(real_idx);
+      }
+    }
+    new_sendnum[s] = static_cast<int>(new_sendlist[s].size());
+
+    // Independently compute recvnum using firstrecv range.
+    // In MPI, sendnum and recvnum are independent per process.
+    int firstrecv = lmp_list.firstrecv[s];
+    int orig_recvnum = lmp_list.recvnum[s];
+    int recv_count = 0;
+    for (int k = 0; k < orig_recvnum; ++k) {
+      int orig_idx = firstrecv + k;
+      if (fwd_map[orig_idx] >= 0) {
+        ++recv_count;
+      }
+    }
+    new_recvnum[s] = recv_count;
+  }
+}
+
 void deepmd::NeighborListData::copy_from_nlist(const InputNlist& inlist,
                                                const int natoms) {
   int inum = natoms >= 0 ? natoms : inlist.inum;

source/lmp/fix_dplr.cpp

@@ -185,7 +185,7 @@ FixDPLR::FixDPLR(LAMMPS* lmp, int narg, char** arg)
         }
       }
       if (!found_element && "NULL" == type_name) {
-        type_idx_map.push_back(type_map.size());  // ghost type
+        type_idx_map.push_back(-1);  // virtual atom
         found_element = true;
       }
       if (!found_element) {

source/lmp/pair_deepmd.cpp

@@ -853,7 +853,7 @@ void PairDeepMD::coeff(int narg, char** arg) {
         }
       }
       if (!found_element && "NULL" == type_name) {
-        type_idx_map.push_back(type_map.size());  // ghost type
+        type_idx_map.push_back(-1);  // virtual atom
         found_element = true;
       }
       if (!found_element) {

source/lmp/pair_deepspin.cpp

@@ -883,7 +883,7 @@ void PairDeepSpin::coeff(int narg, char** arg) {
         }
       }
       if (!found_element && "NULL" == type_name) {
-        type_idx_map.push_back(type_map.size());  // ghost type
+        type_idx_map.push_back(-1);  // virtual atom
         found_element = true;
       }
       if (!found_element) {