Feat(psi): add PsiStorageMode and device_memory_mode for GPU with PW base

source/source_estate/test/CMakeLists.txt

@@ -38,7 +38,7 @@ AddTest(
 AddTest(
   TARGET MODULE_ESTATE_elecstate_base
   LIBS parameter  base device
-  SOURCES elecstate_base_test.cpp ../elecstate.cpp  ../elecstate_tools.cpp ../occupy.cpp ../../source_psi/psi.cpp
+  SOURCES elecstate_base_test.cpp ../elecstate.cpp  ../elecstate_tools.cpp ../occupy.cpp ../../source_psi/psi.cpp ../../source_psi/psi_paging.cpp
   ../../source_base/module_fft/fft_bundle.cpp ../../source_base/module_fft/fft_cpu.cpp
 )
 
@@ -52,6 +52,7 @@ AddTest(
     ../occupy.cpp
     ../module_charge/charge_mpi.cpp
     ../../source_psi/psi.cpp
+    ../../source_psi/psi_paging.cpp
     ../../source_base/module_device/memory_op.cpp
 )
 

source/source_hsolver/hsolver_pw.cpp

@@ -109,13 +109,14 @@ 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()) {
+            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);
             }
 
@@ -136,6 +137,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 +155,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);
@@ -174,6 +177,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)
             {

source/source_io/module_parameter/input_parameter.h

@@ -70,6 +70,7 @@ struct Input_para
     double min_dist_coef = 0.2;                     ///< allowed minimum distance between two atoms
 
     std::string device = "cpu";
+    std::string device_memory_mode = "";    ///< GPU memory mode: "" (auto), "full_gpu" (all on GPU), "paged" (CPU storage + k-point paging)
     std::string precision = "double";
     std::string gint_precision = "double";
     bool timer_enable_nvtx = false;

source/source_io/module_parameter/read_input_item_system.cpp

@@ -743,6 +743,31 @@ Available options are:
         };
         this->add_item(item);
     }
+    {
+        Input_Item item("device_memory_mode");
+        item.annotation = "GPU memory strategy for wavefunction storage";
+        item.category = "System variables";
+        item.type = "String";
+        item.description = R"(Controls GPU memory strategy for wavefunction storage.
+* "" (empty/auto): Automatic selection based on system size.
+* "full_gpu": All wavefunction data resides on GPU (default behavior).
+* "paged": CPU storage with single k-point paging to GPU, reduces GPU memory at the cost of transfer overhead.
+
+Only relevant when device=gpu and basis_type=pw.)";
+        item.default_value = "";
+        item.read_value = [](const Input_Item& item, Parameter& para) {
+            para.input.device_memory_mode = item.get_size() > 0 ? strvalue : "";
+        };
+        sync_string(input.device_memory_mode);
+        item.check_value = [](const Input_Item& item, const Parameter& para) {
+            const std::vector<std::string> avail_list = {"", "full_gpu", "paged"};
+            if (std::find(avail_list.begin(), avail_list.end(), para.input.device_memory_mode) == avail_list.end())
+            {
+                ModuleBase::WARNING_QUIT("ReadInput", "device_memory_mode must be empty, full_gpu, or paged.");
+            }
+        };
+        this->add_item(item);
+    }
     {
         Input_Item item("precision");
         item.annotation = "the computing precision for ABACUS";

source/source_io/test/CMakeLists.txt

@@ -253,7 +253,7 @@ add_test(
 AddTest(
   TARGET MODULE_IO_read_wfc_nao_test
   LIBS parameter base device
-  SOURCES read_wfc_nao_test.cpp ../module_wf/read_wfc_nao.cpp ../../source_psi/psi.cpp ../../source_basis/module_ao/parallel_orbitals.cpp
+  SOURCES read_wfc_nao_test.cpp ../module_wf/read_wfc_nao.cpp ../../source_psi/psi.cpp ../../source_psi/psi_paging.cpp ../../source_basis/module_ao/parallel_orbitals.cpp
 )
 
 add_test(

source/source_psi/CMakeLists.txt

@@ -2,6 +2,7 @@ add_library(
     psi
     OBJECT
     psi.cpp
+    psi_paging.cpp
 )
 
 add_library(