fix: auto-load k-point to GPU in PAGED_GPU fix_k/fix_b/fix_kb

When device_memory_mode defaults to 'paged' (auto mode on GPU PW),
fix_k() returned CPU buffer pointer for k-points not yet loaded,
causing psiToRho and other GPU consumers to read CPU memory instead
of GPU memory. This produced wrong charge density and incorrect total
energy for multi-k-point GPU tests (003, 078, etc.).

Fix: make current_k_gpu_ mutable, call load_k_to_gpu in fix_k/fix_b/
fix_kb when the requested k-point differs from the loaded one. Non-PAGED
modes are unaffected (load_k_to_gpu is a no-op for ALL_GPU).

Author: dyzheng

source/source_psi/psi.cpp

@@ -536,16 +536,13 @@ void Psi<T, Device>::fix_k(const int ik) const
     if (storage_mode_ == PsiStorageMode::PAGED_GPU)
     {
         this->psi_bias = 0;
-        if (this->current_k_gpu_ == ik)
+        if (this->current_k_gpu_ != ik)
         {
-            // This k-point is loaded on GPU by load_k_to_gpu
-            this->psi_current = const_cast<T*>(this->psi);
+            const_cast<Psi*>(this)->load_k_to_gpu(ik);
         }
         else
         {
-            // Read from CPU full buffer (psi_cpu_ has all k-points)
-            this->psi_current = const_cast<T*>(psi_cpu_
-                + static_cast<size_t>(ik) * this->nbands * this->nbasis);
+            this->psi_current = const_cast<T*>(this->psi);
         }
         return;
     }
@@ -572,7 +569,14 @@ void Psi<T, Device>::fix_b(const int ib) const
     if (storage_mode_ == PsiStorageMode::PAGED_GPU)
     {
         this->psi_bias = 0;
-        this->psi_current = const_cast<T*>(this->psi);
+        if (this->current_k_gpu_ != this->current_k)
+        {
+            const_cast<Psi*>(this)->load_k_to_gpu(this->current_k);
+        }
+        else
+        {
+            this->psi_current = const_cast<T*>(this->psi);
+        }
         return;
     }
 
@@ -604,7 +608,14 @@ void Psi<T, Device>::fix_kb(const int ik, const int ib) const
     if (storage_mode_ == PsiStorageMode::PAGED_GPU)
     {
         this->psi_bias = 0;
-        this->psi_current = const_cast<T*>(this->psi);
+        if (this->current_k_gpu_ != ik)
+        {
+            const_cast<Psi*>(this)->load_k_to_gpu(ik);
+        }
+        else
+        {
+            this->psi_current = const_cast<T*>(this->psi);
+        }
         return;
     }
 

source/source_psi/psi.h

@@ -190,7 +190,7 @@ class Psi
     bool psi_cpu_owned_ = true;
     T* psi_gpu_buffer_ = nullptr;
     T* psi_gpu_transfer_buffer_ = nullptr;
-    int current_k_gpu_ = -1;
+    mutable int current_k_gpu_ = -1;
 
 #ifdef __DSP
     using delete_memory_op = base_device::memory::delete_memory_op_mt<T, Device>;