restore

Author: HydrogenSulfate

source/api_cc/src/DeepPotPD.cc

@@ -164,23 +164,15 @@ inline void enableTimestamp(bool enable = true) {
 }
 }  // namespace logg
 
-void fillNlistTensor(const std::vector<std::vector<int>>& data,
-                     std::unique_ptr<paddle_infer::Tensor>& flat_tensor) {
-  size_t total_size = 0;
+std::vector<int> createNlistTensorPD(
+    const std::vector<std::vector<int>>& data) {
+  std::vector<int> ret;
   for (const auto& row : data) {
-    total_size += row.size();
+    ret.insert(ret.end(), row.begin(), row.end());
   }
-  std::vector<int> flat_data;
-  flat_data.reserve(total_size);
-  for (const auto& row : data) {
-    flat_data.insert(flat_data.end(), row.begin(), row.end());
-  }
-
-  int nloc = data.size();
-  int nnei = nloc > 0 ? total_size / nloc : 0;
-  flat_tensor->Reshape({1, nloc, nnei});
-  flat_tensor->CopyFromCpu(flat_data.data());
+  return ret;
 }
+
 DeepPotPD::DeepPotPD() : inited(false) {}
 DeepPotPD::DeepPotPD(const std::string& model,
                      const int& gpu_rank,
@@ -383,14 +375,16 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
   auto coord_wrapped_Tensor = predictor_fl->GetInputHandle("coord");
   coord_wrapped_Tensor->Reshape({1, nall_real, 3});
   coord_wrapped_Tensor->CopyFromCpu(coord_wrapped.data());
+
   auto atype_Tensor = predictor_fl->GetInputHandle("atype");
   atype_Tensor->Reshape({1, nall_real});
   atype_Tensor->CopyFromCpu(datype.data());
+
   if (ago == 0) {
-    nlist_data.copy_from_nlist(lmp_list, nall - nghost);
+    nlist_data.copy_from_nlist(lmp_list);
     nlist_data.shuffle_exclude_empty(fwd_map);
     nlist_data.padding();
-    if (do_message_passing) {
+    if (do_message_passing == 1 && nghost > 0) {
       auto sendproc_tensor = predictor_fl->GetInputHandle("send_proc");
       auto recvproc_tensor = predictor_fl->GetInputHandle("recv_proc");
       auto recvnum_tensor = predictor_fl->GetInputHandle("recv_num");
@@ -452,17 +446,20 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
       this->mapping_tensor->CopyFromCpu(mapping.data());
     }
   }
+  std::vector<int> firstneigh = createNlistTensorPD(nlist_data.jlist);
   this->firstneigh_tensor = predictor_fl->GetInputHandle("nlist");
-  fillNlistTensor(nlist_data.jlist, this->firstneigh_tensor);
+  this->firstneigh_tensor->Reshape(
+      {1, nloc, (int)firstneigh.size() / (int)nloc});
+  this->firstneigh_tensor->CopyFromCpu(firstneigh.data());
   bool do_atom_virial_tensor = atomic;
-  std::unique_ptr<paddle_infer::Tensor> fparam_tensor;
   if (!fparam.empty()) {
+    std::unique_ptr<paddle_infer::Tensor> fparam_tensor;
     fparam_tensor = predictor_fl->GetInputHandle("fparam");
     fparam_tensor->Reshape({1, static_cast<int>(fparam.size())});
     fparam_tensor->CopyFromCpu(fparam.data());
   }
-  std::unique_ptr<paddle_infer::Tensor> aparam_tensor;
   if (!aparam_.empty()) {
+    std::unique_ptr<paddle_infer::Tensor> aparam_tensor;
     aparam_tensor = predictor_fl->GetInputHandle("aparam");
     aparam_tensor->Reshape(
         {1, lmp_list.inum, static_cast<int>(aparam_.size()) / lmp_list.inum});
@@ -513,7 +510,7 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
   }
 }
 template void DeepPotPD::compute<double, std::vector<ENERGYTYPE>>(
-    std::vector<ENERGYTYPE>& ener,
+    std::vector<ENERGYTYPE>& dener,
     std::vector<double>& force,
     std::vector<double>& virial,
     std::vector<double>& atom_energy,
@@ -525,10 +522,11 @@ template void DeepPotPD::compute<double, std::vector<ENERGYTYPE>>(
     const InputNlist& lmp_list,
     const int& ago,
     const std::vector<double>& fparam,
-    const std::vector<double>& aparam,
+    const std::vector<double>& aparam_,
     const bool atomic);
+
 template void DeepPotPD::compute<float, std::vector<ENERGYTYPE>>(
-    std::vector<ENERGYTYPE>& ener,
+    std::vector<ENERGYTYPE>& dener,
     std::vector<float>& force,
     std::vector<float>& virial,
     std::vector<float>& atom_energy,
@@ -540,8 +538,9 @@ template void DeepPotPD::compute<float, std::vector<ENERGYTYPE>>(
     const InputNlist& lmp_list,
     const int& ago,
     const std::vector<float>& fparam,
-    const std::vector<float>& aparam,
+    const std::vector<float>& aparam_,
     const bool atomic);
+
 // ENERGYVTYPE: std::vector<ENERGYTYPE> or ENERGYTYPE
 template <typename VALUETYPE, typename ENERGYVTYPE>
 void DeepPotPD::compute(ENERGYVTYPE& ener,
@@ -563,9 +562,10 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
   coord_wrapped_Tensor->Reshape({1, natoms, 3});
   coord_wrapped_Tensor->CopyFromCpu(coord_wrapped.data());
 
+  std::vector<std::int64_t> atype_64(atype.begin(), atype.end());
   auto atype_Tensor = predictor->GetInputHandle("atype");
   atype_Tensor->Reshape({1, natoms});
-  atype_Tensor->CopyFromCpu(atype.data());
+  atype_Tensor->CopyFromCpu(atype_64.data());
 
   std::unique_ptr<paddle_infer::Tensor> box_Tensor;
   if (!box.empty()) {
@@ -575,13 +575,13 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
   }
   std::unique_ptr<paddle_infer::Tensor> fparam_tensor;
   if (!fparam.empty()) {
-    fparam_tensor = predictor->GetInputHandle("fparam");
+    fparam_tensor = predictor->GetInputHandle("box");
     fparam_tensor->Reshape({1, static_cast<int>(fparam.size())});
     fparam_tensor->CopyFromCpu((fparam.data()));
   }
   std::unique_ptr<paddle_infer::Tensor> aparam_tensor;
   if (!aparam.empty()) {
-    aparam_tensor = predictor->GetInputHandle("aparam");
+    aparam_tensor = predictor->GetInputHandle("box");
     aparam_tensor->Reshape(
         {1, natoms, static_cast<int>(aparam.size()) / natoms});
     aparam_tensor->CopyFromCpu((aparam.data()));
@@ -628,11 +628,11 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
 
 template void DeepPotPD::compute<double, std::vector<ENERGYTYPE>>(
     std::vector<ENERGYTYPE>& ener,
-    std::vector<double>& force,
+    std::vector<double>& dforce,
     std::vector<double>& virial,
     std::vector<double>& atom_energy,
     std::vector<double>& atom_virial,
-    const std::vector<double>& coord,
+    const std::vector<double>& dcoord,
     const std::vector<int>& atype,
     const std::vector<double>& box,
     const std::vector<double>& fparam,
@@ -645,7 +645,7 @@ template void DeepPotPD::compute<float, std::vector<ENERGYTYPE>>(
     std::vector<float>& virial,
     std::vector<float>& atom_energy,
     std::vector<float>& atom_virial,
-    const std::vector<float>& coord,
+    const std::vector<float>& dcoord,
     const std::vector<int>& atype,
     const std::vector<float>& box,
     const std::vector<float>& fparam,