Refactor: new memory record interface (#1794)

* Refactor: new memory record interface

* Fix: CUDA and ROCM compiler

* Fix: UT related to memory.cpp

Co-authored-by: dyzheng <[email protected]>

Author: dyzheng

source/module_base/math_ylmreal.cpp

@@ -6,6 +6,7 @@
 #include <cassert>
 #include "ylm.h"
 #include "module_base/kernels/math_op.h"
+#include "module_psi/kernels/memory_op.h"
 
 namespace ModuleBase
 {
@@ -323,7 +324,7 @@ void YlmReal::Ylm_Real(Device * ctx, const int lmax2, const int ng, const FPTYPE
         ModuleBase::WARNING_QUIT("YLM_REAL","l>30 or l<0");
     }
     FPTYPE * p = nullptr, * phi = nullptr, * cost = nullptr;
-    resmem_var_op()(ctx, p, (lmax + 1) * (lmax + 1) * ng);
+    resmem_var_op()(ctx, p, (lmax + 1) * (lmax + 1) * ng, "YlmReal::Ylm_Real");
 
     cal_ylm_real_op()(
         ctx,

source/module_base/math_ylmreal.h

@@ -3,7 +3,6 @@
 
 #include "vector3.h"
 #include "matrix.h"
-#include "module_psi/psi.h"
 
 namespace ModuleBase
 {

source/module_base/memory.cpp

@@ -3,6 +3,8 @@
 // DATE : 2008-11-18
 //==========================================================
 #include "memory.h"
+#include "global_variable.h"
+#include "src_parallel/parallel_reduce.h"
 
 namespace ModuleBase
 {
@@ -18,7 +20,7 @@ int Memory::bool_memory = sizeof(bool); // 1.0 Byte
 int Memory::float_memory = sizeof(float); // 4.0 Byte
 int Memory::short_memory = sizeof(short); // 2.0 Byte
 
-int Memory::n_memory = 500;
+int Memory::n_memory = 1000;
 int Memory::n_now = 0;
 bool Memory::init_flag =  false;
 
@@ -140,10 +142,75 @@ double Memory::record
 	return consume[find];
 }
 
+void Memory::record
+(
+	const std::string &name_in,
+	const size_t &n_in,
+	const bool accumulate
+)
+{
+	if(!Memory::init_flag)
+	{
+		name = new std::string[n_memory];
+		class_name = new std::string[n_memory];
+		consume = new double[n_memory];
+		for(int i=0;i<n_memory;i++)
+		{
+			consume[i] = 0.0;
+		}
+		Memory::init_flag = true;
+	}
+
+	int find = 0;
+	for(find = 0; find < n_now; find++)
+	{
+		if( name_in == name[find] )
+		{
+			break;
+		}
+	}
+
+	// find == n_now : found a new record.	
+	if(find == n_now)
+	{
+		n_now++;
+		name[find] = name_in;
+		class_name[find] = "";
+	}
+	if(n_now >= n_memory)
+	{
+		std::cout<<" Error! Too many memories has been recorded.";
+		return;
+	}
+
+	const double factor = 1.0/1024.0/1024.0;
+	double size_mb = n_in * factor;
+
+	if(accumulate)
+	{
+		consume[find] += size_mb;
+		Memory::total += size_mb;
+	}
+	else
+	{
+		if(consume[find] < size_mb)
+		{
+			Memory::total += size_mb - consume[find];
+			consume[find] = size_mb;
+			if(consume[find] > 5)
+			{
+				print(find);
+			}
+		}
+	}
+
+	return;
+}
+
 void Memory::print(const int find)
 {
-//	std::cout <<"\n Warning_Memory_Consuming : "
-//	<<class_name[find]<<" "<<name[find]<<" "<<consume[find]<<" MB" << std::endl;
+	GlobalV::ofs_running <<"\n Warning_Memory_Consuming allocated: "
+	<<" "<<name[find]<<" "<<consume[find]<<" MB" << std::endl;
 	return;
 }
 
@@ -167,10 +234,12 @@ void Memory::print_all(std::ofstream &ofs)
 	if(!init_flag) return;
 
 	const double small = 1.0; 
-//    std::cout<<"\n CLASS_NAME---------|NAME---------------|MEMORY(MB)--------";
-    ofs <<"\n CLASS_NAME---------|NAME---------------|MEMORY(MB)--------" << std::endl;
+#ifdef __MPI
+		Parallel_Reduce::reduce_double_all(Memory::total);
+#endif
+    ofs <<"\n NAME---------------|MEMORY(MB)--------" << std::endl;
 //	std::cout<<"\n"<<std::setw(41)<< " " <<std::setprecision(4)<<total;
-	ofs <<std::setw(41)<< " " <<std::setprecision(4)<<total << std::endl;
+	ofs <<std::setw(20)<< "total" << std::setw(15) <<std::setprecision(4)<< Memory::total << std::endl;
 
 	bool *print_flag = new bool[n_memory];
 	for(int i=0; i<n_memory; i++) print_flag[i] = false;
@@ -192,16 +261,16 @@ void Memory::print_all(std::ofstream &ofs)
 			}
 		}
 		print_flag[k] = true;
-
+#ifdef __MPI
+		Parallel_Reduce::reduce_double_all(consume[k]);
+#endif
 	    if ( consume[k] < small ) 
         {
             continue;
         }
   		else
   		{
-        	ofs  << " "
-             << std::setw(20) << class_name[k]
-             << std::setw(20) << name[k]
+        	ofs << std::setw(20) << name[k]
              << std::setw(15) << consume[k] << std::endl;
 
 //        	std::cout  << "\n "
@@ -211,6 +280,7 @@ void Memory::print_all(std::ofstream &ofs)
 		}
     }
 //    std::cout<<"\n ----------------------------------------------------------"<<std::endl;
+	ofs<<" -------------   < 1.0 MB has been ignored ----------------"<<std::endl;
     ofs<<" ----------------------------------------------------------"<<std::endl;
 	delete[] print_flag; //mohan fix by valgrind at 2012-04-02
 	return;

source/module_base/memory.h

@@ -37,6 +37,19 @@ class Memory
                          const std::string &type,
                          const bool accumulate = false);
 
+    /**
+     * @brief Record memory consumed during computation
+     *
+     * @param name The name of a quantity
+     * @param n The number of the quantity
+     * @param accumulate Useless, always set false
+     */
+    static void record(
+      const std::string &name_in,
+      const size_t &n_in,
+      const bool accumulate = false
+    );
+
     static double &get_total(void)
     {
         return total;

source/module_base/test/CMakeLists.txt

@@ -68,6 +68,7 @@ AddTest(
   LIBS ${math_libs} device
   SOURCES math_ylmreal_test.cpp ../math_ylmreal.cpp ../complexmatrix.cpp ../global_variable.cpp ../ylm.cpp ../realarray.cpp ../timer.cpp ../matrix.cpp ../vector3.h
           ../../src_parallel/parallel_reduce.cpp ../../src_parallel/parallel_kpoints.cpp ../../src_parallel/parallel_global.cpp ../../src_parallel/parallel_common.cpp
+          ../memory.cpp
 )
 AddTest(
   TARGET base_math_sphbes

source/module_base/test/memory_test.cpp

@@ -4,6 +4,11 @@
 #include <fstream>
 #include <cstdio>
 
+namespace GlobalV
+{
+	std::ofstream ofs_running;
+}
+
 /************************************************
  *  unit test of class Memory
  ***********************************************/

source/module_deepks/test/klist_1.cpp

@@ -105,10 +105,10 @@ namespace Test_Deepks
 		wk = new double[kpoint_number];
 		isk = new int[kpoint_number];
 
-		ModuleBase::Memory::record("K_Vectors","kvec_c",kpoint_number*3,"double");
-		ModuleBase::Memory::record("K_Vectors","kvec_d",kpoint_number*3,"double");
-		ModuleBase::Memory::record("K_Vectors","wk",kpoint_number*3,"double");
-		ModuleBase::Memory::record("K_Vectors","isk",kpoint_number*3,"int");
+		ModuleBase::Memory::record("KV::kvec_c",sizeof(double) * kpoint_number*3);
+		ModuleBase::Memory::record("KV::kvec_d",sizeof(double) * kpoint_number*3);
+		ModuleBase::Memory::record("KV::wk",sizeof(double) * kpoint_number*3);
+		ModuleBase::Memory::record("KV::isk",sizeof(int) * kpoint_number*3);
 
 		return;
 	}

source/module_deepks/test/parallel_orbitals.cpp

@@ -35,8 +35,8 @@ namespace Test_Deepks
             trace_loc_col[i] = -1;
         }
 
-        ModuleBase::Memory::record("Parallel_Orbitals","trace_loc_row",GlobalV::NLOCAL,"int");
-        ModuleBase::Memory::record("Parallel_Orbitals","trace_loc_col",GlobalV::NLOCAL,"int");
+        ModuleBase::Memory::record("PO::trace_loc_row",sizeof(int) * GlobalV::NLOCAL);
+        ModuleBase::Memory::record("PO::trace_loc_col",sizeof(int) * GlobalV::NLOCAL);
 
         for (int i=0; i<GlobalV::NLOCAL; i++)
         {

source/module_dftu/dftu.cpp

@@ -211,7 +211,7 @@ void DFTU::init(UnitCell& cell, // unitcell class
         }
     }
 
-    ModuleBase::Memory::record("DFTU", "locale", num_locale, "double");
+    ModuleBase::Memory::record("DFTU::locale", sizeof(double) * num_locale);
     return;
 }
 

source/module_elecstate/elecstate_pw.cpp

@@ -51,8 +51,8 @@ void ElecStatePW<FPTYPE, Device>::init_rho_data()
             this->kin_r = reinterpret_cast<FPTYPE **>(this->charge->kin_r);
         }
     }
-    resmem_complex_op()(this->ctx, this->wfcr, this->basis->nmaxgr);
-    resmem_complex_op()(this->ctx, this->wfcr_another_spin, this->charge->nrxx);
+    resmem_complex_op()(this->ctx, this->wfcr, this->basis->nmaxgr, "ElecSPW::wfcr");
+    resmem_complex_op()(this->ctx, this->wfcr_another_spin, this->charge->nrxx, "ElecSPW::wfcr_a");
     this->init_rho = true;
 }