fix(pd): shape send list by swap count

source/api_cc/include/common.h

@@ -177,6 +177,19 @@ void remap_comm_sendlist(std::vector<std::vector<int>>& new_sendlist,
                          const InputNlist& lmp_list,
                          const std::vector<int>& fwd_map);
 
+/**
+ * @brief Pack one communication send-list pointer address per MPI swap.
+ *
+ * The corresponding send counts describe the number of atom indices behind
+ * each pointer; they do not change the length of the pointer-address array.
+ *
+ * @param[in] sendlist Array of send-list pointers indexed by swap.
+ * @param[in] nswap Number of communication swaps.
+ * @return Integer pointer addresses with exactly nswap entries.
+ */
+std::vector<std::intptr_t> pack_comm_sendlist_pointers(int* const* sendlist,
+                                                       int nswap);
+
 /**
  * @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/src/DeepPotPD.cc

@@ -3,7 +3,6 @@
 #include "DeepPotPD.h"
 
 #include <cstdint>
-#include <numeric>
 
 #include "common.h"
 #include "device.h"
@@ -441,21 +440,15 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
       }
 
       assert(sizeof(std::intptr_t) == 8);
-      int total_send = std::accumulate(eff_sendnum, eff_sendnum + nswap, 0);
-      sendlist_tensor->Reshape({total_send});
 
       /**
       ** NOTE: paddle do not support construct a Tensor with from_blob(T**, ...)
-      ** from a double pointer, so we convert int* pointer to indptr_t for each
-      ** entry and wrap it into int64 Tensor as a workaround.
+      ** from a double pointer, so we convert each int* pointer to an intptr_t
+      ** entry and wrap the addresses in an int64 Tensor as a workaround.
       */
-      std::vector<std::intptr_t> pointer_addresses;
-      pointer_addresses.reserve(nswap);
-      for (int iswap = 0; iswap < nswap; ++iswap) {
-        std::intptr_t addr =
-            reinterpret_cast<std::intptr_t>(eff_sendlist[iswap]);
-        pointer_addresses.push_back(addr);
-      }
+      std::vector<std::intptr_t> pointer_addresses =
+          pack_comm_sendlist_pointers(eff_sendlist, nswap);
+      sendlist_tensor->Reshape({static_cast<int>(pointer_addresses.size())});
       sendlist_tensor->CopyFromCpu(pointer_addresses.data());
     }
     if (lmp_list.mapping) {

source/api_cc/src/common.cc

@@ -277,6 +277,17 @@ void deepmd::remap_comm_sendlist(std::vector<std::vector<int>>& new_sendlist,
   }
 }
 
+std::vector<std::intptr_t> deepmd::pack_comm_sendlist_pointers(
+    int* const* sendlist, int nswap) {
+  std::vector<std::intptr_t> pointer_addresses;
+  pointer_addresses.reserve(nswap);
+  for (int iswap = 0; iswap < nswap; ++iswap) {
+    pointer_addresses.push_back(
+        reinterpret_cast<std::intptr_t>(sendlist[iswap]));
+  }
+  return pointer_addresses;
+}
+
 void deepmd::NeighborListData::copy_from_nlist(const InputNlist& inlist,
                                                const int natoms) {
   int inum = natoms >= 0 ? natoms : inlist.inum;

source/api_cc/tests/test_select_real_atoms_sendlist.cc

@@ -149,3 +149,20 @@ TEST_F(TestRemapCommSendlist, all_virtual_in_swap) {
   EXPECT_EQ(new_recvnum[0], 0);
   EXPECT_EQ(new_recvnum[1], 0);
 }
+
+// Paddle passes sendlist as an int64 tensor of pointer addresses. Its length
+// is the swap count, while sendnum controls how many atom indices each pointer
+// references. This fixture has sum(sendnum)=7 and nswap=2 to catch accidental
+// shaping by the total atom count.
+TEST_F(TestRemapCommSendlist, pointer_array_is_shaped_by_swap_count) {
+  const std::vector<std::intptr_t> pointer_addresses =
+      deepmd::pack_comm_sendlist_pointers(lmp_list.sendlist, lmp_list.nswap);
+
+  ASSERT_EQ(pointer_addresses.size(), 2U);
+  EXPECT_EQ(pointer_addresses[0],
+            reinterpret_cast<std::intptr_t>(sendlist_ptrs[0]));
+  EXPECT_EQ(pointer_addresses[1],
+            reinterpret_cast<std::intptr_t>(sendlist_ptrs[1]));
+  EXPECT_NE(pointer_addresses.size(),
+            static_cast<size_t>(sendnum_arr[0] + sendnum_arr[1]));
+}