Fix(paged): resolve Psi PAGED_GPU mode issues with psi_prepare, operator(), force, and stress

- psi_prepare: in PAGED_GPU mode, fix_k returns CPU pointer when k-point
  is not loaded, causing cudaMemcpy H2D with invalid CPU destination.
  Fix by copying init data to kspw_psi CPU backing store before load_k_to_gpu.
- psi::operator()(ik,iband,ibasis): PAGED_GPU always returned CPU pointer,
  but GPU callers (force/stress via CUBLAS) need GPU pointers. Fix by
  returning GPU buffer pointer when current_k_gpu_ matches requested ik.
- forces_nl, stress_nl, stress_kin: add ensure_k_on_gpu(ik) call before
  accessing psi via operator() in k-point loops.

Author: dyzheng

source/source_psi/psi.cpp

@@ -628,8 +628,12 @@ T& Psi<T, Device>::operator()(const int ikb1, const int ikb2, const int ibasis)
 
     if (storage_mode_ == PsiStorageMode::PAGED_GPU)
     {
-        // PAGED_GPU: GPU buffer is single-k, read from CPU full buffer (k_first layout)
+        // PAGED_GPU: GPU buffer is single-k, return from GPU if loaded, else CPU
         assert(psi_cpu_ != nullptr);
+        if (this->current_k_gpu_ == ikb1)
+        {
+            return this->psi[ikb2 * this->nbasis + ibasis];
+        }
         return const_cast<T*>(psi_cpu_)[(ikb1 * this->nbands + ikb2) * this->nbasis + ibasis];
     }
 

source/source_psi/psi_prepare.cpp

@@ -178,7 +178,9 @@ void PSIPrepare<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
             {
                 if (kspw_psi->get_storage_mode() == psi::PsiStorageMode::PAGED_GPU)
                 {
-                    syncmem_h2d_op()(psi_device->get_pointer(), psi_cpu->get_pointer(), nbands_start * nbasis);
+                    T* dst_cpu = kspw_psi->get_cpu_pointer(ik);
+                    std::memcpy(dst_cpu, psi_cpu->get_pointer(), sizeof(T) * nbands_start * nbasis);
+                    kspw_psi->load_k_to_gpu(ik);
                 }
                 else
                 {

source/source_pw/module_pwdft/forces_nl.cpp

@@ -37,6 +37,9 @@ void Forces<FPTYPE, Device>::cal_force_nl(ModuleBase::matrix& forcenl,
     const int max_nbands = wg.nc;
     for (int ik = 0; ik < nks; ik++) // loop k points
     {
+        // Ensure k-point is loaded to GPU for PAGED_GPU mode
+        const_cast<psi::Psi<std::complex<FPTYPE>, Device>*>(psi_in)->ensure_k_on_gpu(ik);
+
         // skip zero weights to speed up
         int nbands_occ = wg.nc;
         while (wg(ik, nbands_occ - 1) == 0.0)

source/source_pw/module_pwdft/stress_kin.cpp

@@ -18,9 +18,12 @@ void Stress_Func<FPTYPE, Device>::stress_kin(ModuleBase::matrix& sigma,
 
 	this->ucell = &ucell_in;
 
-	hamilt::FS_Kin_tools<FPTYPE, Device> kin_tool(*this->ucell, p_kv, wfc_basis, wg);
+    hamilt::FS_Kin_tools<FPTYPE, Device> kin_tool(*this->ucell, p_kv, wfc_basis, wg);
     for (int ik = 0; ik < wfc_basis->nks; ++ik)
     {
+        // Ensure k-point is loaded to GPU for PAGED_GPU mode
+        const_cast<psi::Psi<std::complex<FPTYPE>, Device>*>(psi_in)->ensure_k_on_gpu(ik);
+
         int nbands_occ = wg.nc;
         while (wg(ik, nbands_occ - 1) == 0.0)
         {

source/source_pw/module_pwdft/stress_nl.cpp

@@ -39,6 +39,9 @@ void Stress_Func<FPTYPE, Device>::stress_nl(ModuleBase::matrix& sigma,
     const int max_nbands = wg.nc;
     for (int ik = 0; ik < nks; ik++) // loop k points
     {
+        // Ensure k-point is loaded to GPU for PAGED_GPU mode
+        const_cast<psi::Psi<std::complex<FPTYPE>, Device>*>(psi_in)->ensure_k_on_gpu(ik);
+
         // skip zero weights to speed up
         int nbands_occ = wg.nc;
         while (wg(ik, nbands_occ - 1) == 0.0)