Merge pull request #540 from Qianruipku/planewave

< prof >

Author: Qianruipku

source/module_base/global_function.h

@@ -116,7 +116,7 @@ void DONE(std::ofstream &ofs,const std::string &description, bool only_rank0 = f
 template<class T, class TI>
 inline void ZEROS(std::complex<T> *u,const TI n)		// Peize Lin change int to TI at 2020.03.03
 {
-    if(n <= 0) return;
+    assert(n>=0);
     for (TI i=0;i<n;i++)
     {
         u[i] = std::complex<T>(0.0,0.0);

source/module_pw/fft.cpp

@@ -37,7 +37,7 @@ FFT::~FFT()
 #endif
 }
 
-void FFT:: initfft(int nx_in, int bigny_in, int nz_in, int ns_in, int nplane_in, bool gamma_only_in, bool mpifft_in)
+void FFT:: initfft(int nx_in, int bigny_in, int nz_in, int ns_in, int nplane_in, int nproc_in, bool gamma_only_in, bool mpifft_in)
 {
 	this->gamma_only = gamma_only_in;
 	this->nx = nx_in;
@@ -47,14 +47,16 @@ void FFT:: initfft(int nx_in, int bigny_in, int nz_in, int ns_in, int nplane_in,
 	this->nz = nz_in;
 	this->ns = ns_in;
 	this->nplane = nplane_in;
+	this->nproc = nproc_in;
 	this->mpifft = mpifft_in;
 	this->nxy = this->nx * this-> ny;
 	this->bignxy = this->nx * this->bigny;
 	this->maxgrids = (this->nz * this->ns > this->bignxy * nplane) ? this->nz * this->ns : this->bignxy * nplane;
 	if(!this->mpifft)
 	{
 		//It seems in-place fft is faster than out-of-place fft
-		c_gspace  = (std::complex<double> *) fftw_malloc(sizeof(fftw_complex) * this->nz * this->ns);
+		if(this->nproc == 1) 	c_gspace  = (std::complex<double> *) fftw_malloc(sizeof(fftw_complex) * this->nz * this->ns);
+		else					c_gspace  = (std::complex<double> *) fftw_malloc(sizeof(fftw_complex) * maxgrids);
 		//c_gspace2  = (std::complex<double> *) fftw_malloc(sizeof(fftw_complex) * this->nz * this->ns);
 		if(this->gamma_only)
 		{
@@ -69,8 +71,8 @@ void FFT:: initfft(int nx_in, int bigny_in, int nz_in, int ns_in, int nplane_in,
 		}
 		else
 		{
-			c_rspace  = (std::complex<double> *) fftw_malloc(sizeof(fftw_complex) * this->bignxy * nplane);
-			//c_rspace2  = (std::complex<double> *) fftw_malloc(sizeof(fftw_complex) * this->bignxy * nplane);
+			if(this->nproc == 1) c_rspace  = (std::complex<double> *) fftw_malloc(sizeof(fftw_complex) * this->bignxy * nplane);
+			else 				 c_rspace  = (std::complex<double> *) fftw_malloc(sizeof(fftw_complex) * maxgrids);
 		}
 #ifdef __MIX_PRECISION
 		cf_gspace  = (std::complex<float> *)fftw_malloc(sizeof(fftwf_complex) * this->nz * this->ns);

source/module_pw/fft.h

@@ -27,7 +27,7 @@ class FFT
 
 	FFT();
 	~FFT();
-	void initfft(int nx_in, int bigny_in, int nz_in, int ns_in, int nplane_in, bool gamma_only_in, bool mpifft_in = false);
+	void initfft(int nx_in, int bigny_in, int nz_in, int ns_in, int nplane_in, int nproc_in, bool gamma_only_in, bool mpifft_in = false);
 	void setupFFT();
 	void cleanFFT();
 
@@ -52,6 +52,7 @@ class FFT
 	int ns; //number of sticks
 	int nplane; //number of x-y planes
 	int maxgrids; // max between nz * ns and bignxy * nplane
+	int nproc; // number of proc.
 	std::complex<double> * c_gspace, *c_gspace2; //complex number space for g, [ns * nz]
 	std::complex<double> * c_rspace, *c_rspace2;//complex number space for r, [nplane * nx *ny]
 	double *r_rspace; //real number space for r, [nplane * nx *ny]

source/module_pw/pw_basis.cpp

@@ -14,11 +14,16 @@ PW_Basis::PW_Basis()
     ixy2ip = NULL; 
     startnsz_per = NULL;
     nstnz_per = NULL;
+    nst_per = NULL;
     gdirect = NULL;		
     gcar = NULL; 
     gg = NULL;
     startz = NULL;
     numz = NULL;  
+    this->numg = NULL;
+	this->startg = NULL;
+	this->startr = NULL;
+	this->numr = NULL;
     poolnproc = 1;
     poolrank = 0;
 }
@@ -32,21 +37,62 @@ PW_Basis:: ~PW_Basis()
     if(ixy2ip != NULL) delete[] ixy2ip;
     if(startnsz_per != NULL) delete[] startnsz_per;
     if(nstnz_per != NULL) delete[] nstnz_per;
+    if(nst_per != NULL) delete[] nst_per;
     if(gdirect != NULL) delete[] gdirect;
     if(gcar != NULL) delete[] gcar;
     if(gg != NULL) delete[] gg;
     if(startz != NULL) delete[] startz;
     if(numz != NULL) delete[] numz;
+    if(numg != NULL) delete[] numg;
+    if(numr != NULL) delete[] numr;
+    if(startg != NULL) delete[] startg;
+    if(startr != NULL) delete[] startr;
 }
 
 void PW_Basis::setuptransform()
 {
     this->distribute_r();
     this->distribute_g();
-    this->ft.initfft(this->nx,this->bigny,this->nz,this->nst,this->nplane,this->gamma_only);
+    this->getstartgr();
+    this->ft.initfft(this->nx,this->bigny,this->nz,this->nst,this->nplane,this->poolnproc,this->gamma_only);
     this->ft.setupFFT();
 }
 
+void PW_Basis::getstartgr()
+{
+    this->maxgrids = (this->nz * this->nst > this->bignxy * nplane) ? this->nz * this->nst : this->bignxy * nplane;
+    
+    //---------------------------------------------
+	// sum : starting plane of FFT box.
+	//---------------------------------------------
+    this->numg = new int[poolnproc];
+	this->startg = new int[poolnproc];
+	this->startr = new int[poolnproc];
+	this->numr = new int[poolnproc];
+
+	// Each processor has a set of full sticks,
+	// 'rank_use' processor send a piece(npps[ip]) of these sticks(nst_per[rank_use])
+	// to all the other processors in this pool
+	for (int ip = 0;ip < poolnproc; ++ip) this->numg[ip] = this->nst_per[poolrank] * this->numz[ip];
+
+
+	// Each processor in a pool send a piece of each stick(nst_per[ip]) to
+	// other processors in this pool
+	// rank_use processor receive datas in npps[rank_p] planes.
+	for (int ip = 0;ip < poolnproc; ++ip) this->numr[ip] = this->nst_per[ip] * this->numz[poolrank];
+
+
+	// startg record the starting 'numg' position in each processor.
+	this->startg[0] = 0;
+	for (int ip = 1;ip < poolnproc; ++ip) this->startg[ip] = this->startg[ip-1] + this->numg[ip-1];
+
+
+	// startr record the starting 'numr' position
+	this->startr[0] = 0;
+	for (int ip = 1;ip < poolnproc; ++ip) this->startr[ip] = this->startr[ip-1] + this->numr[ip-1];
+    return;
+}
+
 //
 // Collect planewaves on current core, and construct gg, gdirect, gcar according to ig2isz and is2ixy.
 // is2ixy contains the x-coordinate and y-coordinate of sticks on current core.

source/module_pw/pw_basis.h

@@ -55,6 +55,7 @@ class PW_Basis
     // only on first proc.
     int *startnsz_per; // startnsz_per[ip]: starting is * nz stick in the ip^th proc.
     int *nstnz_per; // nz * nst(number of sticks) on each core.
+    int *nst_per;// nst on each core
     // on all proc.
     int *ig2isz; // map ig to (is, iz).
     int *istot2ixy; // istot2ixy[is]: ix + iy * nx of is^th stick among all sticks.
@@ -69,6 +70,10 @@ class PW_Basis
     int nrxx; //num. of real space grids
     int *startz; //startz[ip]: starting z plane in the ip-th proc. in current POOL_WORLD 
 	int *numz; //numz[ip]: num. of z planes in the ip-th proc. in current POOL_WORLD
+    int *numg; //numg[ip] :  nst_per[poolrank] * numz[ip] 
+    int *numr; //numr[ip] :  numz[poolrank] * nst_per[ip]
+    int *startg;  // startg[ip] = numg[ip-1] + startg[ip-1]
+    int *startr;  // startr[ip] = numr[ip-1] + startr[ip-1]
     int nplane; //num. of planes in current proc.
 
     ModuleBase::Vector3<double> *gdirect;		//(= *G1d) ; // ig = new Vector igc[ngmc]
@@ -86,6 +91,8 @@ class PW_Basis
     //distribute real-space grids to different processors
     void distribute_r();
 
+    void getstartgr();
+
     //distribute plane waves to different processors
     void distribution_method1(); // x varies fast
     void distribution_method2(); // sticks sorted according to ixy
@@ -135,8 +142,7 @@ class PW_Basis
         int* st_i,          // x or x + nx (if x < 0) of stick.
         int* st_j,          // y or y + ny (if y < 0) of stick.
         int* st_length,     // the stick on (x, y) consists of st_length[x*ny+y] planewaves.
-        int* npw_per,       // number of planewaves on each core.
-        int* nst_per       // number of sticks on each core.
+        int* npw_per       // number of planewaves on each core.
     );
     void get_istot2ixy(
         int* st_i,          // x or x + nx (if x < 0) of stick.
@@ -180,6 +186,7 @@ class PW_Basis
 	// FFT dimensions for wave functions.
 	int nx, ny, nz, nxyz, nxy;
     int bigny, bignxyz, bignxy; // Gamma_only: ny = int(bigny/2)-1 , others: ny = bigny
+    int maxgrids; // max between nz * ns and bignxy * nplane
     FFT ft;
 
     void real2recip(double * in, std::complex<double> * out); //in:(nplane,nx*ny)  ; out(nz, ns)
@@ -189,6 +196,8 @@ class PW_Basis
 
     void gatherp_scatters(std::complex<double> *in, std::complex<double> *out); //gather planes and scatter sticks of all processors
     void gathers_scatterp(std::complex<double> *in, std::complex<double> *out); //gather sticks of and scatter planes of all processors
+    void gathers_scatterp2(std::complex<double> *in, std::complex<double> *out); //gather sticks of and scatter planes of all processors
+    void gatherp_scatters2(std::complex<double> *in, std::complex<double> *out); //gather sticks of and scatter planes of all processors
     void gatherp_scatters_gamma(std::complex<double> *in, std::complex<double> *out); //gather planes and scatter sticks of all processors, used when gamma_only
     void gathers_scatterp_gamma(std::complex<double> *in, std::complex<double> *out); //gather sticks of and scatter planes of all processors, used when gamma only
 

source/module_pw/pw_distributeg_method1.cpp

@@ -30,7 +30,7 @@ void PW_Basis::distribution_method1()
     this->nstot = 0;                     // total number of sticks.
     int *st_bottom2D = NULL;             // st_bottom2D[ixy], minimum z of stick on (x, y).
     int *st_length2D = NULL;             // st_length2D[ixy], number of planewaves in stick on (x, y).
-
+    this->nst_per = new int[this->poolnproc]; // number of sticks on each core.
     if (poolrank == 0)
     {
         // (1) Count the total number of planewaves (tot_npw) and sticks (this->nstot).                  
@@ -73,11 +73,10 @@ void PW_Basis::distribution_method1()
 
         // (3) Distribute sticks to cores.
         int *npw_per = new int[this->poolnproc];  // number of planewaves on each core.
-        int *nst_per = new int[this->poolnproc]; // number of sticks on each core.
         this->nstnz_per = new int[this->poolnproc]; // nz * nst(number of sticks) on each core.
         this->startnsz_per = new int[this->poolnproc];
         ModuleBase::GlobalFunc::ZEROS(npw_per, poolnproc);
-        ModuleBase::GlobalFunc::ZEROS(nst_per, poolnproc);
+        ModuleBase::GlobalFunc::ZEROS(this->nst_per, poolnproc);
         ModuleBase::GlobalFunc::ZEROS(this->nstnz_per, poolnproc);
         ModuleBase::GlobalFunc::ZEROS(startnsz_per, poolnproc);
 
@@ -86,7 +85,7 @@ void PW_Basis::distribution_method1()
         {
             this->ixy2ip[ixy] = -1;                 // meaning this stick has not been distributed or there is no stick on (x, y).
         }
-        this->divide_sticks(st_i, st_j, st_length, npw_per, nst_per);
+        this->divide_sticks(st_i, st_j, st_length, npw_per);
         delete[] st_length;
 
          // for test -----------------------------------------------------------------------------
@@ -113,7 +112,6 @@ void PW_Basis::distribution_method1()
             MPI_Send(&nst_per[ip], 1, MPI_INT, ip, 0, POOL_WORLD);
         }
         delete[] npw_per;
-        delete[] nst_per;
 
 #else
         // Serial line
@@ -172,6 +170,7 @@ void PW_Basis::distribution_method1()
     MPI_Bcast(this->ixy2ip, this->nxy, MPI_INT, 0, POOL_WORLD);
     MPI_Bcast(this->istot2ixy, this->nstot, MPI_INT, 0, POOL_WORLD);
     MPI_Bcast(this->ixy2istot, this->nxy, MPI_INT, 0, POOL_WORLD);
+    MPI_Bcast(this->nst_per, this->poolnproc, MPI_INT, 0 , POOL_WORLD);
 
     std::cout << "Bcast done\n";
 #endif
@@ -304,8 +303,7 @@ void PW_Basis::divide_sticks(
     int* st_i,          // x or x + nx (if x < 0) of stick.
     int* st_j,          // y or y + ny (if y < 0) of stick.
     int* st_length,     // the stick on (x, y) consists of st_length[x*ny+y] planewaves.
-    int* npw_per,       // number of planewaves on each core.
-    int* nst_per       // number of sticks on each core.
+    int* npw_per       // number of planewaves on each core.
 )
 {
     int ipmin = 0; // The ip of core containing least number of planewaves.
@@ -322,25 +320,25 @@ void PW_Basis::divide_sticks(
 
             if (npw_ip == 0)
             {
-                if (non_zero_grid + nz < this->nrxx) // assert reciprocal planewaves is less than real space planewaves.
-                {
+                // if (non_zero_grid + nz < this->nrxx) // assert reciprocal planewaves is less than real space planewaves.
+                // {
                     ipmin = ip;
                     break;
-                }
+                // }
             }
             else if (npw_ip < npwmin)
             {
-                if (non_zero_grid + nz < this->nrxx) // assert reciprocal planewaves is less than real space planewaves.
-                {
+                // if (non_zero_grid + nz < this->nrxx) // assert reciprocal planewaves is less than real space planewaves.
+                // {
                     ipmin = ip;
-                }
+                // }
             }
             else if (npw_ip == npwmin && nst_ip < nstmin)
             {
-                if (non_zero_grid + nz < this->nrxx)
-                {
+                // if (non_zero_grid + nz < this->nrxx)
+                // {
                     ipmin = ip;
-                }
+                // }
             }
         }
         nst_per[ipmin]++;