Test: supplementary use of Memory::record to record memory (deepmodel…

…ing#4129)

* Test: supplementary use of Memory::record to record memory

* Test: supplementary use of ModuleBase::Memory::record (revise)

* Add records of init and matrix

* Fix the record of init_mixing_data

* Add records of Array_Pool and fftw

* Add records of Array_Pool and TensorBuffer

source/module_base/mathzone_add1.cpp

@@ -5,6 +5,7 @@
 #include "constants.h"
 #include "global_function.h"
 #include "math_sphbes.h"
+#include "module_base/memory.h"
 
 #if defined __FFTW2
 #include "fftw.h"
@@ -279,6 +280,7 @@ void Mathzone_Add1::Uni_Deriv_Phi
     fftw_complex *fft_phik = new fftw_complex[FFT_NR];
     fftw_complex *fft_ndphik = new fftw_complex[FFT_NR];
     fftw_complex *fft_ndphir = new fftw_complex[FFT_NR];
+	ModuleBase::Memory::record("Mathzone_Add1::Uni_Deriv_Phi",sizeof(fftw_complex) * FFT_NR * 4);
 	fftw_plan p1;
 	fftw_plan p2;
 

source/module_base/module_container/ATen/core/tensor.cpp

@@ -2,6 +2,7 @@
 #include <ATen/core/tensor_map.h>
 #include <ATen/core/tensor_utils.h>
 #include <base/core/cpu_allocator.h>
+#include "module_base/memory.h"
 #if defined(__CUDA) || defined(__ROCM)
 #include <base/core/gpu_allocator.h>
 #endif // __CUDA || __ROCM
@@ -216,6 +217,21 @@ void Tensor::resize(const TensorShape& new_shape) {
     shape_ = new_shape;
 }
 
+// Resize tensor object with the given tensor_shape ( Using for ModuleBase::Memory::record ) 
+void Tensor::resize(const TensorShape& new_shape, const std::string& record_str) {
+    if (shape_ == new_shape) {
+        return;
+    }
+    REQUIRES_OK(buffer_->OwnsMemory() || this->NumElements() == 0,
+        "Cannot resize a tensor that mapped from a given data buffer")
+    if (buffer_ && buffer_->GetAllocatedBytes() < new_shape.NumElements() * SizeOfType(data_type_)) {
+        buffer_->unref();
+        this->buffer_ = new TensorBuffer(GetAllocator(device_), new_shape.NumElements() * SizeOfType(data_type_));
+        ModuleBase::Memory::record(record_str,new_shape.NumElements() * SizeOfType(data_type_));
+    }
+    shape_ = new_shape;
+}
+
 Tensor& Tensor::operator=(const Tensor& other) {
     if (this == &other) {
         return *this;

source/module_base/module_container/ATen/core/tensor.h

@@ -350,6 +350,8 @@ class Tensor {
      */
     void resize(const TensorShape& new_shape);
 
+    void resize(const TensorShape& new_shape, const std::string& record_str); 
+
     /**
      * @brief Get the Allocator object according to the given device type.
      *

source/module_base/module_container/ATen/core/tensor_buffer.h

@@ -4,6 +4,7 @@
 #include <base/core/refcount.h>
 #include <base/core/allocator.h>
 #include <ATen/core/tensor_types.h>
+#include "module_base/memory.h"
 
 namespace container {
 

source/module_base/module_mixing/broyden_mixing.cpp

@@ -45,6 +45,7 @@ void Broyden_Mixing::tem_push_data(Mixing_Data& mdata,
 
     // container::Tensor data = data_in + mixing_beta * F;
     std::vector<FPTYPE> data(length);
+    ModuleBase::Memory::record("Broyden_Mixing::F_tmp&data",sizeof(FPTYPE)*length*2);
     mix(data.data(), data_in, F_tmp.data());
 
     mdata.push(data.data());
@@ -70,6 +71,7 @@ void Broyden_Mixing::tem_push_data(Mixing_Data& mdata,
         if (dF != nullptr)
             free(dF);
         dF = malloc(sizeof(FPTYPE) * length * mixing_ndim);
+        ModuleBase::Memory::record("Broyden_Mixing::F&DF",sizeof(FPTYPE)*length*(mixing_ndim+1));
         FP_dF = static_cast<FPTYPE*>(dF);
 #ifdef _OPENMP
 #pragma omp parallel for schedule(static, 4096 / sizeof(FPTYPE))

source/module_base/module_mixing/broyden_mixing.h

@@ -2,6 +2,7 @@
 #define BROYDEN_MIXING_H_
 #include "mixing.h"
 #include "module_base/matrix.h"
+#include "module_base/memory.h"
 
 namespace Base_Mixing
 {

source/module_basis/module_nao/two_center_table.cpp

@@ -9,6 +9,7 @@
 #include "module_base/constants.h"
 #include "module_base/math_integral.h"
 #include "module_base/cubic_spline.h"
+#include "module_base/memory.h"
 
 void TwoCenterTable::build(const RadialCollection& bra,
                            const RadialCollection& ket,
@@ -44,8 +45,8 @@ void TwoCenterTable::build(const RadialCollection& bra,
     ntab_ = 0;
     two_center_loop(bra, ket, &TwoCenterTable::_indexing);
 
-    table_.resize({ntab_, nr_});
-    dtable_.resize({ntab_, nr_});
+    table_.resize({ntab_, nr_},"TwoCenterTable::table_");
+    dtable_.resize({ntab_, nr_},"TwoCenterTable::dtable_");
     two_center_loop(bra, ket, &TwoCenterTable::_tabulate);
 }
 

source/module_cell/module_symmetry/symmetry.cpp

@@ -1,10 +1,12 @@
+#include <complex>
 #include <memory>
 #include <array>
 #include "symmetry.h"
 #include "module_base/libm/libm.h"
 #include "module_base/mathzone.h"
 #include "module_base/constants.h"
 #include "module_base/timer.h"
+#include "module_base/memory.h"
 
 namespace ModuleSymmetry
 {
@@ -300,6 +302,7 @@ int Symmetry::standard_lat(
     ModuleBase::Vector3<double> &c,
     double *cel_const) const
 {
+    // ModuleBase::TITLE("Symmetry", "standard_lat");
     static bool first = true;
     // there are only 14 types of Bravais lattice.
     int type = 15;
@@ -921,6 +924,7 @@ void Symmetry::checksym(ModuleBase::Matrix3 &s, ModuleBase::Vector3<double> &gtr
 	// is a valid symmetry operation on a supercell
 	//----------------------------------------------
     // the start atom index.
+    ModuleBase::TITLE("Symmetry", "checksym");
     bool no_diff = 0;
     ModuleBase::Vector3<double> trans(2.0, 2.0, 2.0);
     s_flag = 0;
@@ -1094,6 +1098,7 @@ void Symmetry::checksym(ModuleBase::Matrix3 &s, ModuleBase::Vector3<double> &gtr
 
 void Symmetry::pricell(double* pos, const Atom* atoms)
 {
+    ModuleBase::TITLE("Symmetry", "pricell");
     bool no_diff = 0;
     s_flag = 0;
     ptrans.clear();
@@ -1209,6 +1214,7 @@ void Symmetry::pricell(double* pos, const Atom* atoms)
 
     //sort ptrans:
     double* ptrans_array = new double[ntrans*3];
+    ModuleBase::Memory::record("Symmetry::ptrans_array",sizeof(double)*ntrans*3);
     for(int i=0;i<ntrans;++i)
     {
         ptrans_array[i*3]=ptrans[i].x;
@@ -1402,6 +1408,7 @@ void Symmetry::pricell(double* pos, const Atom* atoms)
 void Symmetry::rho_symmetry( double *rho,
                              const int &nr1, const int &nr2, const int &nr3)
 {
+    ModuleBase::TITLE("Symmetry", "rho_symmetry");
 //  if (GlobalV::test_symmetry)ModuleBase::TITLE("Symmetry","rho_symmetry");
     ModuleBase::timer::tick("Symmetry","rho_symmetry");
 
@@ -1460,6 +1467,7 @@ void Symmetry::rhog_symmetry(std::complex<double> *rhogtot,
     int* ixyz2ipw, const int &nx, const int &ny, const int &nz, 
     const int &fftnx, const int &fftny, const int &fftnz)
 {
+    ModuleBase::TITLE("Symmetry", "rhog_symmetry");
 //  if (GlobalV::test_symmetry)ModuleBase::TITLE("Symmetry","rho_symmetry");
     ModuleBase::timer::tick("Symmetry","rhog_symmetry");
 // ----------------------------------------------------------------------
@@ -1472,6 +1480,7 @@ void Symmetry::rhog_symmetry(std::complex<double> *rhogtot,
     int(*isymflag)[48] = new int[fftnx*fftny*fftnz][48];//which rotration operation the grid corresponds to
     int(*table_xyz)[48] = new int[fftnx * fftny * fftnz][48];// group information
     int* count_xyz = new int[fftnx * fftny * fftnz];// how many symmetry operations has been covered
+    ModuleBase::Memory::record("Symmetry::rhog_symmetry",sizeof(int) *fftnx*fftny*fftnz*98);
     for (int i = 0; i < fftnx * fftny * fftnz; i++)
     {
         symflag[i] = -1;
@@ -1766,6 +1775,7 @@ void Symmetry::symmetrize_vec3_nat(double* v)const   // pengfei 2016-12-20
 
 void Symmetry::symmetrize_mat3(ModuleBase::matrix& sigma, const Lattice& lat)const   //zhengdy added 2017
 {
+    ModuleBase::TITLE("Symmetry", "symmetrize_mat3");
     ModuleBase::matrix A = lat.latvec.to_matrix();
     ModuleBase::matrix AT = lat.latvec.Transpose().to_matrix();
     ModuleBase::matrix invA = lat.GT.to_matrix();
@@ -1781,6 +1791,7 @@ void Symmetry::symmetrize_mat3(ModuleBase::matrix& sigma, const Lattice& lat)con
 void Symmetry::gmatrix_convert_int(const ModuleBase::Matrix3* sa, ModuleBase::Matrix3* sb, 
         const int n, const ModuleBase::Matrix3 &a, const ModuleBase::Matrix3 &b) const
 {
+    ModuleBase::TITLE("Symmetry", "gmatrix_convert_int");
     auto round = [](double x){return (x>0.0)?floor(x+0.5):ceil(x-0.5);};
     ModuleBase::Matrix3 ai = a.Inverse();
     ModuleBase::Matrix3 bi = b.Inverse();
@@ -1802,6 +1813,7 @@ void Symmetry::gmatrix_convert_int(const ModuleBase::Matrix3* sa, ModuleBase::Ma
 void Symmetry::gmatrix_convert(const ModuleBase::Matrix3* sa, ModuleBase::Matrix3* sb, 
         const int n, const ModuleBase::Matrix3 &a, const ModuleBase::Matrix3 &b)const
 {
+    ModuleBase::TITLE("Symmetry", "gmatrix_convert");
     ModuleBase::Matrix3 ai = a.Inverse();
     ModuleBase::Matrix3 bi = b.Inverse();
     for (int i=0;i<n;++i)
@@ -1820,6 +1832,7 @@ void Symmetry::gtrans_convert(const ModuleBase::Vector3<double>* va, ModuleBase:
 }
 void Symmetry::gmatrix_invmap(const ModuleBase::Matrix3* s, const int n, int* invmap)
 {
+    ModuleBase::TITLE("Symmetry", "gmatrix_invmap");
     ModuleBase::Matrix3 eig(1, 0, 0, 0, 1, 0, 0, 0, 1);
     ModuleBase::Matrix3 tmp;
     for (int i=0;i<n;++i)
@@ -1842,6 +1855,7 @@ void Symmetry::gmatrix_invmap(const ModuleBase::Matrix3* s, const int n, int* in
 void Symmetry::get_shortest_latvec(ModuleBase::Vector3<double> &a1, 
         ModuleBase::Vector3<double> &a2, ModuleBase::Vector3<double> &a3) const
 {
+    ModuleBase::TITLE("Symmetry", "get_shortest_latvec");
     double len1=a1.norm();
     double len2=a2.norm();
     double len3=a3.norm();
@@ -1891,6 +1905,7 @@ void Symmetry::get_optlat(ModuleBase::Vector3<double> &v1, ModuleBase::Vector3<d
         ModuleBase::Vector3<double> &w2, ModuleBase::Vector3<double> &w3, 
         int& real_brav, double* cel_const, double* tmp_const) const
 {
+    ModuleBase::TITLE("Symmetry", "get_optlat");
     ModuleBase::Vector3<double> r1, r2, r3;
     double cos1 = 1;
     double cos2 = 1;
@@ -2152,6 +2167,7 @@ bool Symmetry::magmom_same_check(const Atom* atoms)const
 
 bool Symmetry::is_all_movable(const Atom* atoms, const Statistics& st)const
 {
+    ModuleBase::TITLE("Symmetry", "is_all_movable");
     bool all_mbl = true;
     for (int iat = 0;iat < st.nat;++iat)
     {

source/module_elecstate/module_charge/charge_mixing.cpp

@@ -8,6 +8,7 @@
 #include "module_base/parallel_reduce.h"
 #include "module_base/timer.h"
 #include "module_hamilt_pw/hamilt_pwdft/global.h"
+#include "module_base/memory.h"
 
 Charge_Mixing::Charge_Mixing()
 {
@@ -123,23 +124,27 @@ void Charge_Mixing::init_mixing()
             this->mixing->init_mixing_data(this->rho_mdata,
                                         this->rhopw->npw * 2,
                                         sizeof(std::complex<double>));
+            ModuleBase::Memory::record("Charge_Mixing::init_mixing_data",sizeof(std::complex<double>)*this->mixing->data_ndim*this->rhopw->npw * 2);
         }
         else
         {
             this->mixing->init_mixing_data(this->rho_mdata,
                                         this->rhopw->npw * GlobalV::NSPIN,
                                         sizeof(std::complex<double>));
+            ModuleBase::Memory::record("Charge_Mixing::init_mixing_data",sizeof(std::complex<double>)*this->mixing->data_ndim*this->rhopw->npw * GlobalV::NSPIN);
         }
     }
     else
     {
         if (GlobalV::NSPIN == 4 && GlobalV::MIXING_ANGLE > 0 )
         {
             this->mixing->init_mixing_data(this->rho_mdata, this->rhopw->nrxx * 2, sizeof(double));
+            ModuleBase::Memory::record("Charge_Mixing::init_mixing_data",sizeof(double)*this->mixing->data_ndim*this->rhopw->nrxx * 2);
         }
         else
         {
             this->mixing->init_mixing_data(this->rho_mdata, this->rhopw->nrxx * GlobalV::NSPIN, sizeof(double));
+            ModuleBase::Memory::record("Charge_Mixing::init_mixing_data",sizeof(double)*this->mixing->data_ndim*this->rhopw->nrxx * GlobalV::NSPIN);
         }
     }
 
@@ -151,10 +156,12 @@ void Charge_Mixing::init_mixing()
             this->mixing->init_mixing_data(this->tau_mdata,
                                            this->rhopw->npw * GlobalV::NSPIN,
                                            sizeof(std::complex<double>));
+            ModuleBase::Memory::record("Charge_Mixing::init_mixing_data",sizeof(std::complex<double>)*this->mixing->data_ndim*this->rhopw->npw * GlobalV::NSPIN);
         }
         else
         {
             this->mixing->init_mixing_data(this->tau_mdata, this->rhopw->nrxx * GlobalV::NSPIN, sizeof(double));
+            ModuleBase::Memory::record("Charge_Mixing::init_mixing_data",sizeof(double)*this->mixing->data_ndim*this->rhopw->nrxx * GlobalV::NSPIN);
         }
     }
 
@@ -185,6 +192,7 @@ void Charge_Mixing::allocate_mixing_dmr(int nnr)
     else if (GlobalV::SCF_THR_TYPE == 2)
     {
         this->mixing->init_mixing_data(this->dmr_mdata, nnr * dmr_nspin, sizeof(double));
+        ModuleBase::Memory::record("Charge_Mixing::init_mixing_dat_dmr",sizeof(double)*this->mixing->data_ndim*nnr*dmr_nspin);
     }
 
     this->dmr_mdata.reset();
@@ -218,6 +226,7 @@ double Charge_Mixing::get_drho(Charge* chr, const double nelec)
 
         ModuleBase::GlobalFunc::NOTE("Calculate the charge difference between rho(G) and rho_save(G)");
         std::vector<std::complex<double>> drhog(GlobalV::NSPIN * this->rhopw->npw);
+        ModuleBase::Memory::record("Charge_Mixing::drhog",sizeof(std::complex<double>)*GlobalV::NSPIN * this->rhopw->npw);
 #ifdef _OPENMP
 #pragma omp parallel for collapse(2) schedule(static, 512)
 #endif
@@ -343,6 +352,7 @@ void Charge_Mixing::mix_rho_recip(Charge* chr)
         // allocate rhog_mag[is*ngmc] and rhog_mag_save[is*ngmc]
         rhog_mag = new std::complex<double>[npw * GlobalV::NSPIN];
         rhog_mag_save = new std::complex<double>[npw * GlobalV::NSPIN];
+        ModuleBase::Memory::record("Charge_Mixing::rhog_mag*",sizeof(std::complex<double>)*npw * GlobalV::NSPIN*2);
         ModuleBase::GlobalFunc::ZEROS(rhog_mag, npw * GlobalV::NSPIN);
         ModuleBase::GlobalFunc::ZEROS(rhog_mag_save, npw * GlobalV::NSPIN);
         // get rhog_mag[is*ngmc] and rhog_mag_save[is*ngmc]
@@ -447,6 +457,7 @@ void Charge_Mixing::mix_rho_recip(Charge* chr)
         const int nrxx = this->rhopw->nrxx;
         double* rho_magabs = new double[nrxx];
         double* rho_magabs_save = new double[nrxx];
+        ModuleBase::Memory::record("Charge_Mixing::rho_magabs*",sizeof(double)*nrxx*2);
         ModuleBase::GlobalFunc::ZEROS(rho_magabs, nrxx);
         ModuleBase::GlobalFunc::ZEROS(rho_magabs_save, nrxx);
         // calculate rho_magabs and rho_magabs_save
@@ -461,6 +472,7 @@ void Charge_Mixing::mix_rho_recip(Charge* chr)
         const int npw = this->rhopw->npw;
         std::complex<double>* rhog_magabs = new std::complex<double>[npw * 2];
         std::complex<double>* rhog_magabs_save = new std::complex<double>[npw * 2];
+        ModuleBase::Memory::record("Charge_Mixing::rhog_magabs_2*",sizeof(std::complex<double>)*npw*4);
         ModuleBase::GlobalFunc::ZEROS(rhog_magabs, npw * 2);
         ModuleBase::GlobalFunc::ZEROS(rhog_magabs_save, npw * 2);
         // calculate rhog_magabs and rhog_magabs_save
@@ -632,6 +644,7 @@ void Charge_Mixing::mix_rho_real(Charge* chr)
         // allocate rho_mag[is*nnrx] and rho_mag_save[is*nnrx]
         rho_mag = new double[nrxx * GlobalV::NSPIN];
         rho_mag_save = new double[nrxx * GlobalV::NSPIN];
+        ModuleBase::Memory::record("Charge_Mixing::rho_mag_2*",sizeof(double)*nrxx * GlobalV::NSPIN*2);
         ModuleBase::GlobalFunc::ZEROS(rho_mag, nrxx * GlobalV::NSPIN);
         ModuleBase::GlobalFunc::ZEROS(rho_mag_save, nrxx * GlobalV::NSPIN);
         // get rho_mag[is*nnrx] and rho_mag_save[is*nnrx]
@@ -726,6 +739,7 @@ void Charge_Mixing::mix_rho_real(Charge* chr)
         // allocate memory for rho_magabs and rho_magabs_save
         double* rho_magabs = new double[nrxx * 2];
         double* rho_magabs_save = new double[nrxx * 2];
+        ModuleBase::Memory::record("Charge_Mixing::rho_magabs_real*",sizeof(double)*nrxx*4);
         ModuleBase::GlobalFunc::ZEROS(rho_magabs, nrxx * 2);
         ModuleBase::GlobalFunc::ZEROS(rho_magabs_save, nrxx * 2);
         // calculate rho_magabs and rho_magabs_save
@@ -822,6 +836,7 @@ void Charge_Mixing::mix_dmr(elecstate::DensityMatrix<double, double>* DM)
         // allocate dmr_mag[is*nnrx] and dmr_mag_save[is*nnrx]
         dmr_mag = new double[nnr * GlobalV::NSPIN];
         dmr_mag_save = new double[nnr * GlobalV::NSPIN];
+        ModuleBase::Memory::record("Charge_Mixing::dmr_mag*",sizeof(double)*nnr*GlobalV::NSPIN*2);
         ModuleBase::GlobalFunc::ZEROS(dmr_mag, nnr * GlobalV::NSPIN);
         ModuleBase::GlobalFunc::ZEROS(dmr_mag_save, nnr * GlobalV::NSPIN);
         double* dmr_up;
@@ -921,6 +936,7 @@ void Charge_Mixing::mix_dmr(elecstate::DensityMatrix<std::complex<double>, doubl
         // allocate dmr_mag[is*nnrx] and dmr_mag_save[is*nnrx]
         dmr_mag = new double[nnr * GlobalV::NSPIN];
         dmr_mag_save = new double[nnr * GlobalV::NSPIN];
+        ModuleBase::Memory::record("Charge_Mixing::dmr_mag_2*",sizeof(double)*nnr*GlobalV::NSPIN*2);
         ModuleBase::GlobalFunc::ZEROS(dmr_mag, nnr * GlobalV::NSPIN);
         ModuleBase::GlobalFunc::ZEROS(dmr_mag_save, nnr * GlobalV::NSPIN);
         double* dmr_up;
@@ -1015,6 +1031,7 @@ void Charge_Mixing::mix_rho(Charge* chr)
     // the charge before mixing.
     const int nrxx = chr->rhopw->nrxx;
     std::vector<double> rho123(GlobalV::NSPIN * nrxx);
+    ModuleBase::Memory::record("Charge_Mixing::rho123",sizeof(double)*GlobalV::NSPIN * nrxx);
     for (int is = 0; is < GlobalV::NSPIN; ++is)
     {
         if (is == 0 || is == 3 || !GlobalV::DOMAG_Z)
@@ -1262,6 +1279,7 @@ double Charge_Mixing::inner_product_recip_rho(std::complex<double>* rho1, std::c
 
     std::complex<double>** rhog1 = new std::complex<double>*[GlobalV::NSPIN];
     std::complex<double>** rhog2 = new std::complex<double>*[GlobalV::NSPIN];
+    ModuleBase::Memory::record("Charge_Mixing::rhog1&2",sizeof(std::complex<double>)*GlobalV::NSPIN*2);
     for (int is = 0; is < GlobalV::NSPIN; is++)
     {
         rhog1[is] = rho1 + is * this->rhopw->npw;
@@ -1623,10 +1641,13 @@ void Charge_Mixing::divide_data(std::complex<double>* data_d,
         const int ndims = this->rhopw->npw;
         const int ndimhf = ndimd - ndims;
         data_s = new std::complex<double>[GlobalV::NSPIN * ndims];
+        ModuleBase::Memory::record("Charge_Mixing::data_s",sizeof(std::complex<double>)*GlobalV::NSPIN * ndims);
         data_hf = nullptr;
         if (ndimhf > 0)
         {
             data_hf = new std::complex<double>[GlobalV::NSPIN * ndimhf];
+
+            ModuleBase::Memory::record("Charge_Mixing::data_hf",sizeof(std::complex<double>)*GlobalV::NSPIN * ndimhf);
         }
         for (int is = 0; is < GlobalV::NSPIN; ++is)
         {

source/module_elecstate/module_charge/symmetry_rhog.cpp

@@ -24,6 +24,7 @@ void Symmetry_rho::psymmg(std::complex<double>* rhog_part, const ModulePW::PW_Ba
 		ModuleBase::GlobalFunc::ZEROS(rhogtot, rho_basis->npwtot);
 		ig2isztot = new int[rho_basis->npwtot];
 		ModuleBase::GlobalFunc::ZEROS(rhogtot, rho_basis->npwtot);
+		ModuleBase::Memory::record("Symmetry_rho::psymmg",sizeof(std::complex<double>)*rho_basis->npwtot + sizeof(int)*(rho_basis->npwtot+rho_basis->fftnxy));
 	}
 	// find max_npw
 	int max_npw=0;