Refactor hpsi_func of dav_subspace

Author: Cstandardlib

source/module_hsolver/diago_dav_subspace.cpp

@@ -124,7 +124,7 @@ int Diago_DavSubspace<T, Device>::diag_once(const HPsiFunc& hpsi_func,
 
     // compute h*psi_in_iter
     // NOTE: bands after the first n_band should yield zero
-    hpsi_func(this->psi_in_iter, this->hphi, this->nbase_x, this->dim, 0, this->nbase_x - 1);
+    hpsi_func(this->psi_in_iter, this->hphi, this->dim, this->nbase_x);
 
     // at this stage, notconv = n_band and nbase = 0
     // note that nbase of cal_elem is an inout parameter: nbase := nbase + notconv
@@ -421,7 +421,7 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
     }
 
     // update hpsi[:, nbase:nbase+notconv]
-    hpsi_func(psi_iter, &hphi[nbase * this->dim], this->nbase_x, this->dim, nbase, nbase + notconv - 1);
+    hpsi_func(psi_iter + nbase * dim, hphi + nbase * this->dim, this->dim, notconv);
 
     ModuleBase::timer::tick("Diago_DavSubspace", "cal_grad");
     return;
@@ -886,7 +886,7 @@ void Diago_DavSubspace<T, Device>::diagH_subspace(T* psi_pointer, // [in] & [out
 
     {
         // do hPsi for all bands
-        hpsi_func(psi_pointer, hphi, n_band, dmax, 0, nstart - 1);
+        hpsi_func(psi_pointer, hphi, dmax, nstart);
 
         gemm_op<T, Device>()(ctx,
                              'C',

source/module_hsolver/diago_dav_subspace.h

@@ -31,7 +31,7 @@ class Diago_DavSubspace : public DiagH<T, Device>
 
     virtual ~Diago_DavSubspace() override;
 
-    using HPsiFunc = std::function<void(T*, T*, const int, const int, const int, const int)>;
+    using HPsiFunc = std::function<void(T*, T*, const int, const int)>;
 
     int diag(const HPsiFunc& hpsi_func,
              T* psi_in,

source/module_hsolver/hsolver_pw.cpp

@@ -434,18 +434,17 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
     else if (this->method == "dav_subspace")
     {
         auto ngk_pointer = psi.get_ngk_pointer();
-        auto hpsi_func = [hm, ngk_pointer](T* psi_in,
-                                           T* hpsi_out,
-                                           const int nband_in,
-                                           const int nbasis_in,
-                                           const int band_index1,
-                                           const int band_index2) {
+        // hpsi_func (X, HX, ld, nvec) -> HX = H(X), X and HX blockvectors of size ld x nvec
+        auto hpsi_func = [hm, ngk_pointer](T *psi_in,
+                                           T *hpsi_out,
+                                           const int ldPsi,
+                                           const int nvec) {
             ModuleBase::timer::tick("DavSubspace", "hpsi_func");
 
             // Convert "pointer data stucture" to a psi::Psi object
-            auto psi_iter_wrapper = psi::Psi<T, Device>(psi_in, 1, nband_in, nbasis_in, ngk_pointer);
+            auto psi_iter_wrapper = psi::Psi<T, Device>(psi_in, 1, nvec, ldPsi, ngk_pointer);
 
-            psi::Range bands_range(true, 0, band_index1, band_index2);
+            psi::Range bands_range(true, 0, 0, nvec-1);
 
             using hpsi_info = typename hamilt::Operator<T, Device>::hpsi_info;
             hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out);

source/module_lr/hsolver_lrtd.cpp

@@ -118,16 +118,14 @@ namespace LR
                     false, //always do the subspace diag (check the implementation)
                     comm_info);
 
-                std::function<void(T*, T*, const int, const int, const int, const int)> hpsi_func = [pHamilt](
+                auto hpsi_func = [pHamilt](
                     T* psi_in,
                     T* hpsi_out,
-                    const int nband_in,
-                    const int nbasis_in,
-                    const int band_index1,
-                    const int band_index2)
+                    const int ldPsi,
+                    const int nvec)
                     {
-                        auto psi_iter_wrapper = psi::Psi<T, Device>(psi_in, 1, nband_in, nbasis_in, nullptr);
-                        psi::Range bands_range(true, 0, band_index1, band_index2);
+                        auto psi_iter_wrapper = psi::Psi<T, Device>(psi_in, 1, nvec, ldPsi, nullptr);
+                        psi::Range bands_range(true, 0, 0, nvec-1);
                         using hpsi_info = typename hamilt::Operator<T, Device>::hpsi_info;
                         hpsi_info info(&psi_iter_wrapper, bands_range, hpsi_out);
                         pHamilt->ops->hPsi(info);