use openmp to accelerate fvnl_dbeta_gamma.cpp (#4814)

source/module_hamilt_lcao/hamilt_lcaodft/fvnl_dbeta_gamma.cpp

@@ -20,33 +20,38 @@ void Force_LCAO<double>::cal_fvnl_dbeta(const elecstate::DensityMatrix<double, d
     ModuleBase::TITLE("Force_LCAO", "cal_fvnl_dbeta");
     ModuleBase::timer::tick("Force_LCAO", "cal_fvnl_dbeta");
 
-    double r0[3];
-    double r1[3];
-
+// use schedule(dynamic) for load balancing because adj_num is various
+#pragma omp parallel
+{ 
+    ModuleBase::matrix local_svnl_dbeta(3, 3);
+    #pragma omp for schedule(dynamic)
     for (int iat = 0; iat < ucell.nat; iat++)
     {
         const int it = ucell.iat2it[iat];
         const int ia = ucell.iat2ia[iat];
         const int T0 = it;
         const int I0 = ia;
+        double r0[3]{0};
+        double r1[3]{0};
 
         const ModuleBase::Vector3<double> tau0 = ucell.atoms[it].tau[ia];
         // find ajacent atom of atom ia
         // gd.Find_atom( ucell.atoms[it].tau[ia] );
-        gd.Find_atom(ucell, tau0, it, ia);
+        AdjacentAtomInfo adjs;
+        gd.Find_atom(ucell, tau0, it, ia, &adjs);
         const double Rcut_Beta = ucell.infoNL.Beta[it].get_rcut_max();
 
         std::vector<std::unordered_map<int, std::vector<std::vector<double>>>> nlm_tot;
-        nlm_tot.resize(gd.getAdjacentNum() + 1); // this saves <psi_i|beta> and <psi_i|\nabla|beta>
+        nlm_tot.resize(adjs.adj_num + 1); // this saves <psi_i|beta> and <psi_i|\nabla|beta>
 
         // Step 1 : Calculate and save <psi_i|beta> and <psi_i|\nabla|beta>
-        for (int ad1 = 0; ad1 < gd.getAdjacentNum() + 1; ad1++)
+        for (int ad1 = 0; ad1 < adjs.adj_num + 1; ad1++)
         {
-            const int T1 = gd.getType(ad1);
+            const int T1 = adjs.ntype[ad1];
             const Atom* atom1 = &ucell.atoms[T1];
-            const int I1 = gd.getNatom(ad1);
+            const int I1 = adjs.natom[ad1];
             const int start1 = ucell.itiaiw2iwt(T1, I1, 0);
-            const ModuleBase::Vector3<double> tau1 = gd.getAdjacentTau(ad1);
+            const ModuleBase::Vector3<double> tau1 = adjs.adjacent_tau[ad1];
             const double Rcut_AO1 = orb.Phi[T1].getRcut();
 
             nlm_tot[ad1].clear();
@@ -87,22 +92,22 @@ void Force_LCAO<double>::cal_fvnl_dbeta(const elecstate::DensityMatrix<double, d
         } // ad
 
         // Step 2 : sum to get beta<psi_i|beta><beta|\nabla|psi_j>
-        for (int ad1 = 0; ad1 < gd.getAdjacentNum() + 1; ++ad1)
+        for (int ad1 = 0; ad1 < adjs.adj_num + 1; ++ad1)
         {
-            const int T1 = gd.getType(ad1);
+            const int T1 = adjs.ntype[ad1];
             const Atom* atom1 = &ucell.atoms[T1];
-            const int I1 = gd.getNatom(ad1);
+            const int I1 = adjs.natom[ad1];
             const int start1 = ucell.itiaiw2iwt(T1, I1, 0);
-            const ModuleBase::Vector3<double> tau1 = gd.getAdjacentTau(ad1);
+            const ModuleBase::Vector3<double> tau1 = adjs.adjacent_tau[ad1];
             const double Rcut_AO1 = orb.Phi[T1].getRcut();
 
-            for (int ad2 = 0; ad2 < gd.getAdjacentNum() + 1; ad2++)
+            for (int ad2 = 0; ad2 < adjs.adj_num + 1; ad2++)
             {
-                const int T2 = gd.getType(ad2);
+                const int T2 = adjs.ntype[ad2];
                 const Atom* atom2 = &ucell.atoms[T2];
-                const int I2 = gd.getNatom(ad2);
+                const int I2 = adjs.natom[ad2];
                 const int start2 = ucell.itiaiw2iwt(T2, I2, 0);
-                const ModuleBase::Vector3<double> tau2 = gd.getAdjacentTau(ad2);
+                const ModuleBase::Vector3<double> tau2 = adjs.adjacent_tau[ad2];
                 const double Rcut_AO2 = orb.Phi[T2].getRcut();
 
                 const double dist1 = (tau1 - tau0).norm() * ucell.lat0;
@@ -207,8 +212,8 @@ void Force_LCAO<double>::cal_fvnl_dbeta(const elecstate::DensityMatrix<double, d
                             {
                                 for (int jpol = ipol; jpol < 3; jpol++)
                                 {
-                                    svnl_dbeta(ipol, jpol)
-                                        += sum / 2.0 * (nlm[ipol] * r0[jpol] + nlm_t[ipol] * r1[jpol]);
+                                    local_svnl_dbeta(ipol, jpol)
+                                    += sum / 2.0 * (nlm[ipol] * r0[jpol] + nlm_t[ipol] * r1[jpol]);
                                 }
                             }
                         }
@@ -218,6 +223,21 @@ void Force_LCAO<double>::cal_fvnl_dbeta(const elecstate::DensityMatrix<double, d
         }             // ad1
     }                 // iat
 
+    // sum up local_svnl_dbeta to svnl_dbeta
+    if (isstress)
+    {
+        #pragma omp critical
+        {
+            for (int ipol = 0; ipol < 3; ipol++)
+            {
+                for (int jpol = ipol; jpol < 3; jpol++)
+                {
+                    svnl_dbeta(ipol, jpol) += local_svnl_dbeta(ipol, jpol);
+                }
+            }
+        }
+    }
+}
     if (isstress)
     {
         StressTools::stress_fill(ucell.lat0, ucell.omega, svnl_dbeta);