Skip to content

Commit cfe3238

Browse files
author
dyzheng
committed
Feat(psi): add PsiStorageMode and device_memory_mode for GPU k-point paging
- Add PsiStorageMode enum (ALL_GPU, ALL_CPU, PAGED_GPU) to psi namespace - Add paging interface: load_k_to_gpu, store_k_from_gpu, ensure_k_on_gpu - Modify resize() to allocate CPU storage + single-k GPU buffer in PAGED_GPU mode - Create psi_paging.cpp with paging method implementations - Add device_memory_mode INPUT parameter (''/full_gpu/paged) - Add device_memory_mode reading in read_input_item_system.cpp
1 parent bb2cb4b commit cfe3238

6 files changed

Lines changed: 567 additions & 47 deletions

File tree

source/source_io/module_parameter/input_parameter.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ struct Input_para
7070
double min_dist_coef = 0.2; ///< allowed minimum distance between two atoms
7171

7272
std::string device = "cpu";
73+
std::string device_memory_mode = ""; ///< GPU memory mode: "" (auto), "full_gpu" (all on GPU), "paged" (CPU storage + k-point paging)
7374
std::string precision = "double";
7475
std::string gint_precision = "double";
7576
bool timer_enable_nvtx = false;

source/source_io/module_parameter/read_input_item_system.cpp

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -743,6 +743,28 @@ Available options are:
743743
};
744744
this->add_item(item);
745745
}
746+
{
747+
Input_Item item("device_memory_mode");
748+
item.annotation = "GPU memory strategy for wavefunction storage";
749+
item.category = "System variables";
750+
item.type = "String";
751+
item.description = R"(Controls GPU memory strategy for wavefunction storage.
752+
* "" (empty/auto): Automatic selection based on system size.
753+
* "full_gpu": All wavefunction data resides on GPU (default behavior).
754+
* "paged": CPU storage with single k-point paging to GPU, reduces GPU memory at the cost of transfer overhead.
755+
756+
Only relevant when device=gpu and basis_type=pw.)";
757+
item.default_value = "";
758+
read_sync_string(input.device_memory_mode);
759+
item.check_value = [](const Input_Item& item, const Parameter& para) {
760+
const std::vector<std::string> avail_list = {"", "full_gpu", "paged"};
761+
if (std::find(avail_list.begin(), avail_list.end(), para.input.device_memory_mode) == avail_list.end())
762+
{
763+
ModuleBase::WARNING_QUIT("ReadInput", "device_memory_mode must be empty, full_gpu, or paged.");
764+
}
765+
};
766+
this->add_item(item);
767+
}
746768
{
747769
Input_Item item("precision");
748770
item.annotation = "the computing precision for ABACUS";

source/source_psi/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ add_library(
22
psi
33
OBJECT
44
psi.cpp
5+
psi_paging.cpp
56
)
67

78
add_library(

0 commit comments

Comments
 (0)