Perf: Davidson method and pw_gatherscatter.h

dyzheng · dyzheng · commit 0b05ace7e27a · 2022-07-23T16:26:30.000+08:00
diff --git a/source/module_hsolver/diago_david.cpp b/source/module_hsolver/diago_david.cpp
@@ -68,24 +68,27 @@ void DiagoDavid::diag_mock(hamilt::Hamilt* phm_in, psi::Psi<std::complex<double>
     for (int m = 0; m < nband; m++)
     {
         // psi_m = psi(m)
-        for (int ig = 0; ig < dim; ig++)
+        ModuleBase::GlobalFunc::COPYARRAY(&psi(m, 0), psi_m.data(), dim);
+        /*for (int ig = 0; ig < dim; ig++)
         {
             psi_m[ig] = psi(m, ig);
-        }
+        }*/
 
         phm_in->sPsi(psi_m.data(), spsi.data(),  (size_t)dim);
         this->SchmitOrth(dim, nband, m, basis, psi_m.data(), spsi.data());
         phm_in->sPsi(psi_m.data(), spsi.data(),  (size_t)dim);
 
         // basis(m) = psi_m, hp(m) = H |psi_m>, sp(m) = S |psi_m>
-        std::complex<double>* sp_p = &sp(m, 0);
+        ModuleBase::GlobalFunc::COPYARRAY(psi_m.data(), &basis(m, 0), dim);
+        ModuleBase::GlobalFunc::COPYARRAY(spsi.data(), &sp(m, 0), dim);
+        /*std::complex<double>* sp_p = &sp(m, 0);
         std::complex<double>* basis_p = &basis(m, 0);
         for (int ig = 0; ig < dim; ig++)
         {
             basis_p[ig] = psi_m[ig];
             //hp(m, ig) = hpsi[ig];
             sp_p[ig] = spsi[ig];
-        }
+        }*/
     }
     hp_info dav_hpsi_in(&basis, psi::Range(1, 0, 0, nband-1));
     auto hp_psi = std::get<0>(phm_in->ops->hPsi(dav_hpsi_in));
@@ -150,7 +153,22 @@ void DiagoDavid::diag_mock(hamilt::Hamilt* phm_in, psi::Psi<std::complex<double>
 
             // updata eigenvectors of Hamiltonian
             ModuleBase::GlobalFunc::ZEROS(psi.get_pointer(), psi.get_nbands() * psi.get_nbasis());
-            for (int m = 0; m < nband; m++)
+            char transa = 'N';
+            char transb = 'T';
+            zgemm_(&transa,
+                    &transb,
+                    &dim, // m: row of A,C
+                    &nband, // n: col of B,C
+                    &nbase, // k: col of A, row of B
+                    &ModuleBase::ONE, // alpha
+                    basis.get_pointer(), // A
+                    &basis.get_nbasis(), // LDA: if(N) max(1,m) if(T) max(1,k)
+                    vc.c, // B
+                    &nbase_x, // LDB: if(N) max(1,k) if(T) max(1,n)
+                    &ModuleBase::ZERO, // belta
+                    psi.get_pointer(), // C
+                    &psi.get_nbasis()); // LDC: if(N) max(1, m)
+            /*for (int m = 0; m < nband; m++)
             {
                 for (int j = 0; j < nbase; j++)
                 {
@@ -159,7 +177,7 @@ void DiagoDavid::diag_mock(hamilt::Hamilt* phm_in, psi::Psi<std::complex<double>
                         psi(m, ig) += vc(j, m) * basis(j, ig);
                     }
                 }
-            }
+            }*/
 
             if (!this->notconv || (dav_iter == DiagoIterAssist::PW_DIAG_NMAX))
             {
@@ -594,18 +612,31 @@ void DiagoDavid::SchmitOrth(const int &npw,
     std::complex<double> *lagrange = new std::complex<double>[m + 1];
     ModuleBase::GlobalFunc::ZEROS(lagrange, m + 1);
 
-    const int one = 1;
-    for (int j = 0; j < m; j++)
+    int inc = 1;
+    int mp = m;
+    char trans = 'C';
+    zgemv_(&trans,
+           &npw,
+           &mp,
+           &ModuleBase::ONE,
+           psi.get_pointer(),
+           &psi.get_nbasis(),
+           spsi,
+           &inc,
+           &ModuleBase::ZERO,
+           lagrange,
+           &inc);
+    /*for (int j = 0; j < m; j++)
     {
         const std::complex<double>* psi_p = &(psi(j, 0));
         zdotc_(&lagrange[j], &npw, psi_p, &one, spsi, &one);
-        /*for (int ig = 0; ig < npw; ig++)
+        for (int ig = 0; ig < npw; ig++)
         {
             lagrange[j] += conj(psi(j, ig)) * spsi[ig];
-        }*/
+        }
         //	lagrange[j] = Diago_CG::ddot( npw, psi, j, spsi );
-    }
-    zdotc_(&lagrange[m], &npw, psi_m, &one, spsi, &one);
+    }*/
+    zdotc_(&lagrange[m], &npw, psi_m, &inc, spsi, &inc);
     /*for (int ig = 0; ig < npw; ig++)
     {
         lagrange[m] += conj(psi_m[ig]) * spsi[ig];
@@ -623,7 +654,7 @@ void DiagoDavid::SchmitOrth(const int &npw,
     for (int j = 0; j < m; j++)
     {
         const std::complex<double> alpha = -1 * lagrange[j];
-        zaxpy_(&npw, &alpha, &psi(j,0), &one, psi_m, &one);
+        zaxpy_(&npw, &alpha, &psi(j,0), &inc, psi_m, &inc);
         /*for (int ig = 0; ig < npw; ig++)
         {
             psi_m[ig] -= lagrange[j] * psi(j, ig);
diff --git a/source/module_pw/pw_gatherscatter.h b/source/module_pw/pw_gatherscatter.h
@@ -16,28 +16,34 @@ void PW_Basis:: gatherp_scatters(std::complex<T> *in, std::complex<T> *out)
     
     if(this->poolnproc == 1) //In this case nst=nstot, nz = nplane, 
     {
+        std::complex<T> * outp, *inp;
         for(int is = 0 ; is < this->nst ; ++is)
         {
             int ixy = this->istot2ixy[is];
             //int ixy = (ixy / fftny)*ny + ixy % fftny;
+            outp = &out[is*nz];
+            inp = &in[ixy*nz];
             for(int iz = 0 ; iz < this->nz ; ++iz)
             {
-                out[is*nz+iz] = in[ixy*nz+iz];
+                outp[iz] = inp[iz];
             }
         }
         ModuleBase::timer::tick(this->classname, "gatherp_scatters");
         return;
     }
 #ifdef __MPI
     //change (nplane fftnxy) to (nplane,nstot)
-    // Hence, we can send them at one time.  
+    // Hence, we can send them at one time.
+    std::complex<T> * outp, *inp;  
 	for (int istot = 0;istot < nstot; ++istot)
 	{
 		int ixy = this->istot2ixy[istot];
         //int ixy = (ixy / fftny)*ny + ixy % fftny;
+        outp = &out[istot*nplane];
+        inp = &in[ixy*nplane];
 		for (int iz = 0; iz < nplane; ++iz)
 		{
-			out[istot*nplane+iz] = in[ixy*nplane+iz];
+			outp[iz] = inp[iz];
 		}
 	}
 
@@ -48,14 +54,19 @@ void PW_Basis:: gatherp_scatters(std::complex<T> *in, std::complex<T> *out)
     else if(typeid(T) == typeid(float))
         MPI_Alltoallv(out, numr, startr, MPI_COMPLEX, in, numg, startg, MPI_COMPLEX, this->pool_world);
     // change (nz,ns) to (numz[ip],ns, poolnproc)
+    std::complex<T> * outp0, *inp0;
     for (int ip = 0; ip < this->poolnproc ;++ip)
 	{
         int nzip = this->numz[ip];
+        outp0 = &out[startz[ip]];
+        inp0 = &in[startg[ip]];
 		for (int is = 0; is < this->nst; ++is)
 		{
+            outp = &outp0[is * nz];
+            inp = &inp0[is * nzip ];
 			for (int izip = 0; izip < nzip; ++izip)
 			{
-				out[ is * nz + startz[ip] + izip] = in[startg[ip] + is*nzip + izip];
+				outp[izip] = inp[izip];
 			}
 		}
 	}
@@ -78,29 +89,38 @@ void PW_Basis:: gathers_scatterp(std::complex<T> *in, std::complex<T> *out)
     if(this->poolnproc == 1) //In this case nrxx=fftnx*fftny*nz, nst = nstot, 
     {
         ModuleBase::GlobalFunc::ZEROS(out, this->nrxx);
+        std::complex<T> * outp, *inp;
         for(int is = 0 ; is < this->nst ; ++is)
         {
             int ixy = istot2ixy[is];
             //int ixy = (ixy / fftny)*ny + ixy % fftny;
+            outp = &out[ixy*nz];
+            inp = &in[is*nz];
             for(int iz = 0 ; iz < this->nz ; ++iz)
             {
-                out[ixy*nz+iz] = in[is*nz+iz];
+                outp[iz] = inp[iz];
             }
         }
         ModuleBase::timer::tick(this->classname, "gathers_scatterp");
         return;
     }
 #ifdef __MPI
     // change (nz,ns) to (numz[ip],ns, poolnproc)
-    // Hence, we can send them at one time.  
+    // Hence, we can send them at one time. 
+    std::complex<T> * outp, *inp; 
+    std::complex<T> * outp0, *inp0;
     for (int ip = 0; ip < this->poolnproc ;++ip)
 	{
         int nzip = this->numz[ip];
+        outp0 = &out[startg[ip]];
+        inp0 = &in[startz[ip]];
 		for (int is = 0; is < this->nst; ++is)
 		{
+            outp = &outp0[is * nzip];
+            inp = &inp0[is * nz ];
 			for (int izip = 0; izip < nzip; ++izip)
 			{
-				out[startg[ip] + is*nzip + izip] = in[ is * nz + startz[ip] + izip];
+				outp[izip] = inp[izip];
 			}
 		}
 	}
@@ -117,9 +137,11 @@ void PW_Basis:: gathers_scatterp(std::complex<T> *in, std::complex<T> *out)
 	{
 		int ixy = this->istot2ixy[istot];
         //int ixy = (ixy / fftny)*ny + ixy % fftny;
+        outp = &out[ixy * nplane];
+        inp = &in[istot * nplane];
 		for (int iz = 0; iz < nplane; ++iz)
 		{
-			out[ixy*nplane+iz] = in[istot*nplane+iz];
+			outp[iz] = inp[iz];
 		}
 	}
 

Original file line number	Diff line number	Diff line change
`@@ -16,28 +16,34 @@ void PW_Basis:: gatherp_scatters(std::complex<T> in, std::complex<T> out)`
`16`	`16`
`17`	`17`	`if(this->poolnproc == 1) //In this case nst=nstot, nz = nplane,`
`18`	`18`	`{`
	`19`	`+ std::complex<T> * outp, *inp;`
`19`	`20`	`for(int is = 0 ; is < this->nst ; ++is)`
`20`	`21`	`{`
`21`	`22`	`int ixy = this->istot2ixy[is];`
`22`	`23`	`//int ixy = (ixy / fftny)*ny + ixy % fftny;`
	`24`	`+ outp = &out[is*nz];`
	`25`	`+ inp = &in[ixy*nz];`
`23`	`26`	`for(int iz = 0 ; iz < this->nz ; ++iz)`
`24`	`27`	`{`
`25`		`- out[isnz+iz] = in[ixynz+iz];`
	`28`	`+ outp[iz] = inp[iz];`
`26`	`29`	`}`
`27`	`30`	`}`
`28`	`31`	`ModuleBase::timer::tick(this->classname, "gatherp_scatters");`
`29`	`32`	`return;`
`30`	`33`	`}`
`31`	`34`	`#ifdef __MPI`
`32`	`35`	`//change (nplane fftnxy) to (nplane,nstot)`
`33`		`- // Hence, we can send them at one time.`
	`36`	`+ // Hence, we can send them at one time.`
	`37`	`+ std::complex<T> * outp, *inp;`
`34`	`38`	`for (int istot = 0;istot < nstot; ++istot)`
`35`	`39`	`{`
`36`	`40`	`int ixy = this->istot2ixy[istot];`
`37`	`41`	`//int ixy = (ixy / fftny)*ny + ixy % fftny;`
	`42`	`+ outp = &out[istot*nplane];`
	`43`	`+ inp = &in[ixy*nplane];`
`38`	`44`	`for (int iz = 0; iz < nplane; ++iz)`
`39`	`45`	`{`
`40`		`- out[istotnplane+iz] = in[ixynplane+iz];`
	`46`	`+ outp[iz] = inp[iz];`
`41`	`47`	`}`
`42`	`48`	`}`
`43`	`49`
`@@ -48,14 +54,19 @@ void PW_Basis:: gatherp_scatters(std::complex<T> in, std::complex<T> out)`
`48`	`54`	`else if(typeid(T) == typeid(float))`
`49`	`55`	`MPI_Alltoallv(out, numr, startr, MPI_COMPLEX, in, numg, startg, MPI_COMPLEX, this->pool_world);`
`50`	`56`	`// change (nz,ns) to (numz[ip],ns, poolnproc)`
	`57`	`+ std::complex<T> * outp0, *inp0;`
`51`	`58`	`for (int ip = 0; ip < this->poolnproc ;++ip)`
`52`	`59`	`{`
`53`	`60`	`int nzip = this->numz[ip];`
	`61`	`+ outp0 = &out[startz[ip]];`
	`62`	`+ inp0 = &in[startg[ip]];`
`54`	`63`	`for (int is = 0; is < this->nst; ++is)`
`55`	`64`	`{`
	`65`	`+ outp = &outp0[is * nz];`
	`66`	`+ inp = &inp0[is * nzip ];`
`56`	`67`	`for (int izip = 0; izip < nzip; ++izip)`
`57`	`68`	`{`
`58`		`- out[ is * nz + startz[ip] + izip] = in[startg[ip] + is*nzip + izip];`
	`69`	`+ outp[izip] = inp[izip];`
`59`	`70`	`}`
`60`	`71`	`}`
`61`	`72`	`}`
`@@ -78,29 +89,38 @@ void PW_Basis:: gathers_scatterp(std::complex<T> in, std::complex<T> out)`
`78`	`89`	`if(this->poolnproc == 1) //In this case nrxx=fftnxfftnynz, nst = nstot,`
`79`	`90`	`{`
`80`	`91`	`ModuleBase::GlobalFunc::ZEROS(out, this->nrxx);`
	`92`	`+ std::complex<T> * outp, *inp;`
`81`	`93`	`for(int is = 0 ; is < this->nst ; ++is)`
`82`	`94`	`{`
`83`	`95`	`int ixy = istot2ixy[is];`
`84`	`96`	`//int ixy = (ixy / fftny)*ny + ixy % fftny;`
	`97`	`+ outp = &out[ixy*nz];`
	`98`	`+ inp = &in[is*nz];`
`85`	`99`	`for(int iz = 0 ; iz < this->nz ; ++iz)`
`86`	`100`	`{`
`87`		`- out[ixynz+iz] = in[isnz+iz];`
	`101`	`+ outp[iz] = inp[iz];`
`88`	`102`	`}`
`89`	`103`	`}`
`90`	`104`	`ModuleBase::timer::tick(this->classname, "gathers_scatterp");`
`91`	`105`	`return;`
`92`	`106`	`}`
`93`	`107`	`#ifdef __MPI`
`94`	`108`	`// change (nz,ns) to (numz[ip],ns, poolnproc)`
`95`		`- // Hence, we can send them at one time.`
	`109`	`+ // Hence, we can send them at one time.`
	`110`	`+ std::complex<T> * outp, *inp;`
	`111`	`+ std::complex<T> * outp0, *inp0;`
`96`	`112`	`for (int ip = 0; ip < this->poolnproc ;++ip)`
`97`	`113`	`{`
`98`	`114`	`int nzip = this->numz[ip];`
	`115`	`+ outp0 = &out[startg[ip]];`
	`116`	`+ inp0 = &in[startz[ip]];`
`99`	`117`	`for (int is = 0; is < this->nst; ++is)`
`100`	`118`	`{`
	`119`	`+ outp = &outp0[is * nzip];`
	`120`	`+ inp = &inp0[is * nz ];`
`101`	`121`	`for (int izip = 0; izip < nzip; ++izip)`
`102`	`122`	`{`
`103`		`- out[startg[ip] + isnzip + izip] = in[ is nz + startz[ip] + izip];`
	`123`	`+ outp[izip] = inp[izip];`
`104`	`124`	`}`
`105`	`125`	`}`
`106`	`126`	`}`
`@@ -117,9 +137,11 @@ void PW_Basis:: gathers_scatterp(std::complex<T> in, std::complex<T> out)`
`117`	`137`	`{`
`118`	`138`	`int ixy = this->istot2ixy[istot];`
`119`	`139`	`//int ixy = (ixy / fftny)*ny + ixy % fftny;`
	`140`	`+ outp = &out[ixy * nplane];`
	`141`	`+ inp = &in[istot * nplane];`
`120`	`142`	`for (int iz = 0; iz < nplane; ++iz)`
`121`	`143`	`{`
`122`		`- out[ixynplane+iz] = in[istotnplane+iz];`
	`144`	`+ outp[iz] = inp[iz];`
`123`	`145`	`}`
`124`	`146`	`}`
`125`	`147`