Fix memory recorder (deepmodeling#6952)

Author: Critsium-xy

source/source_base/memory.cpp

@@ -159,7 +159,7 @@ double Memory::record
 void Memory::record
 (
 	const std::string &name_in,
-	const size_t &n_in,
+	const long &n_in,
 	const bool accumulate
 )
 {

source/source_base/memory.h

@@ -45,7 +45,7 @@ class Memory
      */
     static void record(
       const std::string &name_in,
-      const size_t &n_in,
+      const long &n_in,
       const bool accumulate = false
     );
 

source/source_lcao/module_gint/biggrid_info.cpp

@@ -1,6 +1,7 @@
 #include "biggrid_info.h"
 #include "gint_helper.h"
 #include "gint_type.h"
+#include "source_base/memory.h"
 
 namespace ModuleGint
 {
@@ -44,6 +45,12 @@ BigGridInfo::BigGridInfo(
             meshgrid_coords_[index_1d] = 
                 meshgrid_info_->get_cartesian_coord(mgrid_idx_1Dto3D(index_1d));
         }
+        ModuleBase::Memory::record("BigGridInfo::meshgrid_coords", (long long)nmxyz_ * sizeof(Vec3d), true);
+    }
+
+    BigGridInfo::~BigGridInfo()
+    {
+        ModuleBase::Memory::record("BigGridInfo::meshgrid_coords", -(long long)nmxyz_ * sizeof(Vec3d), true);
     }
 
     Vec3i BigGridInfo::max_ext_bgrid_num(double r) const

source/source_lcao/module_gint/biggrid_info.h

@@ -21,6 +21,8 @@ class BigGridInfo
             Vec3d biggrid_vec2,
             Vec3d biggrid_vec3,
             int nmx, int nmy, int nmz);
+
+        ~BigGridInfo();
 
         Vec3d get_cartesian_coord(const Vec3d& index_3d) const { return index_3d * biggrid_latvec0_; }
         Vec3d get_cartesian_coord(const Vec3i& index_3d) const { return index_3d * biggrid_latvec0_; }

source/source_lcao/module_gint/gint_info.cpp

@@ -4,6 +4,7 @@
 #include "source_base/timer.h"
 #include "gint_info.h"
 #include "gint_type.h"
+#include "source_base/memory.h"
 
 namespace ModuleGint
 {
@@ -62,6 +63,12 @@ GintInfo::GintInfo(
     #endif
 }
 
+GintInfo::~GintInfo()
+{
+    ModuleBase::Memory::record("GintInfo::trace_lo_", -(long long)(sizeof(int) * trace_lo_.size()), true);
+    ModuleBase::Memory::record("GintInfo::ijr_info_", -(long long)(sizeof(int) * ijr_info_.size()), true);
+}
+
 template <typename T>
 HContainer<T> GintInfo::get_hr(int npol) const
 {
@@ -194,6 +201,7 @@ void GintInfo::init_trace_lo_(const UnitCell& ucell, const int nspin)
             ++iat;
         }
     }
+    ModuleBase::Memory::record("GintInfo::trace_lo_", (long long)(sizeof(int) * trace_lo_.size()), true);
 }
 
 void GintInfo::init_ijr_info_(const UnitCell& ucell, Grid_Driver& gd)
@@ -260,6 +268,7 @@ void GintInfo::init_ijr_info_(const UnitCell& ucell, Grid_Driver& gd)
             }
     }
     this->ijr_info_ = hr_gint_local.get_ijr_info();
+    ModuleBase::Memory::record("GintInfo::ijr_info_", (long long)(sizeof(int) * ijr_info_.size()), true);
     return;
 }
 

source/source_lcao/module_gint/gint_info.h

@@ -33,6 +33,8 @@ class GintInfo
         const Numerical_Orbital* Phi,
         const UnitCell& ucell, Grid_Driver& gd);
 
+    ~GintInfo();
+
     // getter functions
     const std::vector<std::shared_ptr<BigGrid>>& get_biggrids() { return biggrids_; }
     int get_bgrids_num() const { return static_cast<int>(biggrids_.size()); }

source/source_lcao/module_hcontainer/hcontainer.cpp

@@ -1,4 +1,5 @@
 #include "hcontainer.h"
+#include "source_base/memory.h"
 
 namespace hamilt
 {
@@ -12,12 +13,16 @@ HContainer<T>::~HContainer()
 {
     if(this->allocated)
     {
+        if(this->allocated_size > 0)
+        {
+            ModuleBase::Memory::record("HContainer", -(long long)this->allocated_size, true);
+        }
         delete[] this->wrapper_pointer;
     }
 }
 
 template <typename T>
-HContainer<T>::HContainer() {}
+HContainer<T>::HContainer() : allocated_size(0) {}
 
 // copy constructor
 template <typename T>
@@ -30,6 +35,7 @@ HContainer<T>::HContainer(const HContainer<T>& HR_in, T* data_array)
     this->current_R = -1;
     this->wrapper_pointer = data_array;
     this->allocated = false;
+    this->allocated_size = 0;
     this->atom_pairs = HR_in.atom_pairs;
     // data of HR_in will not be copied, please call add() after this constructor to copy data.
     this->allocate(this->wrapper_pointer, true);
@@ -47,8 +53,11 @@ HContainer<T>::HContainer(HContainer<T>&& HR_in) noexcept
     this->gamma_only = HR_in.gamma_only;
     this->paraV = HR_in.paraV;
     this->allocated = HR_in.allocated;
+    this->allocated_size = HR_in.allocated_size;
     this->current_R = -1;
     HR_in.wrapper_pointer = nullptr;
+    HR_in.allocated = false;
+    HR_in.allocated_size = 0;
     // tmp terms not moved
 }
 
@@ -65,9 +74,12 @@ HContainer<T>& HContainer<T>::operator=(HContainer<T>&& HR_in) noexcept
         this->gamma_only = HR_in.gamma_only;
         this->paraV = HR_in.paraV;
         this->allocated = HR_in.allocated;
+        this->allocated_size = HR_in.allocated_size;
         this->current_R = -1;
 
         HR_in.wrapper_pointer = nullptr;
+        HR_in.allocated = false;
+        HR_in.allocated_size = 0;
     }
     return *this;
 }
@@ -80,6 +92,7 @@ HContainer<T>::HContainer(int natom)
     this->current_R = -1;
     this->sparse_ap.resize(natom);
     this->sparse_ap_index.resize(natom);
+    this->allocated_size = 0;
 }
 
 // use unitcell to initialize atom_pairs
@@ -88,6 +101,7 @@ HContainer<T>::HContainer(const UnitCell& ucell_, const Parallel_Orbitals* paraV
 {
     this->gamma_only = false;
     this->current_R = -1;
+    this->allocated_size = 0;
     std::vector<int> atom_begin_row(ucell_.nat+1, 0);
     std::vector<int> atom_begin_col(ucell_.nat+1, 0);
     int begin = 0;
@@ -148,6 +162,7 @@ template <typename T>
 HContainer<T>::HContainer(const Parallel_Orbitals* paraV_in, T* data_pointer, const std::vector<int>* ijr_info)
 {
     this->current_R = -1;
+    this->allocated_size = 0;
 
     // use HContainer as a wrapper(!nullptr) or container(nullptr)
     this->wrapper_pointer = data_pointer;
@@ -177,13 +192,23 @@ void HContainer<T>::allocate(T* data_array, bool is_zero)
     size_t nnr = this->get_nnr();
     if(this->allocated)
     {// delete existed memory of this->wrapper_pointer
+        if(this->allocated_size > 0)
+        {
+            ModuleBase::Memory::record("HContainer", -(long long)this->allocated_size, true);
+        }
         delete[] this->wrapper_pointer;
         this->allocated = false;
+        this->allocated_size = 0;
     }
     if(data_array == nullptr)
     {
         // use this->wrapper_pointer as data_array
         this->allocated = true;
+        this->allocated_size = nnr * sizeof(T);
+        if(this->allocated_size > 0)
+        {
+            ModuleBase::Memory::record("HContainer", (long long)this->allocated_size, true);
+        }
         this->wrapper_pointer = new T[nnr];
         ModuleBase::GlobalFunc::ZEROS(this->wrapper_pointer, nnr);
         data_array = this->wrapper_pointer;

source/source_lcao/module_hcontainer/hcontainer.h

@@ -475,6 +475,9 @@ class HContainer
     // sparse table for (atom_i, atom_j)->index of atom_pairs
     std::vector<std::vector<int>> sparse_ap;
     std::vector<std::vector<int>> sparse_ap_index;
+    
+    // record allocated memory size
+    size_t allocated_size = 0;
 
     /**
      * @brief temporary atom-pair lists to loop selected R index