change the position of omp (#5327)

Qianruipku · mohanchen · web-flow · commit f0a0c189e22c · 2024-10-24T01:41:18.000+08:00
Co-authored-by: Mohan Chen &lt;mohanchen@pku.edu.cn&gt;
diff --git a/source/module_hamilt_pw/hamilt_pwdft/kernels/ekinetic_op.cpp b/source/module_hamilt_pw/hamilt_pwdft/kernels/ekinetic_op.cpp
@@ -17,39 +17,34 @@ struct ekinetic_pw_op<FPTYPE, base_device::DEVICE_CPU>
     {
         if (is_first_node)
         {
-            for (int ib = 0; ib < nband; ++ib)
-            {
 #ifdef _OPENMP
 #pragma omp parallel for
 #endif
+            for (int ib = 0; ib < nband; ++ib)
+            {
+                const int ig0 = ib * max_npw;
                 for (int ig = 0; ig < npw; ++ig)
                 {
-                    tmhpsi[ig] = gk2_ik[ig] * tpiba2 * tmpsi_in[ig];
+                    tmhpsi[ig + ig0] = gk2_ik[ig] * tpiba2 * tmpsi_in[ig + ig0];
                 }
-#ifdef _OPENMP
-#pragma omp parallel for
-#endif
                 for (int ig = npw; ig < max_npw; ++ig)
                 {
-                    tmhpsi[ig] = 0.0;
+                    tmhpsi[ig + ig0] = 0.0;
                 }
-                tmpsi_in += max_npw;
-                tmhpsi += max_npw;
             }
         }
         else
         {
-            for (int ib = 0; ib < nband; ++ib)
-            {
 #ifdef _OPENMP
 #pragma omp parallel for
 #endif
+            for (int ib = 0; ib < nband; ++ib)
+            {
+                const int ig0 = ib * max_npw;
                 for (int ig = 0; ig < npw; ++ig)
                 {
-                    tmhpsi[ig] += gk2_ik[ig] * tpiba2 * tmpsi_in[ig];
+                    tmhpsi[ig + ig0] += gk2_ik[ig] * tpiba2 * tmpsi_in[ig + ig0];
                 }
-                tmpsi_in += max_npw;
-                tmhpsi += max_npw;
             }
         }
     }
diff --git a/source/module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.cpp b/source/module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.cpp
@@ -55,14 +55,17 @@ void HamiltSdftPW<T, Device>::hPsi_norm(const T* psi_in, T* hpsi_norm, const int
     const Real emax = *this->emax;
     const Real Ebar = (emin + emax) / 2;
     const Real DeltaE = (emax - emin) / 2;
+
+#ifdef _OPENMP
+#pragma omp parallel for
+#endif
     for (int ib = 0; ib < nbands; ++ib)
     {
+        const int ig0 = ib * npwk_max;
         for (int ig = 0; ig < npwk; ++ig)
         {
-            hpsi_norm[ig] = (hpsi_norm[ig] - Ebar * psi_in[ig]) / DeltaE;
+            hpsi_norm[ig + ig0] = (hpsi_norm[ig + ig0] - Ebar * psi_in[ig + ig0]) / DeltaE;
         }
-        hpsi_norm += npwk_max;
-        psi_in += npwk_max;
     }
     ModuleBase::timer::tick("HamiltSdftPW", "hPsi_norm");
 }

Original file line number	Diff line number	Diff line change
`@@ -17,39 +17,34 @@ struct ekinetic_pw_op<FPTYPE, base_device::DEVICE_CPU>`
`17`	`17`	`{`
`18`	`18`	`if (is_first_node)`
`19`	`19`	`{`
`20`		`- for (int ib = 0; ib < nband; ++ib)`
`21`		`- {`
`22`	`20`	`#ifdef _OPENMP`
`23`	`21`	`#pragma omp parallel for`
`24`	`22`	`#endif`
	`23`	`+ for (int ib = 0; ib < nband; ++ib)`
	`24`	`+ {`
	`25`	`+ const int ig0 = ib * max_npw;`
`25`	`26`	`for (int ig = 0; ig < npw; ++ig)`
`26`	`27`	`{`
`27`		`- tmhpsi[ig] = gk2_ik[ig] * tpiba2 * tmpsi_in[ig];`
	`28`	`+ tmhpsi[ig + ig0] = gk2_ik[ig] * tpiba2 * tmpsi_in[ig + ig0];`
`28`	`29`	`}`
`29`		`-#ifdef _OPENMP`
`30`		`-#pragma omp parallel for`
`31`		`-#endif`
`32`	`30`	`for (int ig = npw; ig < max_npw; ++ig)`
`33`	`31`	`{`
`34`		`- tmhpsi[ig] = 0.0;`
	`32`	`+ tmhpsi[ig + ig0] = 0.0;`
`35`	`33`	`}`
`36`		`- tmpsi_in += max_npw;`
`37`		`- tmhpsi += max_npw;`
`38`	`34`	`}`
`39`	`35`	`}`
`40`	`36`	`else`
`41`	`37`	`{`
`42`		`- for (int ib = 0; ib < nband; ++ib)`
`43`		`- {`
`44`	`38`	`#ifdef _OPENMP`
`45`	`39`	`#pragma omp parallel for`
`46`	`40`	`#endif`
	`41`	`+ for (int ib = 0; ib < nband; ++ib)`
	`42`	`+ {`
	`43`	`+ const int ig0 = ib * max_npw;`
`47`	`44`	`for (int ig = 0; ig < npw; ++ig)`
`48`	`45`	`{`
`49`		`- tmhpsi[ig] += gk2_ik[ig] * tpiba2 * tmpsi_in[ig];`
	`46`	`+ tmhpsi[ig + ig0] += gk2_ik[ig] * tpiba2 * tmpsi_in[ig + ig0];`
`50`	`47`	`}`
`51`		`- tmpsi_in += max_npw;`
`52`		`- tmhpsi += max_npw;`
`53`	`48`	`}`
`54`	`49`	`}`
`55`	`50`	`}`
Original file line number	Diff line number	Diff line change
`@@ -55,14 +55,17 @@ void HamiltSdftPW<T, Device>::hPsi_norm(const T* psi_in, T* hpsi_norm, const int`
`55`	`55`	`const Real emax = *this->emax;`
`56`	`56`	`const Real Ebar = (emin + emax) / 2;`
`57`	`57`	`const Real DeltaE = (emax - emin) / 2;`
	`58`	`+`
	`59`	`+#ifdef _OPENMP`
	`60`	`+#pragma omp parallel for`
	`61`	`+#endif`
`58`	`62`	`for (int ib = 0; ib < nbands; ++ib)`
`59`	`63`	`{`
	`64`	`+ const int ig0 = ib * npwk_max;`
`60`	`65`	`for (int ig = 0; ig < npwk; ++ig)`
`61`	`66`	`{`
`62`		`- hpsi_norm[ig] = (hpsi_norm[ig] - Ebar * psi_in[ig]) / DeltaE;`
	`67`	`+ hpsi_norm[ig + ig0] = (hpsi_norm[ig + ig0] - Ebar * psi_in[ig + ig0]) / DeltaE;`
`63`	`68`	`}`
`64`		`- hpsi_norm += npwk_max;`
`65`		`- psi_in += npwk_max;`
`66`	`69`	`}`
`67`	`70`	`ModuleBase::timer::tick("HamiltSdftPW", "hPsi_norm");`
`68`	`71`	`}`