Complete PPCG pool reductions

Author: cheerly-pku

source/source_hsolver/ppcg/diago_ppcg_cg.hpp

@@ -39,11 +39,9 @@ void DiagoPPCG<T, Device>::orth_gradient(
         for (int i = 0; i < n_band_; ++i)
         {
             // Full complex inner product <psi_i | grad_j>
-            T coeff = 0;
             const T* pi = psi + i * ld_psi_;
             const T* gj = grad.data() + j * ld_psi_;
-            for (int ig = 0; ig < n_dim_; ++ig)
-                coeff += std::conj(pi[ig]) * gj[ig];
+            const T coeff = complex_dot(pi, gj);
             if (std::abs(coeff) <= std::numeric_limits<Real>::epsilon())
                 continue;
             // grad_j -= S|psi_i> * coeff
@@ -102,6 +100,8 @@ void DiagoPPCG<T, Device>::update_polak_ribiere(
             beta_num_zr += static_cast<Real>(std::real(z * std::conj(g[ig])));
             beta_num_zo += static_cast<Real>(std::real(z * std::conj(r_old)));
         }
+        reduce_pool_if_mpi_ready(beta_num_zr);
+        reduce_pool_if_mpi_ready(beta_num_zo);
 
         Real beta = 0;
         const Real denom = beta_denom[j];

source/source_hsolver/ppcg/diago_ppcg_diag.hpp

@@ -288,6 +288,7 @@ double DiagoPPCG<T, Device>::diag(const HPsiFunc& hpsi_func,
                 for (int ig = 0; ig < n_dim_; ++ig)
                     nrm2 += static_cast<Real>(
                         std::norm(grad[idx(ig, i, ld_psi_)]));
+                reduce_pool_if_mpi_ready(nrm2);
                 if (std::sqrt(nrm2) > std::max(static_cast<Real>(ethr_band[i]),
                                                diag_thr_))
                 {

source/source_hsolver/ppcg/diago_ppcg_lapack.hpp

@@ -1,5 +1,7 @@
 #include <ATen/kernels/lapack.h>
 
+#include "source_base/parallel_reduce.h"
+
 #include <cstdlib>
 #include <fstream>
 
@@ -10,6 +12,32 @@ namespace hsolver {
 // =============================================================================
 namespace {
 
+template <typename Value>
+void reduce_pool_if_mpi_ready(Value& value)
+{
+#ifdef __MPI
+    int initialized = 0;
+    int finalized = 0;
+    MPI_Initialized(&initialized);
+    MPI_Finalized(&finalized);
+    if (initialized && !finalized)
+        Parallel_Reduce::reduce_pool(value);
+#endif
+}
+
+template <typename Value>
+void reduce_pool_if_mpi_ready(Value* value, const int n)
+{
+#ifdef __MPI
+    int initialized = 0;
+    int finalized = 0;
+    MPI_Initialized(&initialized);
+    MPI_Finalized(&finalized);
+    if (initialized && !finalized)
+        Parallel_Reduce::reduce_pool(value, n);
+#endif
+}
+
 template <typename T, typename Real>
 Real max_generalized_residual(
     const T* hpsi,
@@ -28,6 +56,7 @@ Real max_generalized_residual(
             const T r = hpsi[ig + j * ld] - T(eigenvalue[j]) * spsi[ig + j * ld];
             nrm2 += static_cast<Real>(std::norm(r));
         }
+        reduce_pool_if_mpi_ready(nrm2);
         max_res = std::max(max_res, std::sqrt(nrm2));
     }
     return max_res;

source/source_hsolver/ppcg/diago_ppcg_ops.hpp

@@ -1,38 +1,6 @@
 #include "source_base/kernels/math_kernel_op.h"
-#include "source_base/parallel_reduce.h"
-
 namespace hsolver {
 
-namespace {
-
-template <typename Value>
-void reduce_pool_if_mpi_ready(Value& value)
-{
-#ifdef __MPI
-    int initialized = 0;
-    int finalized = 0;
-    MPI_Initialized(&initialized);
-    MPI_Finalized(&finalized);
-    if (initialized && !finalized)
-        Parallel_Reduce::reduce_pool(value);
-#endif
-}
-
-template <typename Value>
-void reduce_pool_if_mpi_ready(Value* value, const int n)
-{
-#ifdef __MPI
-    int initialized = 0;
-    int finalized = 0;
-    MPI_Initialized(&initialized);
-    MPI_Finalized(&finalized);
-    if (initialized && !finalized)
-        Parallel_Reduce::reduce_pool(value, n);
-#endif
-}
-
-} // anonymous namespace
-
 // =============================================================================
 // Constructor
 // =============================================================================
@@ -220,11 +188,9 @@ void DiagoPPCG<T, Device>::project_against(
         for (const int bc : basis_cols)
         {
             // Full complex inner product <basis_bc | sx_c>
-            T coeff = 0;
             const T* bb = basis + bc * ld_psi_;
             const T* sc = sx.data() + c * ld_psi_;
-            for (int ig = 0; ig < n_dim_; ++ig)
-                coeff += std::conj(bb[ig]) * sc[ig];
+            const T coeff = complex_dot(bb, sc);
             if (std::abs(coeff) <= std::numeric_limits<Real>::epsilon())
                 continue;
             const T* sb = sbasis + bc * ld_psi_;

source/source_hsolver/ppcg/diago_ppcg_subspace.hpp

@@ -19,6 +19,7 @@ void DiagoPPCG<T, Device>::lock_epairs(
         Real nrm2 = 0;
         for (int ig = 0; ig < n_dim_; ++ig)
             nrm2 += static_cast<Real>(std::norm(residual[idx(ig, j, ld_psi_)]));
+        reduce_pool_if_mpi_ready(nrm2);
         const Real rnrm = std::sqrt(std::max(nrm2, static_cast<Real>(0)));
         const Real thr = std::max(static_cast<Real>(ethr_band[j]), diag_thr_);
         if (rnrm > thr)
@@ -80,6 +81,7 @@ void DiagoPPCG<T, Device>::build_small_subspace(
             for (int ig = 0; ig < n_dim_; ++ig)
                 sn2 += std::real(std::conj(x[idx(ig, j, ld_psi_)])
                                  * sx[idx(ig, j, ld_psi_)]);
+            reduce_pool_if_mpi_ready(sn2);
             Real sn = std::sqrt(std::max(sn2, static_cast<Real>(1e-30)));
             // Only scale if the norm is non-negligible; a near-zero
             // column is a converged band whose contribution is harmless.