Feat(psi): add PsiStorageMode and device_memory_mode for GPU paging

Port PAGED_GPU infrastructure from LTS-gpu-paged branch:
- Add PsiStorageMode enum (ALL_GPU, ALL_CPU, PAGED_GPU) to psi namespace
- Add paging members (psi_cpu_, psi_gpu_buffer_, current_k_gpu_, streams)
- Add paging methods: set_storage_mode, load_k_to_gpu, store_k_from_gpu,
  ensure_k_on_gpu, set_psi_cpu_external, get_cpu_pointer
- Modify resize() to allocate CPU storage + single-k GPU buffer in PAGED_GPU
  mode, reducing GPU memory for large k-point calculations
- Modify destructor to clean up paging resources
- Propagate storage_mode_ in copy constructors
- Add device_memory_mode INPUT parameter (''/full_gpu/paged)
- Create psi_paging.cpp with paging method implementations

Note: uses memcpy for CPU-side transfers (DEVICE_CPU), actual GPU-to-CPU
transfers via synchronize_memory_op will be added in a subsequent phase
with GPU FFT pipeline integration.

Author: dyzheng

source/source_io/module_parameter/input_parameter.h

@@ -87,6 +87,7 @@ struct Input_para
                                             ///< use mesh, which is used in QE.
     int nspin = 1;                          ///< LDA ; LSDA ; non-linear spin
     int gga_grad = 3;                       ///< GGA gradient method for noncollinear spin (nspin=4): 1=collinear approx, 2=projected div(h), 3=Scalmani-Frisch transform (default, most accurate)
+    std::string device_memory_mode = "";    ///< GPU memory mode: "" (auto), "full_gpu" (all on GPU), "paged" (CPU storage + k-point paging)
     int pw_diag_nmax = 50;
     double pw_diag_thr = 0.01;      ///< used in cg method
     bool diago_smooth_ethr = false; ///< smooth ethr for iter methods

source/source_io/module_parameter/read_input_item_system.cpp

@@ -701,6 +701,30 @@ 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.unit = "";
+        item.availability = "Only relevant when device=gpu with PW basis";
+        read_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_psi/CMakeLists.txt

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

source/source_psi/psi.cpp

@@ -1,12 +1,14 @@
 #include "psi.h"
 
 #include "source_base/global_variable.h"
+#include "source_base/memory.h"
 #include "source_base/module_device/device.h"
 #include "source_base/tool_quit.h"
 #include "source_io/module_parameter/parameter.h"
 
 #include <cassert>
 #include <complex>
+#include <cstring>
 #include <type_traits>
 
 namespace psi
@@ -41,6 +43,16 @@ Psi<T, Device>::~Psi()
     {
         delete_memory_op()(this->psi);
     }
+
+    if (psi_cpu_ != nullptr && psi_cpu_owned_)
+    {
+        delete[] psi_cpu_;
+    }
+    psi_cpu_ = nullptr;
+
+    psi_gpu_buffer_ = nullptr;
+
+    current_k_gpu_ = -1;
 }
 
 // Constructor 1:
@@ -160,12 +172,21 @@ Psi<T, Device>::Psi(const Psi& psi_in)
     this->current_k = psi_in.get_current_k();
     this->current_b = psi_in.get_current_b();
     this->k_first = psi_in.get_k_first();
-    // this function will copy psi_in.psi to this->psi no matter the device types of each other.
+    this->storage_mode_ = psi_in.get_storage_mode();
 
     this->resize(psi_in.get_nk(), psi_in.get_nbands(), psi_in.get_nbasis());
-    base_device::memory::synchronize_memory_op<T, Device, Device>()(this->psi,
-                                                                    psi_in.get_pointer() - psi_in.get_psi_bias(),
-                                                                    psi_in.size());
+
+    if (this->storage_mode_ == PsiStorageMode::PAGED_GPU && psi_cpu_ != nullptr)
+    {
+        const size_t total_size = static_cast<size_t>(this->nk) * this->nbands * this->nbasis;
+        std::memcpy(this->psi_cpu_, psi_in.get_pointer() - psi_in.get_psi_bias(), sizeof(T) * total_size);
+    }
+    else
+    {
+        base_device::memory::synchronize_memory_op<T, Device, Device>()(this->psi,
+                                                                        psi_in.get_pointer() - psi_in.get_psi_bias(),
+                                                                        psi_in.size());
+    }
     this->psi_bias = psi_in.get_psi_bias();
     this->current_nbasis = psi_in.get_current_nbas();
     this->psi_current = this->psi + psi_in.get_psi_bias();
@@ -184,16 +205,16 @@ Psi<T, Device>::Psi(const Psi<T_in, Device_in>& psi_in)
     this->current_k = psi_in.get_current_k();
     this->current_b = psi_in.get_current_b();
     this->k_first = psi_in.get_k_first();
-    // this function will copy psi_in.psi to this->psi no matter the device types of each other.
+    this->storage_mode_ = psi_in.get_storage_mode();
 
     this->resize(psi_in.get_nk(), psi_in.get_nbands(), psi_in.get_nbasis());
 
-    // Specifically, if the Device_in type is CPU and the Device type is GPU.
-    // Which means we need to initialize a GPU psi from a given CPU psi.
-    // We first malloc a memory in CPU, then cast the memory from T_in to T in CPU.
-    // Finally, synchronize the memory from CPU to GPU.
-    // This could help to reduce the peak memory usage of device.
-    if (std::is_same<Device, base_device::DEVICE_GPU>::value && std::is_same<Device_in, base_device::DEVICE_CPU>::value)
+    if (this->storage_mode_ == PsiStorageMode::PAGED_GPU && psi_cpu_ != nullptr)
+    {
+        const size_t total_size = static_cast<size_t>(this->nk) * this->nbands * this->nbasis;
+        std::memcpy(this->psi_cpu_, psi_in.get_pointer() - psi_in.get_psi_bias(), sizeof(T) * total_size);
+    }
+    else if (std::is_same<Device, base_device::DEVICE_GPU>::value && std::is_same<Device_in, base_device::DEVICE_CPU>::value)
     {
         auto* arr = (T*)malloc(sizeof(T) * psi_in.size());
         // cast the memory from T_in to T in CPU
@@ -250,21 +271,46 @@ Psi<T, Device>& Psi<T, Device>::operator=(const Psi<T, Device>& psi_in)
 }
 
 template <typename T, typename Device>
-void Psi<T, Device>::resize(const int nks_in, const int nbands_in, const int nbasis_in)
+void Psi<T, Device>::resize(const int nks_in, const int nbands_in, const int nbasis_in, const bool skip_psi_cpu_alloc)
 {
     assert(nks_in > 0 && nbands_in >= 0 && nbasis_in > 0);
 
-    // This function will delete the psi array first(if psi exist), then malloc a new memory for it.
-    resize_memory_op()(this->psi, nks_in * static_cast<std::size_t>(nbands_in) * nbasis_in, "no_record");
+    if (storage_mode_ == PsiStorageMode::PAGED_GPU)
+    {
+        const size_t total_size = static_cast<size_t>(nks_in) * nbands_in * nbasis_in;
+        if (!skip_psi_cpu_alloc)
+        {
+            if (psi_cpu_ != nullptr)
+            {
+                delete[] psi_cpu_;
+            }
+            psi_cpu_ = new T[total_size]();
+            ModuleBase::Memory::record("Psi::psi_cpu", sizeof(T) * total_size);
+            psi_cpu_owned_ = true;
+        }
+
+        const size_t k_size = static_cast<size_t>(nbands_in) * nbasis_in;
+        resize_memory_op()(this->psi, k_size, "no_record");
+#if defined(__CUDA) || defined(__ROCM)
+        ModuleBase::Memory::record_gpu("Psi::psi_gpu_k", sizeof(T) * k_size);
+#endif
 
-    // this->zero_out();
+        psi_gpu_buffer_ = this->psi;
+        current_k_gpu_ = -1;
+    }
+    else
+    {
+        resize_memory_op()(this->psi, nks_in * static_cast<std::size_t>(nbands_in) * nbasis_in, "no_record");
+#if defined(__CUDA) || defined(__ROCM)
+        ModuleBase::Memory::record_gpu("Psi_PW", sizeof(T) * nks_in * static_cast<std::size_t>(nbands_in) * nbasis_in);
+#endif
+    }
 
     this->nk = nks_in;
     this->nbands = nbands_in;
     this->nbasis = nbasis_in;
     this->current_nbasis = nbasis_in;
     this->psi_current = this->psi;
-    // GlobalV::ofs_device << "allocated xxx MB memory for psi" << std::endl;
 }
 
 template <typename T, typename Device>

source/source_psi/psi.h

@@ -10,6 +10,13 @@
 namespace psi
 {
 
+enum class PsiStorageMode
+{
+    ALL_GPU,
+    ALL_CPU,
+    PAGED_GPU
+};
+
 // structure for getting range of Psi
 // two display method: k index first or bands index first
 struct Range
@@ -78,7 +85,7 @@ class Psi
     Psi& operator=(const Psi& psi_in);
 
     // allocate psi for three dimensions
-    void resize(const int nks_in, const int nbands_in, const int nbasis_in);
+    void resize(const int nks_in, const int nbands_in, const int nbasis_in, const bool skip_psi_cpu_alloc = false);
 
     // get the pointer for the 1st index
     T* get_pointer() const;
@@ -139,6 +146,19 @@ class Psi
 
     int get_npol() const;
 
+    void set_storage_mode(PsiStorageMode mode);
+    PsiStorageMode get_storage_mode() const { return storage_mode_; }
+
+    void load_k_to_gpu(int ik);
+    void store_k_from_gpu(int ik);
+    void ensure_k_on_gpu(int ik);
+
+    void set_psi_cpu_external(T* ext_cpu_buf);
+
+    int get_current_k_gpu() const { return current_k_gpu_; }
+    T* get_cpu_pointer(int ik = 0);
+    const T* get_cpu_pointer(int ik = 0) const;
+
   private:
     T* psi = nullptr; // avoid using C++ STL
 
@@ -164,6 +184,18 @@ class Psi
 
     bool allocate_inside = true; ///< whether allocate psi inside Psi class
 
+    PsiStorageMode storage_mode_ = PsiStorageMode::ALL_GPU;
+    T* psi_cpu_ = nullptr;
+    bool psi_cpu_owned_ = true;
+    T* psi_gpu_buffer_ = nullptr;
+    T* psi_gpu_transfer_buffer_ = nullptr;
+    int current_k_gpu_ = -1;
+
+#if defined(__CUDA) || defined(__ROCM)
+    void* compute_stream_ = nullptr;
+    void* transfer_stream_ = nullptr;
+#endif
+
 #ifdef __DSP
     using delete_memory_op = base_device::memory::delete_memory_op_mt<T, Device>;
     using resize_memory_op = base_device::memory::resize_memory_op_mt<T, Device>;