fix(c++): fix NULL type in custom op (#4889)

Replaces usage of lmp_list send/recv arrays with new vectors that map
indices using fwd_map and synchronize counts via MPI. Updates tensor
construction to use these new vectors, improving correctness and
flexibility in distributed communication.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Refactor**
* Message-passing now remaps send lists to actual atoms before building
communication tensors, improving consistency for distributed runs.
* **Bug Fixes**
* Invalid or out-of-range send indices are filtered and counts updated
to prevent communication mismatches and related errors.
* **Tests**
* Added a test covering a type-map scenario with a NULL mapping to
exercise the updated handling.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>

Author: 3 people

source/api_cc/include/DeepPotPD.h

@@ -392,6 +392,10 @@ class DeepPotPD : public DeepPotBackend {
   bool gpu_enabled;
   std::unique_ptr<paddle_infer::Tensor> firstneigh_tensor;
   std::unique_ptr<paddle_infer::Tensor> mapping_tensor;
+  std::vector<std::vector<int>> remapped_sendlist;
+  std::vector<int*> remapped_sendlist_ptrs;
+  std::vector<int> remapped_sendnum;
+  std::vector<int> remapped_recvnum;
 };
 
 }  // namespace deepmd

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/DeepPotPD.cc

@@ -378,6 +378,8 @@ void DeepPotPD::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;
   auto coord_wrapped_Tensor = predictor_fl->GetInputHandle("coord");
@@ -391,6 +393,28 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
     nlist_data.shuffle_exclude_empty(fwd_map);
     nlist_data.padding();
     if (do_message_passing) {
+      // Determine the actual sendlist/sendnum/recvnum to use.
+      // When NULL-type atoms exist, remap sendlist indices through fwd_map.
+      int** eff_sendlist;
+      int* eff_sendnum;
+      int* eff_recvnum;
+      if (has_null_atoms) {
+        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();
+        }
+        eff_sendlist = remapped_sendlist_ptrs.data();
+        eff_sendnum = remapped_sendnum.data();
+        eff_recvnum = remapped_recvnum.data();
+      } else {
+        eff_sendlist = lmp_list.sendlist;
+        eff_sendnum = lmp_list.sendnum;
+        eff_recvnum = lmp_list.recvnum;
+      }
+
       auto sendproc_tensor = predictor_fl->GetInputHandle("send_proc");
       auto recvproc_tensor = predictor_fl->GetInputHandle("recv_proc");
       auto recvnum_tensor = predictor_fl->GetInputHandle("recv_num");
@@ -406,26 +430,18 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
       recvproc_tensor->CopyFromCpu(lmp_list.recvproc);
 
       recvnum_tensor->Reshape({nswap});
-      recvnum_tensor->CopyFromCpu(lmp_list.recvnum);
+      recvnum_tensor->CopyFromCpu(eff_recvnum);
 
       sendnum_tensor->Reshape({nswap});
-      if (sizeof(lmp_list.sendnum[0]) != sizeof(int32_t)) {
-        std::vector<int32_t> temp_data(nswap);
-        for (int i = 0; i < nswap; i++) {
-          temp_data[i] = static_cast<int32_t>(lmp_list.sendnum[i]);
-        }
-        sendnum_tensor->CopyFromCpu(temp_data.data());
-      } else {
-        sendnum_tensor->CopyFromCpu(lmp_list.sendnum);
-      }
+      sendnum_tensor->CopyFromCpu(eff_sendnum);
+
       communicator_tensor->Reshape({1});
       if (lmp_list.world) {
         communicator_tensor->CopyFromCpu(static_cast<int*>(lmp_list.world));
       }
 
       assert(sizeof(std::intptr_t) == 8);
-      int total_send =
-          std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0);
+      int total_send = std::accumulate(eff_sendnum, eff_sendnum + nswap, 0);
       sendlist_tensor->Reshape({total_send});
 
       /**
@@ -437,7 +453,7 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
       pointer_addresses.reserve(nswap);
       for (int iswap = 0; iswap < nswap; ++iswap) {
         std::intptr_t addr =
-            reinterpret_cast<std::intptr_t>(lmp_list.sendlist[iswap]);
+            reinterpret_cast<std::intptr_t>(eff_sendlist[iswap]);
         pointer_addresses.push_back(addr);
       }
       sendlist_tensor->CopyFromCpu(pointer_addresses.data());

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,