Add pw_diag_rr_step to decouple the PPCG Rayleigh-Ritz interval from the block size

Author: cheerly-pku

docs/advanced/input_files/input-main.md

@@ -66,6 +66,7 @@
     - [use\_k\_continuity](#use_k_continuity)
     - [pw\_diag\_nmax](#pw_diag_nmax)
     - [pw\_diag\_ndim](#pw_diag_ndim)
+    - [pw\_diag\_rr\_step](#pw_diag_rr_step)
     - [diago\_cg\_prec](#diago_cg_prec)
   - [Numerical atomic orbitals related variables](#numerical-atomic-orbitals-related-variables)
     - [lmaxmax](#lmaxmax)
@@ -1085,9 +1086,15 @@
 
 - **Type**: Integer
 - **Availability**: *[`basis_type`](#basis_type)==pw and [`ks_solver`](#ks_solver) in [dav, dav_subspace, ppcg]*
-- **Description**: Only useful when you use ks_solver = dav, dav_subspace, or ppcg. It indicates dimension of workspace(number of wavefunction packets, at least 2 needed) for the Davidson method, and the PPCG block size/Rayleigh-Ritz interval for the PPCG method. A larger value may yield a smaller number of iterations in the algorithm but uses more memory and more CPU time in subspace diagonalization.
+- **Description**: Only useful when you use ks_solver = dav, dav_subspace, or ppcg. It indicates dimension of workspace(number of wavefunction packets, at least 2 needed) for the Davidson method, and the block size for the PPCG method. A larger value may yield a smaller number of iterations in the algorithm but uses more memory and more CPU time in subspace diagonalization.
 - **Default**: 4
 
+### pw_diag_rr_step
+
+- **Type**: Integer
+- **Availability**: *[`basis_type`](#basis_type)==pw and [`ks_solver`](#ks_solver) == ppcg*- **Description**: Only useful when you use ks_solver = ppcg. It controls how often (in subspace iterations) H and S are re-applied to reset the accumulated rounding drift after the Rayleigh-Ritz rotation. A larger value reduces the number of H/S applications and thus the wall time without changing the iteration count in well-conditioned cases; a smaller value is more robust against rounding drift in ill-conditioned problems.
+- **Default**: 16
+
 ### diago_cg_prec
 
 - **Type**: Integer

docs/parameters.yaml

@@ -1009,10 +1009,18 @@ parameters:
     category: Plane wave related variables
     type: Integer
     description: |
-      Only useful when you use ks_solver = dav, dav_subspace, or ppcg. It indicates dimension of workspace(number of wavefunction packets, at least 2 needed) for the Davidson method, and the PPCG block size/Rayleigh-Ritz interval for the PPCG method. A larger value may yield a smaller number of iterations in the algorithm but uses more memory and more CPU time in subspace diagonalization.
+      Only useful when you use ks_solver = dav, dav_subspace, or ppcg. It indicates dimension of workspace(number of wavefunction packets, at least 2 needed) for the Davidson method, and the block size for the PPCG method. A larger value may yield a smaller number of iterations in the algorithm but uses more memory and more CPU time in subspace diagonalization.
     default_value: "4"
     unit: ""
     availability: "basis_type==pw and ks_solver in [dav, dav_subspace, ppcg]"
+  - name: pw_diag_rr_step
+    category: Plane wave related variables
+    type: Integer
+    description: |
+      Only useful when you use ks_solver = ppcg. It controls how often (in subspace iterations) H and S are re-applied to reset the accumulated rounding drift after the Rayleigh-Ritz rotation. A larger value reduces the number of H/S applications and thus the wall time without changing the iteration count in well-conditioned cases; a smaller value is more robust against rounding drift in ill-conditioned problems.
+    default_value: "16"
+    unit: ""
+    availability: "basis_type==pw and ks_solver==ppcg"
   - name: diago_cg_prec
     category: Plane wave related variables
     type: Integer

source/source_hsolver/diago_iter_assist.h

@@ -22,6 +22,7 @@ class DiagoIterAssist
     static Real PW_DIAG_THR;
     static int PW_DIAG_NMAX;
     static int PW_DIAG_NDIM;
+    static int PW_DIAG_RR_STEP;
 
     static Real LCAO_DIAG_THR;
     static int LCAO_DIAG_NMAX;
@@ -161,6 +162,9 @@ int DiagoIterAssist<T, Device>::PW_DIAG_NMAX = 30;
 template <typename T, typename Device>
 int DiagoIterAssist<T, Device>::PW_DIAG_NDIM = 4;
 
+template <typename T, typename Device>
+int DiagoIterAssist<T, Device>::PW_DIAG_RR_STEP = 16;
+
 template <typename T, typename Device>
 typename DiagoIterAssist<T, Device>::Real DiagoIterAssist<T, Device>::PW_DIAG_THR = 1.0e-2;
 

source/source_hsolver/diago_params.cpp

@@ -16,6 +16,7 @@ void setup_diago_params_pw(const int istep,
     DiagoIterAssist<T, Device>::SCF_ITER = iter;
     DiagoIterAssist<T, Device>::PW_DIAG_THR = ethr;
     DiagoIterAssist<T, Device>::PW_DIAG_NDIM = inp.pw_diag_ndim;
+    DiagoIterAssist<T, Device>::PW_DIAG_RR_STEP = inp.pw_diag_rr_step;
 
     if (inp.calculation != "nscf")
     {
@@ -43,6 +44,7 @@ void setup_diago_params_sdft(const int istep,
     DiagoIterAssist<T, Device>::PW_DIAG_THR = ethr;
     DiagoIterAssist<T, Device>::PW_DIAG_NMAX = inp.pw_diag_nmax;
     DiagoIterAssist<T, Device>::PW_DIAG_NDIM = inp.pw_diag_ndim;
+    DiagoIterAssist<T, Device>::PW_DIAG_RR_STEP = inp.pw_diag_rr_step;
 }
 
 /// Template instantiation for CPU

source/source_hsolver/hsolver_pw.cpp

@@ -40,16 +40,17 @@ double run_ppcg_pw(const HPsiFunc& hpsi_func,
                    const double diag_thr,
                    const int diag_iter_max,
                    const int pw_diag_ndim,
+                   const int rr_step,
                    const bool gamma_only,
                    std::true_type)
 {
     const int sbsize = std::max(1, std::min(nband, pw_diag_ndim));
-    const int rr_step = std::max(1, pw_diag_ndim);
+    const int rr_step_safe = std::max(1, rr_step);
 
     DiagoPPCG<T, Device> ppcg(Real(diag_thr),
                               diag_iter_max,
                               sbsize,
-                              rr_step,
+                              rr_step_safe,
                               gamma_only,
                               PpcgStrategy::BLOCK_SUBSPACE);
 
@@ -77,11 +78,12 @@ double run_ppcg_pw(const HPsiFunc& hpsi_func,
                    const double diag_thr,
                    const int diag_iter_max,
                    const int pw_diag_ndim,
+                   const int rr_step,
                    const bool gamma_only,
                    std::false_type)
 {
     const int sbsize = std::max(1, std::min(nband, pw_diag_ndim));
-    const int rr_step = std::max(1, pw_diag_ndim);
+    const int rr_step_safe = std::max(1, rr_step);
     const int nelem = ld_psi * nband;
 
     // Transitional GPU path: keep PPCG's control logic and small dense solves
@@ -128,7 +130,7 @@ double run_ppcg_pw(const HPsiFunc& hpsi_func,
     DiagoPPCG<T, base_device::DEVICE_CPU> ppcg(Real(diag_thr),
                                                diag_iter_max,
                                                sbsize,
-                                               rr_step,
+                                               rr_step_safe,
                                                gamma_only,
                                                PpcgStrategy::BLOCK_SUBSPACE);
     const double avg_iter = ppcg.diag(bridge_hpsi,
@@ -518,6 +520,7 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
             this->diag_thr,
             this->diag_iter_max,
             DiagoIterAssist<T, Device>::PW_DIAG_NDIM,
+            DiagoIterAssist<T, Device>::PW_DIAG_RR_STEP,
             this->wfc_basis->gamma_only,
             std::is_same<Device, base_device::DEVICE_CPU>());
     }

source/source_hsolver/test/diago_compare_test.cpp

@@ -40,6 +40,12 @@
 using T = std::complex<double>;
 using Real = double;
 
+// Optional PPCG parameter overrides (set from argv) for exploring the block
+// size (sbsize) and Rayleigh-Ritz frequency (rr_step).  A negative value keeps
+// the default used by the comparison benchmark (sbsize = nband, rr_step = 16).
+static int g_sbsize = -1;
+static int g_rr_step = -1;
+
 // Total heap memory currently allocated (bytes).  Used to compare the peak
 // working memory of the solvers: PPCG keeps a bounded subspace, while
 // Davidson grows its basis with the number of iterations.
@@ -217,7 +223,9 @@ static Result run_ppcg(const std::vector<T>& H, int n, int nband, const std::vec
     std::vector<T> psi = psi0;
     std::vector<Real> eval(nband, 0.0);
     long mem0 = heap_bytes();
-    hsolver::DiagoPPCG<T, hsolver::base_device::DEVICE_CPU> solver(1e-8, 500, nband, std::min(nband, 4), false,
+    const int sbsize = (g_sbsize > 0) ? g_sbsize : nband;
+    const int rr_step = (g_rr_step > 0) ? g_rr_step : 16;
+    hsolver::DiagoPPCG<T, hsolver::base_device::DEVICE_CPU> solver(1e-8, 500, sbsize, rr_step, false,
                                                                    hsolver::PpcgStrategy::BLOCK_SUBSPACE);
     auto h_op = [&H, n](T* in, T* out, int ld, int nc) { dense_h_multiply(H.data(), n, in, out, ld, nc); };
     auto t0 = std::chrono::high_resolution_clock::now();
@@ -339,7 +347,7 @@ int main(int argc, char** argv)
         int sparsity;
     };
     // Without arguments a small default grid is used.  To benchmark a single
-    // (possibly large) problem, pass:  <n> <nband> <sparsity_pct>
+    // (possibly large) problem, pass:  <n> <nband> <sparsity_pct> [sbsize] [rr_step]
     std::vector<Case> cases;
     if (argc >= 4)
     {
@@ -351,6 +359,14 @@ int main(int argc, char** argv)
             {50, 10, 0}, {50, 10, 60}, {100, 10, 60}, {200, 10, 80}, {500, 10, 80},
         };
     }
+    if (argc >= 5)
+    {
+        g_sbsize = std::atoi(argv[4]);
+    }
+    if (argc >= 6)
+    {
+        g_rr_step = std::atoi(argv[5]);
+    }
 
     std::printf("\n=== Solver comparison (identical H, psi0, ethr) ===\n");
     std::printf("%-5s %-5s %-6s %-10s %-14s %-12s %-10s %-12s\n", "n", "nband", "spars", "solver", "wall_time(s)", "avg_iter",

source/source_io/module_parameter/input_parameter.h

@@ -91,6 +91,7 @@ struct Input_para
     double pw_diag_thr = 0.01;      ///< used in cg method
     bool diago_smooth_ethr = false; ///< smooth ethr for iter methods
     int pw_diag_ndim = 4;           ///< dimension of workspace for Davidson diagonalization
+    int pw_diag_rr_step = 16;       ///< Rayleigh-Ritz re-application interval for PPCG diagonalization
     int diago_cg_prec = 1;          ///< mohan add 2012-03-31
     int diag_subspace = 0;          // 0: Lapack, 1: elpa, 2: scalapack
     bool use_k_continuity = false;  ///< whether to use k-point continuity for initializing wave functions

source/source_io/module_parameter/read_inp_estruc.cpp

@@ -1104,13 +1104,25 @@ Use case: When experimental or high-level theoretical results suggest that the S
         item.annotation = "dimension of workspace for iterative PW diagonalization";
         item.category = "Plane wave related variables";
         item.type = "Integer";
-        item.description = "Only useful when you use ks_solver = dav, dav_subspace, or ppcg. It indicates dimension of workspace(number of wavefunction packets, at least 2 needed) for the Davidson method, and the PPCG block size/Rayleigh-Ritz interval for the PPCG method. A larger value may yield a smaller number of iterations in the algorithm but uses more memory and more CPU time in subspace diagonalization.";
+        item.description = "Only useful when you use ks_solver = dav, dav_subspace, or ppcg. It indicates dimension of workspace(number of wavefunction packets, at least 2 needed) for the Davidson method, and the block size for the PPCG method. A larger value may yield a smaller number of iterations in the algorithm but uses more memory and more CPU time in subspace diagonalization.";
         item.default_value = "4";
         item.unit = "";
         item.set_availability("basis_type==pw and ks_solver in [dav, dav_subspace, ppcg]");
         read_sync_int(input.pw_diag_ndim);
         this->add_item(item);
     }
+    {
+        Input_Item item("pw_diag_rr_step");
+        item.annotation = "Rayleigh-Ritz re-application interval for PPCG";
+        item.category = "Plane wave related variables";
+        item.type = "Integer";
+        item.description = "Only useful when you use ks_solver = ppcg. It controls how often (in subspace iterations) H and S are re-applied to reset the accumulated rounding drift after the Rayleigh-Ritz rotation. A larger value reduces the number of H/S applications and thus the wall time without changing the iteration count in well-conditioned cases; a smaller value is more robust against rounding drift in ill-conditioned problems.";
+        item.default_value = "16";
+        item.unit = "";
+        item.set_availability("basis_type==pw and ks_solver==ppcg");
+        read_sync_int(input.pw_diag_rr_step);
+        this->add_item(item);
+    }
     {
         Input_Item item("diago_cg_prec");
         item.annotation = "diago_cg_prec";

source/source_io/test/read_input_ptest.cpp

@@ -152,6 +152,7 @@ TEST_F(InputParaTest, ParaRead)
     EXPECT_EQ(param.inp.pw_diag_nmax, 50);
     EXPECT_EQ(param.inp.diago_cg_prec, 1);
     EXPECT_EQ(param.inp.pw_diag_ndim, 4);
+    EXPECT_EQ(param.inp.pw_diag_rr_step, 16);
     EXPECT_DOUBLE_EQ(param.inp.pw_diag_thr, 1.0e-2);
     EXPECT_FALSE(param.inp.diago_smooth_ethr);
     EXPECT_EQ(param.inp.nb2d, 0);