Feat(psi): integrate k-point paging into solver and force/stress paths

- HSolverPW: load_k_to_gpu before each k-point in dav_subspace loop
- psi_prepare: PAGED_GPU mode support, skip GPU prep when paging
- setup_psi_pw: log active GPU memory mode
- diago_dav_subspace: host buffer optimization (hcc_h/scc_h),
  d_eigenvalue sync fix, diag_zhegvx CPU fallback for GPU
- forces/stress: load_k_to_gpu/ensure_k_on_gpu guards before k-point loops
- onsite_proj: ensure_k_on_gpu in cal_occupations
- kernel force/stress_op: add load_k_to_gpu before GPU kernel launch

Author: dyzheng

source/source_hsolver/diago_dav_subspace.cpp

@@ -1,5 +1,7 @@
 #include "diago_dav_subspace.h"
 
+#include <type_traits>
+
 #include "diago_iter_assist.h"
 
 #include "source_base/module_device/device.h"
@@ -103,6 +105,8 @@ Diago_DavSubspace<T, Device>::~Diago_DavSubspace()
         delmem_real_op()(this->d_precondition);
         delmem_complex_op()(this->d_scc);
         delmem_real_op()(this->d_eigenvalue);
+        delmem_complex_h_op()(this->hcc_h);
+        delmem_complex_h_op()(this->scc_h);
     }
 #endif
 }
@@ -275,17 +279,21 @@ int Diago_DavSubspace<T, Device>::diag_once(const HPsiFunc& hpsi_func,
 
 template <typename T, typename Device>
 void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
-                                            const HPsiFunc& spsi_func,
-                                            const int& dim,
-                                            const int& nbase,
-                                            const int& notconv,
-                                            T* psi_iter,
-                                            T* hpsi,
-                                            T* spsi,
-                                            T* vcc,
-                                            const int* unconv,
-                                            std::vector<Real>* eigenvalue_iter)
+                                             const HPsiFunc& spsi_func,
+                                             const int& dim,
+                                             const int& nbase,
+                                             const int& notconv_ref,
+                                             T* psi_iter,
+                                             T* hpsi,
+                                             T* spsi,
+                                             T* vcc,
+                                             const int* unconv,
+                                             std::vector<Real>* eigenvalue_iter)
 {
+    // Local copy: notconv_ref is a reference to this->notconv which may be
+    // modified to 0 inside this function (all_zero path). We need the original
+    // value for the remainder of the function.
+    const int notconv = notconv_ref;
     ModuleBase::timer::start("Diago_DavSubspace", "cal_grad");
 
     for (size_t i = 0; i < notconv; i++)
@@ -436,13 +444,23 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
         {
             if (psi_norm_host[i] <= 1.0e-12)
             {
-                std::cout << "Diago_DavSubspace::cal_grad: psi_norm <= 0 for band " << i << std::endl;
-                std::cout << "This may be due to npwx < nbands: the number of plane waves is less than" << std::endl;
-                std::cout << "the number of bands, leading to a rank-deficient problem." << std::endl;
-                std::cout << "Please increase ecutwfc or reduce nbands." << std::endl;
-                delmem_real_h_op()(psi_norm_host);
-                delmem_real_op()(psi_norm);
-                ModuleBase::WARNING_QUIT("cal_grad", "psi_norm <= 0");
+                std::cout << "Diago_DavSubspace::cal_grad: norm of new direction for band " << i
+                          << " is too small, treating as converged." << std::endl;
+            }
+        }
+        {
+            bool all_zero = true;
+            for (int i = 0; i < notconv; i++)
+            {
+                if (psi_norm_host[i] > 1.0e-12)
+                {
+                    all_zero = false;
+                    break;
+                }
+            }
+            if (all_zero && notconv > 0)
+            {
+                this->notconv = 0;
             }
         }
         delmem_real_h_op()(psi_norm_host);
@@ -466,12 +484,23 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
         {
             if (psi_norm[i] <= 1.0e-12)
             {
-                std::cout << "Diago_DavSubspace::cal_grad: psi_norm <= 0 for band " << i << std::endl;
-                std::cout << "This may be due to npwx < nbands: the number of plane waves is less than" << std::endl;
-                std::cout << "the number of bands, leading to a rank-deficient problem." << std::endl;
-                std::cout << "Please increase ecutwfc or reduce nbands." << std::endl;
-                delmem_real_h_op()(psi_norm);
-                ModuleBase::WARNING_QUIT("cal_grad", "psi_norm <= 0");
+                std::cout << "Diago_DavSubspace::cal_grad: norm of new direction for band " << i
+                          << " is too small, treating as converged." << std::endl;
+            }
+        }
+        {
+            bool all_zero = true;
+            for (int i = 0; i < notconv; i++)
+            {
+                if (psi_norm[i] > 1.0e-12)
+                {
+                    all_zero = false;
+                    break;
+                }
+            }
+            if (all_zero && notconv > 0)
+            {
+                this->notconv = 0;
             }
         }
         delmem_real_h_op()(psi_norm);
@@ -585,8 +614,30 @@ void Diago_DavSubspace<T, Device>::cal_elem(const int& dim,
         mtfunc::dsp_dav_subspace_reduce(hcc, scc, nbase, this->nbase_x, this->notconv, this->diag_comm.comm);
 #else
         assert(this->diag_comm.comm == POOL_WORLD);
-        Parallel_Reduce::reduce_pool(hcc + nbase * this->nbase_x, notconv * this->nbase_x);
-        Parallel_Reduce::reduce_pool(scc + nbase * this->nbase_x, notconv * this->nbase_x);
+        // For GPU, need to copy to CPU before MPI reduction
+        if constexpr (std::is_same<Device, base_device::DEVICE_GPU>::value)
+        {
+            // Allocate CPU buffers if needed
+            if (this->hcc_h == nullptr)
+            {
+                resmem_complex_h_op()(this->hcc_h, this->nbase_x * this->nbase_x, "DAV::hcc_h");
+                resmem_complex_h_op()(this->scc_h, this->nbase_x * this->nbase_x, "DAV::scc_h");
+            }
+            // Copy from GPU to CPU
+            syncmem_d2h_op()(this->hcc_h, hcc + nbase * this->nbase_x, notconv * this->nbase_x);
+            syncmem_d2h_op()(this->scc_h, scc + nbase * this->nbase_x, notconv * this->nbase_x);
+            // Reduce on CPU
+            Parallel_Reduce::reduce_pool(this->hcc_h, notconv * this->nbase_x);
+            Parallel_Reduce::reduce_pool(this->scc_h, notconv * this->nbase_x);
+            // Copy back to GPU
+            syncmem_h2d_op()(hcc + nbase * this->nbase_x, this->hcc_h, notconv * this->nbase_x);
+            syncmem_h2d_op()(scc + nbase * this->nbase_x, this->scc_h, notconv * this->nbase_x);
+        }
+        else
+        {
+            Parallel_Reduce::reduce_pool(hcc + nbase * this->nbase_x, notconv * this->nbase_x);
+            Parallel_Reduce::reduce_pool(scc + nbase * this->nbase_x, notconv * this->nbase_x);
+        }
 #endif
     }
 #endif
@@ -615,9 +666,24 @@ void Diago_DavSubspace<T, Device>::diag_zhegvx(const int& nbase,
 #if defined(__CUDA) || defined(__ROCM)
         if (this->diag_comm.rank == 0)
         {
-            syncmem_complex_op()(this->d_scc, scc, nbase * this->nbase_x);
-            ct::kernels::lapack_hegvd<T, ct_Device>()(nbase, this->nbase_x, this->hcc, this->d_scc, this->d_eigenvalue, this->vcc);
-            syncmem_var_d2h_op()((*eigenvalue_iter).data(), this->d_eigenvalue, this->nbase_x);
+            // Copy hcc and scc from GPU to CPU for robust eigensolving
+            std::vector<T> hcc_h(nbase * this->nbase_x, *this->zero);
+            std::vector<T> scc_h(nbase * this->nbase_x, *this->zero);
+            syncmem_d2h_op()(hcc_h.data(), this->hcc, nbase * this->nbase_x);
+            syncmem_d2h_op()(scc_h.data(), scc, nbase * this->nbase_x);
+
+            // Solve on CPU using hegvx (more robust than hegvd for ill-conditioned overlap)
+            std::vector<T> vcc_h(nbase * this->nbase_x, *this->zero);
+            hegvx_op<T, base_device::DEVICE_CPU>()(this->cpu_ctx,
+                                                    nbase,
+                                                    this->nbase_x,
+                                                    hcc_h.data(),
+                                                    scc_h.data(),
+                                                    nband,
+                                                    (*eigenvalue_iter).data(),
+                                                    vcc_h.data());
+            // Copy eigenvectors back to GPU
+            syncmem_h2d_op()(this->vcc, vcc_h.data(), nbase * this->nbase_x);
         }
 #endif
     }
@@ -715,9 +781,25 @@ void Diago_DavSubspace<T, Device>::diag_zhegvx(const int& nbase,
     if (this->diag_comm.nproc > 1)
     {
         // vcc: nbase * nband
-        for (int i = 0; i < nband; i++)
+        if constexpr (std::is_same<Device, base_device::DEVICE_GPU>::value)
         {
-            MPI_Bcast(&vcc[i * this->nbase_x], nbase, MPI_DOUBLE_COMPLEX, 0, this->diag_comm.comm);
+            // Copy vcc from GPU to CPU for MPI broadcast
+            T* vcc_h = nullptr;
+            resmem_complex_h_op()(vcc_h, nband * this->nbase_x, "DAV::vcc_h");
+            syncmem_d2h_op()(vcc_h, vcc, nband * this->nbase_x);
+            for (int i = 0; i < nband; i++)
+            {
+                MPI_Bcast(&vcc_h[i * this->nbase_x], nbase, MPI_DOUBLE_COMPLEX, 0, this->diag_comm.comm);
+            }
+            syncmem_h2d_op()(vcc, vcc_h, nband * this->nbase_x);
+            delmem_complex_h_op()(vcc_h);
+        }
+        else
+        {
+            for (int i = 0; i < nband; i++)
+            {
+                MPI_Bcast(&vcc[i * this->nbase_x], nbase, MPI_DOUBLE_COMPLEX, 0, this->diag_comm.comm);
+            }
         }
         MPI_Bcast((*eigenvalue_iter).data(), nband, MPI_DOUBLE, 0, this->diag_comm.comm);
     }
@@ -795,6 +877,14 @@ void Diago_DavSubspace<T, Device>::refresh(const int& dim,
 
     if (this->device == base_device::GpuDevice)
     {
+#if defined(__CUDA) || defined(__ROCM)
+        // this->d_eigenvalue was last synced inside cal_grad, i.e. BEFORE the
+        // latest diag_zhegvx. eigenvalue_in_hsolver holds the up-to-date
+        // eigenvalues, so refresh the device copy here; otherwise the restarted
+        // subspace Hamiltonian gets a stale diagonal and the Davidson
+        // iteration diverges on GPU.
+        syncmem_var_h2d_op()(this->d_eigenvalue, eigenvalue_in_hsolver, nband);
+#endif
         refresh_hcc_scc_vcc_op<T, Device>()(nbase, hcc, scc, vcc, this->nbase_x, this->d_eigenvalue, this->one_);
     }
     else

source/source_hsolver/diago_dav_subspace.h

@@ -190,12 +190,18 @@ class Diago_DavSubspace
     using syncmem_h2d_op = base_device::memory::synchronize_memory_op<T, Device, base_device::DEVICE_CPU>;
     using syncmem_d2h_op = base_device::memory::synchronize_memory_op<T, base_device::DEVICE_CPU, Device>;
 
+    using resmem_complex_h_op = base_device::memory::resize_memory_op<T, base_device::DEVICE_CPU>;
+    using delmem_complex_h_op = base_device::memory::delete_memory_op<T, base_device::DEVICE_CPU>;
+
     // Note that ct_Device is different from base_device!
     using ct_Device = typename ct::PsiToContainer<Device>::type;
     // using hegvd_op = container::kernels::lapack_hegvd<T, ct_Device>;
 
     const T *one = nullptr, *zero = nullptr, *neg_one = nullptr;
     const T one_ = static_cast<T>(1.0), zero_ = static_cast<T>(0.0), neg_one_ = static_cast<T>(-1.0);
+
+    T* hcc_h = nullptr;
+    T* scc_h = nullptr;
 };
 
 } // namespace hsolver

source/source_hsolver/hsolver_pw.cpp

@@ -12,6 +12,7 @@
 #include "source_hsolver/diago_dav_subspace.h"
 #include "source_hsolver/diago_david.h"
 #include "source_hsolver/diago_iter_assist.h"
+#include "source_io/module_parameter/parameter.h"
 #include "source_psi/psi.h"
 #include "source_estate/elecstate_tools.h"
 
@@ -109,21 +110,23 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
             // update H(k) for each k point
             pHamilt->updateHk(ik);
 
-
+            psi.load_k_to_gpu(ik);
 
             // update psi pointer for each k point
             psi.fix_k(ik);
 
             // If using k-point continuity and not first k-point, propagate from parent
-            if (ik > 0 && count == 0 && k_parent.find(ik) != k_parent.end()) {
+            // Skip in PAGED_GPU mode: only one k-point resides on GPU at a time
+            if (ik > 0 && count == 0 && k_parent.find(ik) != k_parent.end()
+                && psi.get_storage_mode() != psi::PsiStorageMode::PAGED_GPU) {
                 propagate_psi(psi, k_parent[ik], ik);
             }
 
             // template add precondition calculating here
             update_precondition(precondition, ik, this->wfc_basis->npwk[ik], Real(pes->pot->get_vl_of_0()));
 
             // use smooth threshold for all iter methods
-            if (this->diago_smooth_ethr == true)
+            if (PARAM.inp.diago_smooth_ethr == true)
             {
                 this->cal_smooth_ethr(pes->klist->wk[ik],
                                     &pes->wg(ik, 0),
@@ -136,6 +139,8 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
             // solve eigenvector and eigenvalue for H(k)
             this->hamiltSolvePsiK(pHamilt, psi, precondition, eigenvalues.data() + ik * psi.get_nbands(), this->wfc_basis->nks);
 
+            psi.store_k_from_gpu(ik);
+
             if (skip_charge)
             {
                 GlobalV::ofs_running << " Average iterative diagonalization steps for k-points " << ik
@@ -152,7 +157,7 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
             // update H(k) for each k point
             pHamilt->updateHk(ik);
 
-
+            psi.load_k_to_gpu(ik);
 
             // update psi pointer for each k point
             psi.fix_k(ik);
@@ -161,7 +166,7 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
             update_precondition(precondition, ik, this->wfc_basis->npwk[ik], Real(pes->pot->get_vl_of_0()));
 
             // use smooth threshold for all iter methods
-            if (this->diago_smooth_ethr == true)
+            if (PARAM.inp.diago_smooth_ethr == true)
             {
                 this->cal_smooth_ethr(pes->klist->wk[ik],
                                     &pes->wg(ik, 0),
@@ -174,6 +179,8 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
             // solve eigenvector and eigenvalue for H(k)
             this->hamiltSolvePsiK(pHamilt, psi, precondition, eigenvalues.data() + ik * psi.get_nbands(), this->wfc_basis->nks);
 
+            psi.store_k_from_gpu(ik);
+
             // output iteration information and reset avg_iter
             if (skip_charge)
             {
@@ -211,7 +218,7 @@ void HSolverPW<T, Device>::solve(hamilt::Hamilt<T, Device>* pHamilt,
     elecstate::calEBand(_pes_pw->ekb,_pes_pw->wg,_pes_pw->f_en);
     if (skip_charge)
     {
-        if (this->use_uspp)
+        if (PARAM.globalv.use_uspp)
         {
             reinterpret_cast<elecstate::ElecStatePW<T, Device>*>(pes)->cal_becsum(psi);
         }
@@ -319,7 +326,7 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
         const int nbasis = psi.get_nbasis();
         const int ndim = psi.get_current_ngk();
         DiagoBPCG<T, Device> bpcg(pre_condition.data());
-        bpcg.init_iter(this->nbands, nband_l, nbasis, ndim);
+        bpcg.init_iter(PARAM.inp.nbands, nband_l, nbasis, ndim);
         bpcg.diag(hpsi_func, psi.get_pointer(), eigenvalue, this->ethr_band);
     }
     else if (this->method == "dav_subspace")
@@ -330,12 +337,12 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
                                                   psi.get_nbands(),
                                                   psi.get_k_first() ? psi.get_current_ngk()
                                                                     : psi.get_nk() * psi.get_nbasis(),
-                                                  this->pw_diag_ndim,
+                                                  PARAM.inp.pw_diag_ndim,
                                                   this->diag_thr,
                                                   this->diag_iter_max,
                                                   comm_info,
-                                                  this->diag_subspace,
-                                                  this->nb2d);
+                                                  PARAM.inp.diag_subspace,
+                                                  PARAM.inp.nb2d);
 
         DiagoIterAssist<T, Device>::avg_iter += static_cast<double>(
             dav_subspace.diag(hpsi_func,
@@ -365,7 +372,7 @@ void HSolverPW<T, Device>::hamiltSolvePsiK(hamilt::Hamilt<T, Device>* hm,
         const int nband = psi.get_nbands();            /// number of eigenpairs sought
         const int ld_psi = psi.get_nbasis();           /// leading dimension of psi
 
-        DiagoDavid<T, Device> david(pre_condition.data(), nband, dim, this->pw_diag_ndim, comm_info);
+        DiagoDavid<T, Device> david(pre_condition.data(), nband, dim, PARAM.inp.pw_diag_ndim, comm_info);
         // do diag and add davidson iteration counts up to avg_iter
         DiagoIterAssist<T, Device>::avg_iter += static_cast<double>(
              david.diag(hpsi_func,

source/source_psi/psi_prepare.cpp

@@ -176,7 +176,16 @@ void PSIPrepare<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
             this->psi_initer->init_psig(psi_cpu->get_pointer(), ik);
             if (psi_device->get_pointer() != psi_cpu->get_pointer())
             {
-                syncmem_h2d_op()(psi_device->get_pointer(), psi_cpu->get_pointer(), nbands_start * nbasis);
+                if (kspw_psi->get_storage_mode() == psi::PsiStorageMode::PAGED_GPU)
+                {
+                    T* dst_cpu = kspw_psi->get_cpu_pointer(ik);
+                    std::memcpy(dst_cpu, psi_cpu->get_pointer(), sizeof(T) * nbands_l * nbasis);
+                    kspw_psi->load_k_to_gpu(ik);
+                }
+                else
+                {
+                    syncmem_h2d_op()(psi_device->get_pointer(), psi_cpu->get_pointer(), nbands_start * nbasis);
+                }
             }
 
 
@@ -208,7 +217,13 @@ void PSIPrepare<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
             {
                 if (psi_device->get_pointer() != kspw_psi->get_pointer())
                 {
-                    syncmem_complex_op()(kspw_psi->get_pointer(), psi_device->get_pointer(), nbands_l * nbasis);
+                    if (kspw_psi->get_storage_mode() == psi::PsiStorageMode::PAGED_GPU)
+                    {
+                    }
+                    else
+                    {
+                        syncmem_complex_op()(kspw_psi->get_pointer(), psi_device->get_pointer(), nbands_l * nbasis);
+                    }
                 }
             }
         }