Refactor: add orbital parallel algorithm for sdft

Qianruipku · Qianruipku · commit 8094ec80081e · 2022-05-11T22:11:59.000+08:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -17,7 +17,7 @@ option(USE_CUDA "Enable support to CUDA." OFF)
 option(USE_ROCM "Enable support to ROCm." OFF)
 option(USE_OPENMP " Enable OpenMP in abacus." ON)
 option(ENABLE_ASAN "Enable AddressSanitizer" OFF)
-option(BUILD_TESTING "Build ABACUS unit tests" ON)
+option(BUILD_TESTING "Build ABACUS unit tests" OFF)
 option(GENERATE_TEST_REPORTS "Enable test report generation" OFF)
 
 set(ABACUS_BIN_NAME abacus)
diff --git a/source/Makefile.Objects b/source/Makefile.Objects
@@ -223,7 +223,6 @@ parallel_reduce.o\
 parallel_pw.o\
 ft.o\
 parallel_grid.o\
-parallel_stochi.o\
 
 OBJS_ESOLVER=esolver.o\
     esolver_ks.o\
diff --git a/source/driver.cpp b/source/driver.cpp
@@ -96,7 +96,8 @@ void Driver::atomic_world(void)
 	ModuleESolver::ESolver *p_esolver;
 	if(GlobalV::BASIS_TYPE=="pw" || GlobalV::BASIS_TYPE=="lcao_in_pw")
 	{
-		use_ensol = "ksdft_pw"; 
+		if(GlobalV::CALCULATION.substr(0,3) == "sto")	use_ensol = "sdft_pw";
+		else											use_ensol = "ksdft_pw"; 
 		//We set it temporarily
 		//Finally, we have ksdft_pw, ksdft_lcao, sdft_pw, ofdft, lj, eam, etc.
 		ModuleESolver::init_esolver(p_esolver, use_ensol);
diff --git a/source/input.cpp b/source/input.cpp
@@ -140,7 +140,6 @@ void Input::Default(void)
     pw_seed = 1;
     nche_sto = 0;
     seed_sto = 0;
-    stotype = "pw";
     kpar = 1;
     berry_phase = false;
     gdir = 3;
@@ -552,9 +551,9 @@ bool Input::Read(const std::string &fn)
         {
             read_value(ifs, emin_sto);
         }
-        else if (strcmp("stotype", word) == 0)
+        else if (strcmp("nstogroup", word) == 0)
         {
-            read_value(ifs, stotype);
+            read_value(ifs, nstogroup);
         }
         else if (strcmp("kpar", word) == 0) // number of pools
         {
@@ -1798,6 +1797,7 @@ void Input::Default_2(void) // jiyy add 2019-08-04
             vdw_radius = "95";
         }
     }
+    if(calculation.substr(0,3) != "sto")    nstogroup = 1;
 }
 #ifdef __MPI
 void Input::Bcast()
@@ -1828,7 +1828,7 @@ void Input::Bcast()
     Parallel_Common::bcast_int(pw_seed);
     Parallel_Common::bcast_double(emax_sto);
     Parallel_Common::bcast_double(emin_sto);
-    Parallel_Common::bcast_string(stotype);
+    Parallel_Common::bcast_int(nstogroup);
     Parallel_Common::bcast_int(kpar);
     Parallel_Common::bcast_bool(berry_phase);
     Parallel_Common::bcast_int(gdir);
@@ -2174,7 +2174,7 @@ void Input::Check(void)
         */
         this->relax_nmax = 1;
     }
-    else if (calculation == "scf-sto") // qianrui 2021-2-20
+    else if (calculation == "sto-scf") // qianrui 2021-2-20
     {
         if (mem_saver == 1)
         {
diff --git a/source/input.h b/source/input.h
@@ -68,7 +68,7 @@ class Input
     int seed_sto; // random seed for sDFT
     double emax_sto; // Emax & Emin to normalize H
     double emin_sto;
-    std::string stotype;
+    int nstogroup;
 
     //==========================================================
     // electrons / spin
diff --git a/source/input_conv.cpp b/source/input_conv.cpp
@@ -61,6 +61,7 @@ void Input_Conv::Convert(void)
     GlobalV::KPAR = temp_nproc;
 #else
     GlobalV::KPAR = INPUT.kpar;
+    GlobalV::NSTOGROUP = INPUT.nstogroup;
 #endif
     GlobalV::CALCULATION = INPUT.calculation;
 
diff --git a/source/module_elecstate/elecstate_pw.cpp b/source/module_elecstate/elecstate_pw.cpp
@@ -41,15 +41,19 @@ void ElecStatePW::parallelK()
 {
 #ifdef __MPI
     charge->rho_mpi();
-    if (GlobalV::CALCULATION != "scf-sto" && GlobalV::CALCULATION != "relax-sto"
-        && GlobalV::CALCULATION != "md-sto") // qinarui add it temporarily.
-    {
-        //==================================
-        // Reduce all the Energy in each cpu
-        //==================================
-        this->eband /= GlobalV::NPROC_IN_POOL;
-        Parallel_Reduce::reduce_double_all(this->eband);
-    }
+    if(GlobalV::CALCULATION.substr(0,3) == "sto") //qinarui add it 2021-7-21
+	{
+		GlobalC::en.eband /= GlobalV::NPROC_IN_POOL;
+		MPI_Allreduce(MPI_IN_PLACE, &GlobalC::en.eband, 1, MPI_DOUBLE, MPI_SUM , STO_WORLD);
+	}
+	else
+	{
+    	//==================================
+    	// Reduce all the Energy in each cpu
+    	//==================================
+		GlobalC::en.eband /= GlobalV::NPROC_IN_POOL;
+		Parallel_Reduce::reduce_double_all( GlobalC::en.eband );
+	}
 #endif
     return;
 }
diff --git a/source/module_esolver/esolver_ks.cpp b/source/module_esolver/esolver_ks.cpp
@@ -53,10 +53,17 @@ void ESolver_KS:: Run(const int istep, UnitCell_pseudo& cell)
         clock_t iterstart,iterend;
         iterstart = std::clock();
         set_ethr(istep,iter);
-        eachiterinit(iter); 
+        eachiterinit(iter);
         
         this->hamilt2density(istep, iter, this->diag_ethr);
         
+        //<Temporary> It may be changed when more clever parallel algorithm is put forward.
+        //When parallel algorithm for bands are adopted. Density will only be treated in the first group.
+        //(Different ranks should have abtained the same, but small differences always exist in practice.)
+        //Maybe in the future, density and wavefunctions should use different parallel algorithms, in which 
+        //they do not occupy all processors, for example wavefunctions uses 20 processors while density uses 10.
+        if(GlobalV::MY_STOGROUP == 0)
+        {
         // double drho = this->estate.caldr2(); 
         // EState should be used after it is constructed.
         drho = GlobalC::CHR.get_drho();
@@ -88,6 +95,13 @@ void ESolver_KS:: Run(const int istep, UnitCell_pseudo& cell)
             //conv_elec = this->estate.mix_rho();
             GlobalC::CHR.mix_rho(iter);
         }
+        
+        }
+#ifdef __MPI
+		MPI_Bcast(&drho, 1, MPI_DOUBLE , 0, PARAPW_WORLD);
+		MPI_Bcast(&conv_elec, 1, MPI_DOUBLE , 0, PARAPW_WORLD);
+		MPI_Bcast(GlobalC::CHR.rho[0], GlobalC::pw.nrxx, MPI_DOUBLE, 0, PARAPW_WORLD);
+#endif       
 
         // Hamilt should be used after it is constructed.
         // this->phamilt->update(conv_elec);
diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
@@ -151,7 +151,8 @@ void ESolver_KS_PW::Init(Input &inp, UnitCell_pseudo &ucell)
     //================================
     // Initial start wave functions
     //================================
-    if (GlobalV::NBANDS != 0 || (GlobalV::CALCULATION != "scf-sto" && GlobalV::CALCULATION != "relax-sto" && GlobalV::CALCULATION != "md-sto")) //qianrui add
+    if (GlobalV::NBANDS != 0 || GlobalV::CALCULATION.substr(0,3) != "sto")
+    // qianrui add temporarily. In the future, wfcinit() should be compatible with cases when NBANDS=0
     {
         GlobalC::wf.wfcinit();
     }
diff --git a/source/module_esolver/esolver_sdft_pw.cpp b/source/module_esolver/esolver_sdft_pw.cpp
@@ -205,6 +205,12 @@ void ESolver_SDFT_PW::hamilt2density(int istep, int iter, double ethr)
 		{
 			if(ModuleSymmetry::Symmetry::symm_flag)	MPI_Barrier(MPI_COMM_WORLD);
 		}
+
+		if(GlobalV::MY_STOGROUP == 0)
+		{
+        	GlobalC::en.deband = GlobalC::en.delta_e();
+		}
+
 }
 
 void ESolver_SDFT_PW:: c_bands_k(const int ik, double* h_diag, const int istep, const int iter)
diff --git a/source/src_io/print_info.cpp b/source/src_io/print_info.cpp
@@ -12,7 +12,7 @@ void Print_Info::setup_parameters(UnitCell_pseudo &ucell, K_Vectors &kv)
 	ModuleBase::TITLE("Print_Info","setup_parameters");
 	
     if(GlobalV::CALCULATION=="scf" || GlobalV::CALCULATION=="relax" || GlobalV::CALCULATION=="cell-relax" || GlobalV::CALCULATION=="nscf"
-	        || GlobalV::CALCULATION=="istate" || GlobalV::CALCULATION=="ienvelope" || GlobalV::CALCULATION=="md")
+	        || GlobalV::CALCULATION=="istate" || GlobalV::CALCULATION=="ienvelope" || GlobalV::CALCULATION=="md"||GlobalV::CALCULATION.substr(0,3) == "sto")
 	{
 		std::cout << " ---------------------------------------------------------" << std::endl;
 		if(GlobalV::CALCULATION=="scf")
diff --git a/source/src_io/read_rho.cpp b/source/src_io/read_rho.cpp
@@ -99,7 +99,7 @@ bool Charge::read_rho(const int &is, const std::string &fn, double* rho) //add b
 	for(int iz=0; iz<GlobalC::pw.ncz; iz++)
 	{
 		ModuleBase::GlobalFunc::ZEROS(zpiece, nxy);
-		if(GlobalV::MY_RANK==0)
+		if(GlobalV::MY_RANK==0||(GlobalV::CALCULATION.substr(0,3) == "sto"&&GlobalV::RANK_IN_STOGROUP==0))
 		{
 			//				GlobalV::ofs_running << " Read charge density iz=" << iz << std::endl;
 			for(int j=0; j<GlobalC::pw.ncy; j++)
@@ -115,6 +115,6 @@ bool Charge::read_rho(const int &is, const std::string &fn, double* rho) //add b
 	delete[] zpiece;
 #endif
 
-    if(GlobalV::MY_RANK==0) ifs.close();
+    if(GlobalV::MY_RANK==0||(GlobalV::CALCULATION.substr(0,3) == "sto"&&GlobalV::RANK_IN_STOGROUP==0)) ifs.close();
     return true;
 }
diff --git a/source/src_parallel/CMakeLists.txt b/source/src_parallel/CMakeLists.txt
@@ -6,7 +6,6 @@ list(APPEND objects
     parallel_kpoints.cpp
     parallel_pw.cpp
     parallel_reduce.cpp
-    parallel_stochi.cpp
     subgrid_oper.cpp
 )
 
diff --git a/source/src_parallel/parallel_grid.cpp b/source/src_parallel/parallel_grid.cpp
@@ -67,12 +67,16 @@ void Parallel_Grid::init(
 	assert(GlobalV::KPAR > 0);
 
 	this->nproc_in_pool = new int[GlobalV::KPAR];
-	const int remain_pro = GlobalV::NPROC%GlobalV::KPAR;
+	int nprocgroup;
+	if(GlobalV::CALCULATION.substr(0,3)=="sto")		nprocgroup = GlobalV::NPROC_IN_STOGROUP;
+	else											nprocgroup = GlobalV::NPROC;
+
+	const int remain_pro = nprocgroup%GlobalV::KPAR;
 	for(int i=0; i<GlobalV::KPAR; i++)
 	{
-		nproc_in_pool[i] = GlobalV::NPROC/GlobalV::KPAR;
+		nproc_in_pool[i] = nprocgroup/GlobalV::KPAR;
 		if(i<remain_pro) this->nproc_in_pool[i]++;
-	}	
+	}
 
 	this->numz = new int*[GlobalV::KPAR];
 	this->startz = new int*[GlobalV::KPAR];
@@ -184,6 +188,11 @@ void Parallel_Grid::z_distribution(void)
 #ifdef __MPI
 void Parallel_Grid::zpiece_to_all(double *zpiece, const int &iz, double *rho)
 {
+	if(GlobalV::CALCULATION.substr(0,3)=="sto")
+	{
+		this->zpiece_to_stogroup(zpiece,iz,rho);
+		return;
+	}
 	assert(allocate);	
 	//ModuleBase::TITLE("Parallel_Grid","zpiece_to_all");
 	MPI_Status ierror;
@@ -251,6 +260,73 @@ void Parallel_Grid::zpiece_to_all(double *zpiece, const int &iz, double *rho)
 #endif
 
 #ifdef __MPI
+void Parallel_Grid::zpiece_to_stogroup(double *zpiece, const int &iz, double *rho)
+{
+	assert(allocate);	
+	//TITLE("Parallel_Grid","zpiece_to_all");
+	MPI_Status ierror;
+
+	const int znow = iz - this->startz[GlobalV::MY_POOL][GlobalV::RANK_IN_POOL];
+	const int proc = this->whichpro[GlobalV::MY_POOL][iz];
+	
+	if(GlobalV::MY_POOL==0)
+	{
+		// case 1: the first part of rho in processor 0.
+		// and send zpeice to to other pools.
+		if(proc == 0 && GlobalV::RANK_IN_STOGROUP ==0)
+		{
+			for(int ir=0; ir<ncxy; ir++)
+			{
+				rho[ir*nczp+znow] = zpiece[ir];
+			}
+			for(int ipool=1; ipool < GlobalV::KPAR; ipool++)
+			{
+				MPI_Send(zpiece, ncxy, MPI_DOUBLE, this->whichpro[ipool][iz], iz, STO_WORLD);
+			}
+		}
+
+		// case 2: processor n (n!=0) receive rho from processor 0.
+		// and the receive tag is iz.
+		else if(proc == GlobalV::RANK_IN_POOL )
+		{
+			MPI_Recv(zpiece, ncxy, MPI_DOUBLE, 0, iz, STO_WORLD,&ierror);
+			for(int ir=0; ir<ncxy; ir++)
+			{
+				rho[ir*nczp + znow] = zpiece[ir];
+			}
+		}
+				
+		// case 2: > first part rho: processor 0 send the rho
+		// to all pools. The tag is iz, because processor may
+		// send more than once, and the only tag to distinguish
+		// them is iz.
+		else if(GlobalV::RANK_IN_POOL==0)
+		{
+			for(int ipool=0; ipool < GlobalV::KPAR; ipool++)
+			{
+				MPI_Send(zpiece, ncxy, MPI_DOUBLE, this->whichpro[ipool][iz], iz, STO_WORLD);
+			}
+		}
+	}// MY_POOL == 0
+	else
+	{
+		//ofs_running << "\n Receive charge density iz=" << iz << endl;
+		// the processors in other pools always receive rho from
+		// processor 0. the tag is 'iz'
+		if(proc == GlobalV::RANK_IN_STOGROUP )
+		{
+			MPI_Recv(zpiece, ncxy, MPI_DOUBLE, 0, iz, STO_WORLD,&ierror);
+			for(int ir=0; ir<ncxy; ir++)
+			{
+				rho[ir*nczp+znow] = zpiece[ir];
+			}
+		}
+	}
+
+	//ofs_running << "\n iz = " << iz << " Done.";
+	return;	
+
+}
 void Parallel_Grid::reduce_to_fullrho(double *rhotot, double *rhoin)
 {
 	//ModuleBase::TITLE("Parallel_Grid","reduce_to_fullrho");
diff --git a/source/src_parallel/parallel_grid.h b/source/src_parallel/parallel_grid.h
@@ -21,6 +21,7 @@ class Parallel_Grid
 
 #ifdef __MPI	
 	void zpiece_to_all(double *zpiece, const int &iz, double *rho);
+	void zpiece_to_stogroup(double *zpiece, const int &iz, double *rho); //qainrui add for sto-dft 2021-7-21
 	
 	void reduce_to_fullrho(double *rhotot, double *rhoin);
 #endif
diff --git a/source/src_parallel/parallel_kpoints.cpp b/source/src_parallel/parallel_kpoints.cpp
diff --git a/source/src_parallel/parallel_stochi.cpp b/source/src_parallel/parallel_stochi.cpp
diff --git a/source/src_parallel/parallel_stochi.h b/source/src_parallel/parallel_stochi.h
diff --git a/source/src_pw/charge.cpp b/source/src_pw/charge.cpp
diff --git a/source/src_pw/stress_func_kin.cpp b/source/src_pw/stress_func_kin.cpp
diff --git a/source/src_pw/stress_func_nl.cpp b/source/src_pw/stress_func_nl.cpp
diff --git a/source/src_pw/wavefunc.cpp b/source/src_pw/wavefunc.cpp

Original file line number	Diff line number	Diff line change
`@@ -140,7 +140,6 @@ void Input::Default(void)`
`140`	`140`	`pw_seed = 1;`
`141`	`141`	`nche_sto = 0;`
`142`	`142`	`seed_sto = 0;`
`143`		`- stotype = "pw";`
`144`	`143`	`kpar = 1;`
`145`	`144`	`berry_phase = false;`
`146`	`145`	`gdir = 3;`
`@@ -552,9 +551,9 @@ bool Input::Read(const std::string &fn)`
`552`	`551`	`{`
`553`	`552`	`read_value(ifs, emin_sto);`
`554`	`553`	`}`
`555`		`- else if (strcmp("stotype", word) == 0)`
	`554`	`+ else if (strcmp("nstogroup", word) == 0)`
`556`	`555`	`{`
`557`		`- read_value(ifs, stotype);`
	`556`	`+ read_value(ifs, nstogroup);`
`558`	`557`	`}`
`559`	`558`	`else if (strcmp("kpar", word) == 0) // number of pools`
`560`	`559`	`{`
`@@ -1798,6 +1797,7 @@ void Input::Default_2(void) // jiyy add 2019-08-04`
`1798`	`1797`	`vdw_radius = "95";`
`1799`	`1798`	`}`
`1800`	`1799`	`}`
	`1800`	`+ if(calculation.substr(0,3) != "sto") nstogroup = 1;`
`1801`	`1801`	`}`
`1802`	`1802`	`#ifdef __MPI`
`1803`	`1803`	`void Input::Bcast()`
`@@ -1828,7 +1828,7 @@ void Input::Bcast()`
`1828`	`1828`	`Parallel_Common::bcast_int(pw_seed);`
`1829`	`1829`	`Parallel_Common::bcast_double(emax_sto);`
`1830`	`1830`	`Parallel_Common::bcast_double(emin_sto);`
`1831`		`- Parallel_Common::bcast_string(stotype);`
	`1831`	`+ Parallel_Common::bcast_int(nstogroup);`
`1832`	`1832`	`Parallel_Common::bcast_int(kpar);`
`1833`	`1833`	`Parallel_Common::bcast_bool(berry_phase);`
`1834`	`1834`	`Parallel_Common::bcast_int(gdir);`
`@@ -2174,7 +2174,7 @@ void Input::Check(void)`
`2174`	`2174`	`*/`
`2175`	`2175`	`this->relax_nmax = 1;`
`2176`	`2176`	`}`
`2177`		`- else if (calculation == "scf-sto") // qianrui 2021-2-20`
	`2177`	`+ else if (calculation == "sto-scf") // qianrui 2021-2-20`
`2178`	`2178`	`{`
`2179`	`2179`	`if (mem_saver == 1)`
`2180`	`2180`	`{`
Original file line number	Diff line number	Diff line change
`@@ -151,7 +151,8 @@ void ESolver_KS_PW::Init(Input &inp, UnitCell_pseudo &ucell)`
`151`	`151`	`//================================`
`152`	`152`	`// Initial start wave functions`
`153`	`153`	`//================================`
`154`		`- if (GlobalV::NBANDS != 0 \|\| (GlobalV::CALCULATION != "scf-sto" && GlobalV::CALCULATION != "relax-sto" && GlobalV::CALCULATION != "md-sto")) //qianrui add`
	`154`	`+ if (GlobalV::NBANDS != 0 \|\| GlobalV::CALCULATION.substr(0,3) != "sto")`
	`155`	`+ // qianrui add temporarily. In the future, wfcinit() should be compatible with cases when NBANDS=0`
`155`	`156`	`{`
`156`	`157`	`GlobalC::wf.wfcinit();`
`157`	`158`	`}`
Original file line number	Diff line number	Diff line change
`@@ -205,6 +205,12 @@ void ESolver_SDFT_PW::hamilt2density(int istep, int iter, double ethr)`
`205`	`205`	`{`
`206`	`206`	`if(ModuleSymmetry::Symmetry::symm_flag) MPI_Barrier(MPI_COMM_WORLD);`
`207`	`207`	`}`
	`208`	`+`
	`209`	`+ if(GlobalV::MY_STOGROUP == 0)`
	`210`	`+ {`
	`211`	`+ GlobalC::en.deband = GlobalC::en.delta_e();`
	`212`	`+ }`
	`213`	`+`
`208`	`214`	`}`
`209`	`215`
`210`	`216`	`void ESolver_SDFT_PW:: c_bands_k(const int ik, double* h_diag, const int istep, const int iter)`