diff --git a/docs/advanced/input_files/input-main.md b/docs/advanced/input_files/input-main.md index f04f903448..5a2002bfdc 100644 --- a/docs/advanced/input_files/input-main.md +++ b/docs/advanced/input_files/input-main.md @@ -30,6 +30,7 @@ - [kmesh\_type](#kmesh_type) - [min\_dist\_coef](#min_dist_coef) - [device](#device) + - [device\_memory\_mode](#device_memory_mode) - [precision](#precision) - [gint\_precision](#gint_precision) - [timer\_enable\_nvtx](#timer_enable_nvtx) @@ -799,6 +800,16 @@ > Note: ks_solver must also be set to the algorithms supported. lcao_in_pw currently does not support gpu. - **Default**: cpu +### device_memory_mode + +- **Type**: String +- **Description**: 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. + ### precision - **Type**: String diff --git a/docs/parameters.yaml b/docs/parameters.yaml index 005fad3076..33caddd7eb 100644 --- a/docs/parameters.yaml +++ b/docs/parameters.yaml @@ -277,6 +277,19 @@ parameters: default_value: cpu unit: "" availability: "" + - name: device_memory_mode + category: System variables + type: String + description: | + 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. + default_value: "" + unit: "" + availability: "" - name: precision category: System variables type: String diff --git a/source/source_estate/test/CMakeLists.txt b/source/source_estate/test/CMakeLists.txt index b9118b01a5..4ed5efec24 100644 --- a/source/source_estate/test/CMakeLists.txt +++ b/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 ) diff --git a/source/source_hsolver/hsolver_pw.cpp b/source/source_hsolver/hsolver_pw.cpp index 1fdcd131d6..649757eefb 100644 --- a/source/source_hsolver/hsolver_pw.cpp +++ b/source/source_hsolver/hsolver_pw.cpp @@ -109,13 +109,14 @@ void HSolverPW::solve(hamilt::Hamilt* 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::solve(hamilt::Hamilt* 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::solve(hamilt::Hamilt* 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::solve(hamilt::Hamilt* 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) { diff --git a/source/source_io/module_parameter/input_parameter.h b/source/source_io/module_parameter/input_parameter.h index 127e459548..891fdf0819 100644 --- a/source/source_io/module_parameter/input_parameter.h +++ b/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; diff --git a/source/source_io/module_parameter/read_input_item_system.cpp b/source/source_io/module_parameter/read_input_item_system.cpp index ddda68c657..752b1f3527 100644 --- a/source/source_io/module_parameter/read_input_item_system.cpp +++ b/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 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"; diff --git a/source/source_io/test/CMakeLists.txt b/source/source_io/test/CMakeLists.txt index 7044d7c8b9..203434c6d4 100644 --- a/source/source_io/test/CMakeLists.txt +++ b/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( diff --git a/source/source_psi/CMakeLists.txt b/source/source_psi/CMakeLists.txt index 8be3a2ba39..1bc4ba02a4 100644 --- a/source/source_psi/CMakeLists.txt +++ b/source/source_psi/CMakeLists.txt @@ -2,6 +2,7 @@ add_library( psi OBJECT psi.cpp + psi_paging.cpp ) add_library( diff --git a/source/source_psi/psi.cpp b/source/source_psi/psi.cpp index 8f64b0b8ea..842c2548fa 100644 --- a/source/source_psi/psi.cpp +++ b/source/source_psi/psi.cpp @@ -7,11 +7,37 @@ #include #include +#include #include namespace psi { +namespace detail +{ +// Helper for type-aware memory copy +template +struct TypeCopy +{ + static void copy(T* dst, const T_in* src, size_t size) + { + for (size_t i = 0; i < size; ++i) + { + dst[i] = static_cast(src[i]); + } + } +}; + +template +struct TypeCopy +{ + static void copy(T* dst, const T_in* src, size_t size) + { + std::memcpy(dst, src, sizeof(T) * size); + } +}; +} + Range::Range(const size_t range_in) { k_first = true; @@ -41,6 +67,16 @@ Psi::~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,15 +196,60 @@ Psi::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()(this->psi, - psi_in.get_pointer() - psi_in.get_psi_bias(), - psi_in.size()); - this->psi_bias = psi_in.get_psi_bias(); + + if (this->storage_mode_ == PsiStorageMode::PAGED_GPU && psi_cpu_ != nullptr) + { + const size_t total_size = static_cast(this->nk) * this->nbands * this->nbasis; + const T* src_ptr = psi_in.get_pointer() - psi_in.get_psi_bias(); + + if (psi_in.get_storage_mode() == PsiStorageMode::PAGED_GPU) + { + // Source may have PAGED_GPU flag set but psi_cpu_ not yet allocated + // (e.g. after set_storage_mode() in setup_psi_pw). In that case, + // copy from the source's main buffer (psi) into our psi_cpu_. + const T* cp_ptr = psi_in.get_cpu_pointer_safe(); + if (cp_ptr != nullptr) + { + std::memcpy(this->psi_cpu_, cp_ptr, sizeof(T) * total_size); + } + else + { + // Fallback: source data is in psi buffer (ALL_GPU layout) + std::memcpy(this->psi_cpu_, src_ptr, sizeof(T) * total_size); + } + } + else if (std::is_same::value) + { + std::memcpy(this->psi_cpu_, src_ptr, sizeof(T) * total_size); + } + else + { + base_device::memory::synchronize_memory_op()( + this->psi_cpu_, src_ptr, total_size); + } + this->current_k_gpu_ = -1; + this->psi_bias = 0; + this->psi_current = this->psi; + } + else + { + if (std::is_same::value) + { + std::memcpy(this->psi, psi_in.get_pointer() - psi_in.get_psi_bias(), + sizeof(T) * psi_in.size()); + } + else + { + base_device::memory::synchronize_memory_op()( + this->psi, psi_in.get_pointer() - psi_in.get_psi_bias(), psi_in.size()); + } + this->psi_bias = psi_in.get_psi_bias(); + this->psi_current = this->psi + psi_in.get_psi_bias(); + } this->current_nbasis = psi_in.get_current_nbas(); - this->psi_current = this->psi + psi_in.get_psi_bias(); } // Constructor 2-2: @@ -184,44 +265,83 @@ Psi::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()); - // 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::value && std::is_same::value) + if (this->storage_mode_ == PsiStorageMode::PAGED_GPU && psi_cpu_ != nullptr) + { + const size_t total_size = static_cast(this->nk) * this->nbands * this->nbasis; + + if (psi_in.get_storage_mode() == PsiStorageMode::PAGED_GPU) + { + // Source may have PAGED_GPU flag set but psi_cpu_ not yet allocated. + const T_in* cp_ptr = psi_in.get_cpu_pointer_safe(); + if (cp_ptr != nullptr) + { + detail::TypeCopy::value>::copy(this->psi_cpu_, cp_ptr, total_size); + } + else + { + // Fallback: source data is in main psi buffer (ALL_GPU layout) + auto* arr = (T*)malloc(sizeof(T) * total_size); + base_device::memory::cast_memory_op()( + arr, psi_in.get_pointer() - psi_in.get_psi_bias(), total_size); + std::memcpy(this->psi_cpu_, arr, sizeof(T) * total_size); + free(arr); + } + } + else + { + auto* arr = (T*)malloc(sizeof(T) * total_size); + base_device::memory::cast_memory_op()( + arr, psi_in.get_pointer() - psi_in.get_psi_bias(), total_size); + std::memcpy(this->psi_cpu_, arr, sizeof(T) * total_size); + free(arr); + } + this->current_k_gpu_ = -1; + this->psi_bias = 0; + this->psi_current = this->psi; + } + else if (std::is_same::value && std::is_same::value) { auto* arr = (T*)malloc(sizeof(T) * psi_in.size()); - // cast the memory from T_in to T in CPU base_device::memory::cast_memory_op()(arr, psi_in.get_pointer() - psi_in.get_psi_bias(), psi_in.size()); - // synchronize the memory from CPU to GPU base_device::memory::synchronize_memory_op()(this->psi, arr, psi_in.size()); free(arr); + this->psi_bias = psi_in.get_psi_bias(); + this->psi_current = this->psi + psi_in.get_psi_bias(); } else { base_device::memory::cast_memory_op()(this->psi, psi_in.get_pointer() - psi_in.get_psi_bias(), psi_in.size()); + this->psi_bias = psi_in.get_psi_bias(); + this->psi_current = this->psi + psi_in.get_psi_bias(); } - 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(); } template void Psi::set_all_psi(const T* another_pointer, const std::size_t size_in) { assert(size_in == this->size()); + if (this->storage_mode_ == PsiStorageMode::PAGED_GPU) + { + if (psi_cpu_ != nullptr) + { + const size_t total_size = static_cast(this->nk) * this->nbands * this->nbasis; + base_device::memory::synchronize_memory_op()( + this->psi_cpu_, another_pointer, total_size); + } + return; + } synchronize_memory_op()(this->psi, another_pointer, this->size()); } @@ -238,33 +358,68 @@ Psi& Psi::operator=(const Psi& psi_in) 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()(this->psi, - psi_in.psi, - psi_in.size()); - this->psi_bias = psi_in.get_psi_bias(); + + if (this->storage_mode_ == PsiStorageMode::PAGED_GPU) + { + // PAGED_GPU: copy the full CPU buffer, GPU buffer stays single-k + if (psi_cpu_ != nullptr && psi_in.psi_cpu_ != nullptr) + { + const size_t total_size = static_cast(this->nk) * this->nbands * this->nbasis; + std::memcpy(this->psi_cpu_, psi_in.psi_cpu_, sizeof(T) * total_size); + } + this->current_k_gpu_ = -1; + this->psi_bias = 0; + this->psi_current = this->psi; + } + else + { + base_device::memory::synchronize_memory_op()(this->psi, + psi_in.psi, + psi_in.size()); + this->psi_bias = psi_in.get_psi_bias(); + this->psi_current = this->psi + psi_in.get_psi_bias(); + } this->current_nbasis = psi_in.get_current_nbas(); - this->psi_current = this->psi + psi_in.get_psi_bias(); return *this; } template -void Psi::resize(const int nks_in, const int nbands_in, const int nbasis_in) +void Psi::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(nbands_in) * nbasis_in, "no_record"); - - // this->zero_out(); + if (storage_mode_ == PsiStorageMode::PAGED_GPU) + { + const size_t total_size = static_cast(nks_in) * nbands_in * nbasis_in; + if (!skip_psi_cpu_alloc) + { + if (psi_cpu_ != nullptr) + { + delete[] psi_cpu_; + } + psi_cpu_ = new T[total_size](); + psi_cpu_owned_ = true; + } + + const size_t k_size = static_cast(nbands_in) * nbasis_in; + resize_memory_op()(this->psi, k_size, "no_record"); + + psi_gpu_buffer_ = this->psi; + current_k_gpu_ = -1; + } + else + { + resize_memory_op()(this->psi, nks_in * static_cast(nbands_in) * nbasis_in, "no_record"); + } 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 @@ -377,6 +532,24 @@ void Psi::fix_k(const int ik) const { this->current_b = 0; } + + if (storage_mode_ == PsiStorageMode::PAGED_GPU) + { + this->psi_bias = 0; + if (this->current_k_gpu_ == ik) + { + // This k-point is loaded on GPU by load_k_to_gpu + this->psi_current = const_cast(this->psi); + } + else + { + // Read from CPU full buffer (psi_cpu_ has all k-points) + this->psi_current = const_cast(psi_cpu_ + + static_cast(ik) * this->nbands * this->nbasis); + } + return; + } + int base = this->current_b * this->nk * this->nbasis; if (ik >= this->nk) { @@ -396,6 +569,13 @@ void Psi::fix_b(const int ib) const assert(ib >= 0); this->current_b = ib; + if (storage_mode_ == PsiStorageMode::PAGED_GPU) + { + this->psi_bias = 0; + this->psi_current = const_cast(this->psi); + return; + } + if (!this->k_first) { this->current_k = 0; @@ -420,6 +600,14 @@ void Psi::fix_kb(const int ik, const int ib) const assert(ik >= 0 && ib >= 0); this->current_k = ik; this->current_b = ib; + + if (storage_mode_ == PsiStorageMode::PAGED_GPU) + { + this->psi_bias = 0; + this->psi_current = const_cast(this->psi); + return; + } + if (ik >= this->nk || ib >= this->nbands) { // fix to 0 this->psi_bias = 0; @@ -437,6 +625,18 @@ T& Psi::operator()(const int ikb1, const int ikb2, const int ibasis) { assert(ikb1 >= 0 && ikb2 >= 0 && ibasis >= 0); assert(this->k_first ? ikb1 < this->nk && ikb2 < this->nbands : ikb1 < this->nbands && ikb2 < this->nk); + + if (storage_mode_ == PsiStorageMode::PAGED_GPU) + { + // 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(psi_cpu_)[(ikb1 * this->nbands + ikb2) * this->nbasis + ibasis]; + } + return this->k_first ? this->psi[(ikb1 * this->nbands + ikb2) * this->nbasis + ibasis] : this->psi[(ikb1 * this->nk + ikb2) * this->nbasis + ibasis]; } @@ -485,38 +685,68 @@ const int& Psi::get_ngk(const int ik_in) const template void Psi::zero_out() { - // this->psi.assign(this->psi.size(), T(0)); - set_memory_op()(this->psi, 0, this->size()); + if (storage_mode_ == PsiStorageMode::PAGED_GPU) + { + // Zero full CPU buffer and single-k GPU buffer + if (psi_cpu_ != nullptr) + { + const size_t total_size = static_cast(this->nk) * this->nbands * this->nbasis; + std::memset(psi_cpu_, 0, sizeof(T) * total_size); + } + const size_t k_size = static_cast(this->nbands) * this->nbasis; + set_memory_op()(this->psi, 0, k_size); + } + else + { + set_memory_op()(this->psi, 0, this->size()); + } } template std::tuple Psi::to_range(const Range& range) const { - const int& i1 = range.index_1; - const int& r1 = range.range_1; - const int& r2 = range.range_2; + const size_t i1 = range.index_1; + const size_t r1 = range.range_1; + const size_t r2 = range.range_2; - if (range.k_first != this->k_first || r1 < 0 - || r2 < r1 - // || (range.k_first && (r2 >= this->nbands || i1 >= this->nk)) - // || (!range.k_first && (r2 >= this->nk || i1 >= this->nbands))) - || (range.k_first ? (i1 >= this->nk) : (i1 >= this->nbands)) // illegal index 1 - || (range.k_first ? (i1 > 0 && r2 >= this->nbands) : (i1 > 0 && r2 >= this->nk)) // illegal range of index 2 - || (range.k_first ? (i1 < 0 && r2 >= this->nk) : (i1 < 0 && r2 >= this->nbands))) // illegal range of index 1 + if (range.k_first != this->k_first || r2 < r1 + || (range.k_first ? (i1 >= static_cast(this->nk)) : (i1 >= static_cast(this->nbands))) + || (range.k_first ? (i1 > 0 && r2 >= static_cast(this->nbands)) : (i1 > 0 && r2 >= static_cast(this->nk)))) { return std::tuple(nullptr, 0); } - else if (i1 < 0) // [r1, r2] is the range of index1 with length m + else if (i1 > static_cast(std::numeric_limits::max())) { + // Sentinel case: i1 = (size_t)(-1) means a range of index1 (k-points) + // For PAGED_GPU, this is not supported since GPU buffer is single-k + if (storage_mode_ == PsiStorageMode::PAGED_GPU) + { + ModuleBase::WARNING_QUIT("Psi::to_range", + "In PAGED_GPU mode, multi-k-point range is not supported"); + } const T* p = &this->psi[r1 * (k_first ? this->nbands : this->nk) * this->nbasis]; - int m = (r2 - r1 + 1) * this->get_npol(); + int m = static_cast(r2 - r1 + 1) * this->get_npol(); return std::tuple(p, m); } - else // [r1, r2] is the range of index2 with length m + else // [r1, r2] is the range of band index for a specific k-point { - const T* p = &this->psi[(i1 * (k_first ? this->nbands : this->nk) + r1) * this->nbasis]; - int m = (r2 - r1 + 1) * this->get_npol(); - return std::tuple(p, m); + if (storage_mode_ == PsiStorageMode::PAGED_GPU) + { + if (k_first && static_cast(i1) != current_k) + { + ModuleBase::WARNING_QUIT("Psi::to_range", + "In PAGED_GPU mode, requested k-point must match current_k"); + } + const T* p = &this->psi[static_cast(r1) * static_cast(this->nbasis)]; + int m = static_cast(r2 - r1 + 1) * this->get_npol(); + return std::tuple(p, m); + } + else + { + const T* p = &this->psi[(i1 * (k_first ? this->nbands : this->nk) + r1) * this->nbasis]; + int m = static_cast(r2 - r1 + 1) * this->get_npol(); + return std::tuple(p, m); + } } } diff --git a/source/source_psi/psi.h b/source/source_psi/psi.h index 1aed4f46f6..7940713793 100644 --- a/source/source_psi/psi.h +++ b/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,20 @@ 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; + const T* get_cpu_pointer_safe(int ik = 0) const; + private: T* psi = nullptr; // avoid using C++ STL @@ -164,6 +185,13 @@ 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; + #ifdef __DSP using delete_memory_op = base_device::memory::delete_memory_op_mt; using resize_memory_op = base_device::memory::resize_memory_op_mt; diff --git a/source/source_psi/psi_paging.cpp b/source/source_psi/psi_paging.cpp new file mode 100644 index 0000000000..cbc1711197 --- /dev/null +++ b/source/source_psi/psi_paging.cpp @@ -0,0 +1,233 @@ +#include "psi.h" +#include "source_base/tool_quit.h" + +#include +#include + +namespace psi +{ + +namespace detail +{ +template +struct PagingTransfer +{ + static void load(T* dst, const T* src, size_t size); + static void store(T* dst, const T* src, size_t size); +}; + +template +struct PagingTransfer +{ + static void load(T* dst, const T* src, size_t size) + { + std::memcpy(dst, src, sizeof(T) * size); + } + static void store(T* dst, const T* src, size_t size) + { + std::memcpy(dst, src, sizeof(T) * size); + } +}; + +#if defined(__CUDA) || defined(__ROCM) +template +struct PagingTransfer +{ + static void load(T* dst, const T* src, size_t size) + { + base_device::memory::synchronize_memory_op()(dst, src, size); + } + static void store(T* dst, const T* src, size_t size) + { + base_device::memory::synchronize_memory_op()(dst, src, size); + } +}; +#endif +} + +template +void Psi::set_storage_mode(PsiStorageMode mode) +{ + if (mode == storage_mode_) + { + return; + } + + if (this->psi != nullptr) + { + if ((storage_mode_ == PsiStorageMode::ALL_GPU && mode == PsiStorageMode::PAGED_GPU) + || (storage_mode_ == PsiStorageMode::PAGED_GPU && mode == PsiStorageMode::ALL_GPU)) + { + ModuleBase::WARNING("Psi::set_storage_mode", + "Setting storage mode flag on already-allocated Psi for copy construction propagation."); + } + else + { + ModuleBase::WARNING_QUIT("Psi::set_storage_mode", + "Cannot change storage mode after allocation"); + } + } + + storage_mode_ = mode; +} + +template +T* Psi::get_cpu_pointer(int ik) +{ + if (storage_mode_ == PsiStorageMode::PAGED_GPU) + { + if (ik < 0 || ik >= this->nk) + { + ModuleBase::WARNING_QUIT("Psi::get_cpu_pointer", "Invalid k-point index"); + } + if (psi_cpu_ == nullptr) + { + ModuleBase::WARNING_QUIT("Psi::get_cpu_pointer", + "CPU buffer not allocated in PAGED_GPU mode"); + } + return psi_cpu_ + static_cast(ik) * this->nbands * this->nbasis; + } + else + { + return this->psi + static_cast(ik) * this->nbands * this->nbasis; + } +} + +template +const T* Psi::get_cpu_pointer(int ik) const +{ + if (storage_mode_ == PsiStorageMode::PAGED_GPU) + { + if (ik < 0 || ik >= this->nk) + { + ModuleBase::WARNING_QUIT("Psi::get_cpu_pointer", "Invalid k-point index"); + } + if (psi_cpu_ == nullptr) + { + ModuleBase::WARNING_QUIT("Psi::get_cpu_pointer", + "CPU buffer not allocated in PAGED_GPU mode"); + } + return psi_cpu_ + static_cast(ik) * this->nbands * this->nbasis; + } + else + { + return this->psi + static_cast(ik) * this->nbands * this->nbasis; + } +} + +template +const T* Psi::get_cpu_pointer_safe(int ik) const +{ + if (storage_mode_ == PsiStorageMode::PAGED_GPU && psi_cpu_ != nullptr) + { + if (ik < 0 || ik >= this->nk) return nullptr; + return psi_cpu_ + static_cast(ik) * this->nbands * this->nbasis; + } + return nullptr; +} + +template +void Psi::load_k_to_gpu(int ik) +{ + if (storage_mode_ != PsiStorageMode::PAGED_GPU) + { + return; + } + + if (ik < 0 || ik >= this->nk) + { + ModuleBase::WARNING_QUIT("Psi::load_k_to_gpu", "Invalid k-point index"); + } + + if (psi_cpu_ == nullptr) + { + ModuleBase::WARNING_QUIT("Psi::load_k_to_gpu", "CPU buffer not allocated"); + } + + const size_t k_size = static_cast(this->nbands) * this->nbasis; + const T* src = psi_cpu_ + static_cast(ik) * this->nbands * this->nbasis; + + detail::PagingTransfer::load(this->psi, src, k_size); + + current_k_gpu_ = ik; + this->psi_current = this->psi; +} + +template +void Psi::store_k_from_gpu(int ik) +{ + if (storage_mode_ != PsiStorageMode::PAGED_GPU) + { + return; + } + + if (ik < 0 || ik >= this->nk) + { + ModuleBase::WARNING_QUIT("Psi::store_k_from_gpu", "Invalid k-point index"); + } + + if (psi_cpu_ == nullptr) + { + ModuleBase::WARNING_QUIT("Psi::store_k_from_gpu", "CPU buffer not allocated"); + } + + const size_t k_size = static_cast(this->nbands) * this->nbasis; + T* dst = psi_cpu_ + static_cast(ik) * this->nbands * this->nbasis; + + detail::PagingTransfer::store(dst, this->psi, k_size); +} + +template +void Psi::ensure_k_on_gpu(int ik) +{ + if (storage_mode_ != PsiStorageMode::PAGED_GPU) + { + return; + } + + if (current_k_gpu_ != ik) + { + load_k_to_gpu(ik); + } +} + +template +void Psi::set_psi_cpu_external(T* ext_cpu_buf) +{ + if (storage_mode_ != PsiStorageMode::PAGED_GPU) + { + ModuleBase::WARNING_QUIT("Psi::set_psi_cpu_external", + "Only valid in PAGED_GPU mode"); + } + if (psi_cpu_ != nullptr && psi_cpu_owned_) + { + delete[] psi_cpu_; + } + psi_cpu_ = ext_cpu_buf; + psi_cpu_owned_ = false; +} + +#define INSTANTIATE_PAGING_METHODS(T, Device) \ + template void Psi::set_storage_mode(PsiStorageMode); \ + template T* Psi::get_cpu_pointer(int); \ + template const T* Psi::get_cpu_pointer(int) const; \ + template const T* Psi::get_cpu_pointer_safe(int) const; \ + template void Psi::load_k_to_gpu(int); \ + template void Psi::store_k_from_gpu(int); \ + template void Psi::ensure_k_on_gpu(int); \ + template void Psi::set_psi_cpu_external(T*); + +INSTANTIATE_PAGING_METHODS(float, base_device::DEVICE_CPU) +INSTANTIATE_PAGING_METHODS(std::complex, base_device::DEVICE_CPU) +INSTANTIATE_PAGING_METHODS(double, base_device::DEVICE_CPU) +INSTANTIATE_PAGING_METHODS(std::complex, base_device::DEVICE_CPU) +#if ((defined __CUDA) || (defined __ROCM)) +INSTANTIATE_PAGING_METHODS(float, base_device::DEVICE_GPU) +INSTANTIATE_PAGING_METHODS(std::complex, base_device::DEVICE_GPU) +INSTANTIATE_PAGING_METHODS(double, base_device::DEVICE_GPU) +INSTANTIATE_PAGING_METHODS(std::complex, base_device::DEVICE_GPU) +#endif + +#undef INSTANTIATE_PAGING_METHODS + +} // namespace psi diff --git a/source/source_psi/psi_prepare.cpp b/source/source_psi/psi_prepare.cpp index adc10eeb3f..495c6ee065 100644 --- a/source/source_psi/psi_prepare.cpp +++ b/source/source_psi/psi_prepare.cpp @@ -176,7 +176,16 @@ void PSIPrepare::initialize_psi(Psi>* psi, this->psi_initer->init_psig(psi_cpu->get_pointer(), ik); if (psi_device->get_pointer() != psi_cpu->get_pointer()) { - syncmem_h2d_op()(psi_device->get_pointer(), psi_cpu->get_pointer(), nbands_start * nbasis); + if (kspw_psi->get_storage_mode() == psi::PsiStorageMode::PAGED_GPU) + { + T* dst_cpu = kspw_psi->get_cpu_pointer(ik); + std::memcpy(dst_cpu, psi_cpu->get_pointer(), sizeof(T) * nbands_l * nbasis); + kspw_psi->load_k_to_gpu(ik); + } + else + { + syncmem_h2d_op()(psi_device->get_pointer(), psi_cpu->get_pointer(), nbands_start * nbasis); + } } @@ -208,7 +217,13 @@ void PSIPrepare::initialize_psi(Psi>* psi, { if (psi_device->get_pointer() != kspw_psi->get_pointer()) { - syncmem_complex_op()(kspw_psi->get_pointer(), psi_device->get_pointer(), nbands_l * nbasis); + if (kspw_psi->get_storage_mode() == psi::PsiStorageMode::PAGED_GPU) + { + } + else + { + syncmem_complex_op()(kspw_psi->get_pointer(), psi_device->get_pointer(), nbands_l * nbasis); + } } } } diff --git a/source/source_psi/setup_psi_pw.cpp b/source/source_psi/setup_psi_pw.cpp index f5bc240292..a3a89ab23f 100644 --- a/source/source_psi/setup_psi_pw.cpp +++ b/source/source_psi/setup_psi_pw.cpp @@ -40,7 +40,33 @@ void Setup_Psi_pw::before_runner_impl( } if (inp.device == "gpu" || inp.precision == "single") { + const int nks = kv.get_nks(); + psi::PsiStorageMode target_mode = psi::PsiStorageMode::ALL_GPU; + if (inp.device_memory_mode == "paged") + { + target_mode = psi::PsiStorageMode::PAGED_GPU; + } + else if (inp.device_memory_mode == "full_gpu") + { + target_mode = psi::PsiStorageMode::ALL_GPU; + } + else if (nks > 10) + { + target_mode = psi::PsiStorageMode::PAGED_GPU; + } + if (target_mode != psi::PsiStorageMode::ALL_GPU) + { + this->psi_cpu->set_storage_mode(target_mode); + } + const char* mode_str = (target_mode == psi::PsiStorageMode::PAGED_GPU) ? "PAGED_GPU" : "ALL_GPU"; + GlobalV::ofs_running << " GPU memory mode for Psi: " << mode_str + << " (nks=" << nks << ", device_memory_mode=\"" + << inp.device_memory_mode << "\")" << std::endl; this->psi_t = static_cast(new psi::Psi(this->psi_cpu[0])); + if (target_mode != psi::PsiStorageMode::ALL_GPU) + { + this->psi_cpu->set_storage_mode(psi::PsiStorageMode::ALL_GPU); + } } else { this->psi_t = static_cast(reinterpret_cast*>(this->psi_cpu)); } @@ -178,9 +204,19 @@ template void Setup_Psi_pw::copy_d2h_impl() { auto* psi_t = this->get_psi_t(); - this->castmem_d2h_impl(this->psi_cpu[0].get_pointer() - this->psi_cpu[0].get_psi_bias(), - psi_t->get_pointer() - psi_t->get_psi_bias(), - this->psi_cpu[0].size()); + if (psi_t->get_storage_mode() == psi::PsiStorageMode::PAGED_GPU) + { + // PAGED_GPU: full data is on CPU in psi_cpu_, just memcpy + const size_t total_size = sizeof(T) * psi_t->get_nk() * psi_t->get_nbands() * psi_t->get_nbasis(); + std::memcpy(this->psi_cpu[0].get_pointer() - this->psi_cpu[0].get_psi_bias(), + psi_t->get_cpu_pointer(), total_size); + } + else + { + this->castmem_d2h_impl(this->psi_cpu[0].get_pointer() - this->psi_cpu[0].get_psi_bias(), + psi_t->get_pointer() - psi_t->get_psi_bias(), + this->psi_cpu[0].size()); + } } void Setup_Psi_pw::copy_d2h() diff --git a/source/source_psi/test/CMakeLists.txt b/source/source_psi/test/CMakeLists.txt index 63af5799e1..650ddf333c 100644 --- a/source/source_psi/test/CMakeLists.txt +++ b/source/source_psi/test/CMakeLists.txt @@ -3,7 +3,8 @@ AddTest( LIBS parameter base device SOURCES psi_test.cpp - ../psi.cpp + ../psi.cpp + ../psi_paging.cpp ) if(ENABLE_LCAO) diff --git a/source/source_pw/module_pwdft/forces_nl.cpp b/source/source_pw/module_pwdft/forces_nl.cpp index 599121591c..bb5bfda18b 100644 --- a/source/source_pw/module_pwdft/forces_nl.cpp +++ b/source/source_pw/module_pwdft/forces_nl.cpp @@ -37,6 +37,9 @@ void Forces::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, 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) diff --git a/source/source_pw/module_pwdft/forces_onsite.cpp b/source/source_pw/module_pwdft/forces_onsite.cpp index 255d33f943..71a480ac2b 100644 --- a/source/source_pw/module_pwdft/forces_onsite.cpp +++ b/source/source_pw/module_pwdft/forces_onsite.cpp @@ -32,6 +32,11 @@ void Forces::cal_force_onsite(ModuleBase::matrix& force_onsite, const int nks = wfc_basis->nks; for (int ik = 0; ik < nks; ik++) { + // In PAGED_GPU mode only one k-point resides on the GPU; make sure the + // wavefunction of this k-point is loaded before the onsite projector + // reads psi_ (otherwise psi_(ik,0,0) returns a host pointer that the + // GPU gemm would dereference as a device pointer -> illegal access). + const_cast, Device>*>(psi_in)->ensure_k_on_gpu(ik); int nbands_occ = wg.nc; while (wg(ik, nbands_occ - 1) == 0.0) { diff --git a/source/source_pw/module_pwdft/onsite_proj.cpp b/source/source_pw/module_pwdft/onsite_proj.cpp index d4af4a7c2f..3a9d53d61d 100644 --- a/source/source_pw/module_pwdft/onsite_proj.cpp +++ b/source/source_pw/module_pwdft/onsite_proj.cpp @@ -539,6 +539,7 @@ void projectors::OnsiteProjector::cal_occupations( const int nbands = psi_in->get_nbands(); for(int ik = 0; ik < psi_in->get_nk(); ik++) { + const_cast, Device>*>(psi_in)->ensure_k_on_gpu(ik); psi_in->fix_k(ik); if(ik != 0) { diff --git a/source/source_pw/module_pwdft/stress_kin.cpp b/source/source_pw/module_pwdft/stress_kin.cpp index ba30034d87..d1c003e064 100644 --- a/source/source_pw/module_pwdft/stress_kin.cpp +++ b/source/source_pw/module_pwdft/stress_kin.cpp @@ -18,9 +18,12 @@ void Stress_Func::stress_kin(ModuleBase::matrix& sigma, this->ucell = &ucell_in; - hamilt::FS_Kin_tools kin_tool(*this->ucell, p_kv, wfc_basis, wg); + hamilt::FS_Kin_tools 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, Device>*>(psi_in)->ensure_k_on_gpu(ik); + int nbands_occ = wg.nc; while (wg(ik, nbands_occ - 1) == 0.0) { diff --git a/source/source_pw/module_pwdft/stress_nl.cpp b/source/source_pw/module_pwdft/stress_nl.cpp index 7b9b84925c..9cb6adea6b 100644 --- a/source/source_pw/module_pwdft/stress_nl.cpp +++ b/source/source_pw/module_pwdft/stress_nl.cpp @@ -39,6 +39,9 @@ void Stress_Func::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, 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) diff --git a/source/source_pw/module_pwdft/stress_onsite.cpp b/source/source_pw/module_pwdft/stress_onsite.cpp index 8223be49ef..93c60eb0c8 100644 --- a/source/source_pw/module_pwdft/stress_onsite.cpp +++ b/source/source_pw/module_pwdft/stress_onsite.cpp @@ -60,6 +60,12 @@ void Stress_Func::stress_onsite( // Loop over all k-points for (int ik = 0; ik < nks; ik++) { + // In PAGED_GPU mode only one k-point resides on the GPU; load this + // k-point's wavefunction before the onsite projector reads psi_ + // (otherwise psi_(ik,0,0) returns a host pointer used as a device + // pointer by the GPU gemm -> illegal memory access). + const_cast, Device>*>( + static_cast, Device>*>(psi_in))->ensure_k_on_gpu(ik); // Determine number of occupied bands (skip zero weights) int nbands_occ = wg.nc;