Revert commit

This reverts commit b415c67.

Author: iProzd

source/api_cc/include/DeepPotPD.h

@@ -392,10 +392,6 @@ class DeepPotPD : public DeepPotBackend {
   bool gpu_enabled;
   std::unique_ptr<paddle_infer::Tensor> firstneigh_tensor;
   std::unique_ptr<paddle_infer::Tensor> mapping_tensor;
-  // Owned comm data for remapped sendlist (NULL type filtering)
-  std::vector<std::vector<int>> comm_sendlist_;
-  std::vector<int> comm_sendnum_;
-  std::vector<int> comm_recvnum_;
 };
 
 }  // namespace deepmd

source/api_cc/include/DeepPotPT.h

@@ -340,11 +340,6 @@ class DeepPotPT : public DeepPotBackend {
   at::Tensor firstneigh_tensor;
   c10::optional<torch::Tensor> mapping_tensor;
   torch::Dict<std::string, torch::Tensor> comm_dict;
-  // Owned comm data for remapped sendlist (NULL type filtering)
-  std::vector<std::vector<int>> comm_sendlist_;
-  std::vector<int> comm_sendnum_;
-  std::vector<int> comm_recvnum_;
-  std::vector<int*> comm_sendlist_ptrs_;
   bool profiler_enabled{false};
   std::string profiler_file;
   /**

source/api_cc/include/DeepSpinPT.h

@@ -262,11 +262,6 @@ class DeepSpinPT : public DeepSpinBackend {
   at::Tensor firstneigh_tensor;
   c10::optional<torch::Tensor> mapping_tensor;
   torch::Dict<std::string, torch::Tensor> comm_dict;
-  // Owned comm data for remapped sendlist (NULL type filtering)
-  std::vector<std::vector<int>> comm_sendlist_;
-  std::vector<int> comm_sendnum_;
-  std::vector<int> comm_recvnum_;
-  std::vector<int*> comm_sendlist_ptrs_;
   /**
    * @brief Translate PyTorch exceptions to the DeePMD-kit exception.
    * @param[in] f The function to run.

source/api_cc/include/common.h

@@ -153,20 +153,6 @@ void select_map_inv(typename std::vector<VT>::iterator out,
                     const std::vector<int>& fwd_map,
                     const int& stride);
 
-/**
- * @brief Build filtered/remapped sendlist for NULL type filtering.
- * @param[out] new_sendlist Remapped sendlist vectors (one per swap).
- * @param[out] new_sendnum Number of real atoms to send per swap.
- * @param[out] new_recvnum Number of real atoms to receive per swap.
- * @param[in] inlist The input neighbor list with comm data.
- * @param[in] fwd_map Forward map from original to remapped indices.
- */
-void select_real_atoms_sendlist(std::vector<std::vector<int>>& new_sendlist,
-                                std::vector<int>& new_sendnum,
-                                std::vector<int>& new_recvnum,
-                                const InputNlist& inlist,
-                                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/src/DeepPotPD.cc

@@ -399,29 +399,33 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
       auto sendlist_tensor = predictor_fl->GetInputHandle("send_list");
 
       int nswap = lmp_list.nswap;
-      select_real_atoms_sendlist(comm_sendlist_, comm_sendnum_, comm_recvnum_,
-                                 lmp_list, fwd_map);
       sendproc_tensor->Reshape({nswap});
       sendproc_tensor->CopyFromCpu(lmp_list.sendproc);
 
       recvproc_tensor->Reshape({nswap});
       recvproc_tensor->CopyFromCpu(lmp_list.recvproc);
 
       recvnum_tensor->Reshape({nswap});
-      recvnum_tensor->CopyFromCpu(comm_recvnum_.data());
+      recvnum_tensor->CopyFromCpu(lmp_list.recvnum);
 
       sendnum_tensor->Reshape({nswap});
-      sendnum_tensor->CopyFromCpu(comm_sendnum_.data());
+      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);
+      }
       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 = 0;
-      for (int s = 0; s < nswap; ++s) {
-        total_send += comm_sendnum_[s];
-      }
+      int total_send =
+          std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0);
       sendlist_tensor->Reshape({total_send});
 
       /**
@@ -433,7 +437,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>(comm_sendlist_[iswap].data());
+            reinterpret_cast<std::intptr_t>(lmp_list.sendlist[iswap]);
         pointer_addresses.push_back(addr);
       }
       sendlist_tensor->CopyFromCpu(pointer_addresses.data());

source/api_cc/src/DeepPotPT.cc

@@ -212,20 +212,16 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
     nlist_data.padding();
     if (do_message_passing) {
       int nswap = lmp_list.nswap;
-      select_real_atoms_sendlist(comm_sendlist_, comm_sendnum_, comm_recvnum_,
-                                 lmp_list, fwd_map);
-      comm_sendlist_ptrs_.resize(nswap);
-      for (int s = 0; s < nswap; ++s) {
-        comm_sendlist_ptrs_[s] = comm_sendlist_[s].data();
-      }
       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 sendnum_tensor =
-          torch::from_blob(comm_sendnum_.data(), {nswap}, int32_option);
+      torch::Tensor firstrecv_tensor =
+          torch::from_blob(lmp_list.firstrecv, {nswap}, int32_option);
       torch::Tensor recvnum_tensor =
-          torch::from_blob(comm_recvnum_.data(), {nswap}, int32_option);
+          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);
@@ -234,9 +230,11 @@ void DeepPotPT::compute(ENERGYVTYPE& ener,
             const_cast<void*>(lmp_list.world), {1}, torch::kInt64);
       }
 
-      int ptr_len = nswap * static_cast<int>(sizeof(int*) / sizeof(int32_t));
-      torch::Tensor sendlist_tensor = torch::from_blob(
-          comm_sendlist_ptrs_.data(), {std::max(ptr_len, 1)}, int32_option);
+      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);

source/api_cc/src/DeepSpinPT.cc

@@ -182,30 +182,28 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener,
     nlist_data.padding();
     if (do_message_passing) {
       int nswap = lmp_list.nswap;
-      select_real_atoms_sendlist(comm_sendlist_, comm_sendnum_, comm_recvnum_,
-                                 lmp_list, fwd_map);
-      comm_sendlist_ptrs_.resize(nswap);
-      for (int s = 0; s < nswap; ++s) {
-        comm_sendlist_ptrs_[s] = comm_sendlist_[s].data();
-      }
       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 sendnum_tensor =
-          torch::from_blob(comm_sendnum_.data(), {nswap}, int32_option);
+      torch::Tensor firstrecv_tensor =
+          torch::from_blob(lmp_list.firstrecv, {nswap}, int32_option);
       torch::Tensor recvnum_tensor =
-          torch::from_blob(comm_recvnum_.data(), {nswap}, int32_option);
+          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);
       } else {
         communicator_tensor = torch::from_blob(
             const_cast<void*>(lmp_list.world), {1}, torch::kInt64);
       }
-      int ptr_len = nswap * static_cast<int>(sizeof(int*) / sizeof(int32_t));
-      torch::Tensor sendlist_tensor = torch::from_blob(
-          comm_sendlist_ptrs_.data(), {std::max(ptr_len, 1)}, int32_option);
+      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);

source/api_cc/src/common.cc

@@ -232,32 +232,6 @@ template void deepmd::select_real_atoms_coord<float>(
     const int& nall,
     const bool aparam_nall);
 
-void deepmd::select_real_atoms_sendlist(
-    std::vector<std::vector<int>>& new_sendlist,
-    std::vector<int>& new_sendnum,
-    std::vector<int>& new_recvnum,
-    const InputNlist& inlist,
-    const std::vector<int>& fwd_map) {
-  int nswap = inlist.nswap;
-  new_sendlist.resize(nswap);
-  new_sendnum.resize(nswap);
-  new_recvnum.resize(nswap);
-  for (int s = 0; s < nswap; ++s) {
-    new_sendlist[s].clear();
-    int orig_num = inlist.sendnum[s];
-    new_sendlist[s].reserve(orig_num);
-    for (int i = 0; i < orig_num; ++i) {
-      int idx = inlist.sendlist[s][i];
-      int mapped = fwd_map[idx];
-      if (mapped >= 0) {
-        new_sendlist[s].push_back(mapped);
-      }
-    }
-    new_sendnum[s] = new_sendlist[s].size();
-    new_recvnum[s] = new_sendlist[s].size();
-  }
-}
-
 void deepmd::NeighborListData::copy_from_nlist(const InputNlist& inlist,
                                                const int natoms) {
   int inum = natoms >= 0 ? natoms : inlist.inum;