From 5af9aa510c68452dce46fe0d5154b3ae779c5bf3 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 06:25:57 +0800 Subject: [PATCH] fix(pd): shape send list by swap count Pack exactly one send-list pointer address per communication swap and shape the Paddle tensor from that packed array instead of the total atom send count. Add a regression with two swaps and seven sent atoms to keep pointer-array length independent from per-swap send counts. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/api_cc/include/common.h | 13 +++++++++++++ source/api_cc/src/DeepPotPD.cc | 17 +++++------------ source/api_cc/src/common.cc | 11 +++++++++++ .../tests/test_select_real_atoms_sendlist.cc | 17 +++++++++++++++++ 4 files changed, 46 insertions(+), 12 deletions(-) diff --git a/source/api_cc/include/common.h b/source/api_cc/include/common.h index 283205a028..ae09df007f 100644 --- a/source/api_cc/include/common.h +++ b/source/api_cc/include/common.h @@ -177,6 +177,19 @@ void remap_comm_sendlist(std::vector>& new_sendlist, const InputNlist& lmp_list, const std::vector& 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 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. diff --git a/source/api_cc/src/DeepPotPD.cc b/source/api_cc/src/DeepPotPD.cc index 25128c017c..9e8f05c332 100644 --- a/source/api_cc/src/DeepPotPD.cc +++ b/source/api_cc/src/DeepPotPD.cc @@ -3,7 +3,6 @@ #include "DeepPotPD.h" #include -#include #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 pointer_addresses; - pointer_addresses.reserve(nswap); - for (int iswap = 0; iswap < nswap; ++iswap) { - std::intptr_t addr = - reinterpret_cast(eff_sendlist[iswap]); - pointer_addresses.push_back(addr); - } + std::vector pointer_addresses = + pack_comm_sendlist_pointers(eff_sendlist, nswap); + sendlist_tensor->Reshape({static_cast(pointer_addresses.size())}); sendlist_tensor->CopyFromCpu(pointer_addresses.data()); } if (lmp_list.mapping) { diff --git a/source/api_cc/src/common.cc b/source/api_cc/src/common.cc index dc604b32e9..d4b5988c54 100644 --- a/source/api_cc/src/common.cc +++ b/source/api_cc/src/common.cc @@ -277,6 +277,17 @@ void deepmd::remap_comm_sendlist(std::vector>& new_sendlist, } } +std::vector deepmd::pack_comm_sendlist_pointers( + int* const* sendlist, int nswap) { + std::vector pointer_addresses; + pointer_addresses.reserve(nswap); + for (int iswap = 0; iswap < nswap; ++iswap) { + pointer_addresses.push_back( + reinterpret_cast(sendlist[iswap])); + } + return pointer_addresses; +} + void deepmd::NeighborListData::copy_from_nlist(const InputNlist& inlist, const int natoms) { int inum = natoms >= 0 ? natoms : inlist.inum; diff --git a/source/api_cc/tests/test_select_real_atoms_sendlist.cc b/source/api_cc/tests/test_select_real_atoms_sendlist.cc index f87f5eda6d..5aabccfb53 100644 --- a/source/api_cc/tests/test_select_real_atoms_sendlist.cc +++ b/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 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(sendlist_ptrs[0])); + EXPECT_EQ(pointer_addresses[1], + reinterpret_cast(sendlist_ptrs[1])); + EXPECT_NE(pointer_addresses.size(), + static_cast(sendnum_arr[0] + sendnum_arr[1])); +}