Merge branch 'devel' into pd_grad_acc

Author: HydrogenSulfate

source/api_cc/include/DeepPotPD.h

@@ -282,12 +282,11 @@ class DeepPotPD : public DeepPotBackend {
    * @brief Compute the number of elements in a tensor.
    * @param[in] x Tensor x.
    **/
-  int numel(const paddle_infer::Tensor& x) const {
-    // TODO: There might be a overflow problem here for multiply int numbers.
-    int ret = 1;
+  size_t numel(const paddle_infer::Tensor& x) const {
+    size_t ret = 1;
     std::vector<int> x_shape = x.shape();
     for (std::size_t i = 0, n = x_shape.size(); i < n; ++i) {
-      ret *= x_shape[i];
+      ret *= static_cast<size_t>(x_shape[i]);
     }
     return ret;
   };

source/api_cc/src/DeepPotPD.cc

@@ -474,9 +474,9 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
   auto energy_ = predictor_fl->GetOutputHandle(output_names.at(1));
   auto force_ = predictor_fl->GetOutputHandle(output_names.at(2));
   auto virial_ = predictor_fl->GetOutputHandle(output_names.at(4));
-  int output_energy_size = numel(*energy_);
-  int output_force_size = numel(*force_);
-  int output_virial_size = numel(*virial_);
+  size_t output_energy_size = numel(*energy_);
+  size_t output_force_size = numel(*force_);
+  size_t output_virial_size = numel(*virial_);
   // output energy
   ener.resize(output_energy_size);
   energy_->CopyToCpu(ener.data());
@@ -597,26 +597,26 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
   auto force_ = predictor->GetOutputHandle(output_names.at(3));
   auto virial_ = predictor->GetOutputHandle(output_names.at(5));
 
-  int enery_numel = numel(*energy_);
+  size_t enery_numel = numel(*energy_);
   assert(enery_numel > 0);
   ener.resize(enery_numel);
   energy_->CopyToCpu(ener.data());
 
-  int force_numel = numel(*force_);
+  size_t force_numel = numel(*force_);
   assert(force_numel > 0);
   force.resize(force_numel);
   force_->CopyToCpu(force.data());
 
-  int virial_numel = numel(*virial_);
+  size_t virial_numel = numel(*virial_);
   assert(virial_numel > 0);
   virial.resize(virial_numel);
   virial_->CopyToCpu(virial.data());
 
   if (atomic) {
     auto atom_energy_ = predictor->GetOutputHandle(output_names.at(0));
     auto atom_virial_ = predictor->GetOutputHandle(output_names.at(1));
-    int atom_energy_numel = numel(*atom_energy_);
-    int atom_virial_numel = numel(*atom_virial_);
+    size_t atom_energy_numel = numel(*atom_energy_);
+    size_t atom_virial_numel = numel(*atom_virial_);
     assert(atom_energy_numel > 0);
     assert(atom_virial_numel > 0);
     atom_energy.resize(atom_energy_numel);
@@ -656,7 +656,7 @@ template void DeepPotPD::compute<float, std::vector<ENERGYTYPE>>(
 that need to be postprocessed */
 void DeepPotPD::get_type_map(std::string& type_map) {
   auto type_map_tensor = predictor->GetOutputHandle("buffer_type_map");
-  int type_map_size = numel(*type_map_tensor);
+  size_t type_map_size = numel(*type_map_tensor);
 
   std::vector<int> type_map_arr(type_map_size, 0);
   type_map_tensor->CopyToCpu(type_map_arr.data());
@@ -670,7 +670,7 @@ template <typename BUFFERTYPE>
 void DeepPotPD::get_buffer(const std::string& buffer_name,
                            std::vector<BUFFERTYPE>& buffer_array) {
   auto buffer_tensor = predictor->GetOutputHandle(buffer_name);
-  int buffer_size = numel(*buffer_tensor);
+  size_t buffer_size = numel(*buffer_tensor);
   buffer_array.resize(buffer_size);
   buffer_tensor->CopyToCpu(buffer_array.data());
 }