From 1f41ef5682ac0b4d7eef9a349c5b9d356f91245d Mon Sep 17 00:00:00 2001 From: denghuilu Date: Sun, 6 Nov 2022 00:36:22 +0800 Subject: [PATCH 1/5] Enable GPU support for CG method --- source/module_esolver/esolver_ks.h | 2 +- source/module_esolver/esolver_ks_pw.cpp | 6 +- source/module_hamilt/hamilt.h | 6 +- source/module_hamilt/hamilt_lcao.h | 2 +- source/module_hamilt/hamilt_pw.cpp | 99 +++++- source/module_hamilt/hamilt_pw.h | 8 +- source/module_hamilt/ks_pw/ekinetic_pw.cpp | 24 ++ source/module_hamilt/ks_pw/ekinetic_pw.h | 11 +- source/module_hamilt/ks_pw/meta_pw.cpp | 19 ++ source/module_hamilt/ks_pw/meta_pw.h | 9 + source/module_hamilt/ks_pw/nonlocal_pw.cpp | 27 ++ source/module_hamilt/ks_pw/nonlocal_pw.h | 7 + source/module_hamilt/ks_pw/operator_pw.h | 2 + source/module_hamilt/ks_pw/veff_pw.cpp | 31 ++ source/module_hamilt/ks_pw/veff_pw.h | 10 + source/module_hamilt/operator.cpp | 4 +- source/module_hamilt/operator.h | 11 +- source/module_hsolver/diagh.h | 6 +- source/module_hsolver/diago_blas.cpp | 4 +- source/module_hsolver/diago_blas.h | 6 +- source/module_hsolver/diago_cg.cpp | 311 ++++++------------ source/module_hsolver/diago_cg.h | 28 +- source/module_hsolver/diago_david.cpp | 10 +- source/module_hsolver/diago_david.h | 8 +- source/module_hsolver/diago_elpa.cpp | 4 +- source/module_hsolver/diago_elpa.h | 6 +- source/module_hsolver/diago_iter_assist.cpp | 8 +- source/module_hsolver/diago_iter_assist.h | 7 +- source/module_hsolver/diago_lapack.cpp | 4 +- source/module_hsolver/diago_lapack.h | 6 +- source/module_hsolver/hsolver.h | 9 +- source/module_hsolver/hsolver_lcao.cpp | 10 +- source/module_hsolver/hsolver_lcao.h | 10 +- source/module_hsolver/hsolver_pw.cpp | 50 ++- source/module_hsolver/hsolver_pw.h | 6 +- source/module_hsolver/hsolver_pw_sdft.cpp | 2 +- source/module_hsolver/hsolver_pw_sdft.h | 2 +- source/module_hsolver/src/cuda/math_kernel.cu | 11 +- source/module_hsolver/test/diago_cg_test.cpp | 4 +- .../module_hsolver/test/diago_david_test.cpp | 4 +- .../module_hsolver/test/diago_lcao_test.cpp | 4 +- source/module_hsolver/test/diago_mock.h | 8 +- source/module_psi/psi.cpp | 116 ++++--- source/module_psi/psi.h | 2 + source/src_lcao/ELEC_evolve.cpp | 2 +- source/src_lcao/ELEC_evolve.h | 2 +- source/src_lcao/LCAO_evolve.cpp | 6 +- source/src_lcao/LCAO_evolve.h | 6 +- source/src_pw/sto_iter.cpp | 2 +- source/src_pw/sto_iter.h | 2 +- source/src_pw/wavefunc.cpp | 2 +- source/src_pw/wavefunc.h | 2 +- 52 files changed, 561 insertions(+), 387 deletions(-) diff --git a/source/module_esolver/esolver_ks.h b/source/module_esolver/esolver_ks.h index 36687a15975..9cba1cc8996 100644 --- a/source/module_esolver/esolver_ks.h +++ b/source/module_esolver/esolver_ks.h @@ -68,7 +68,7 @@ namespace ModuleESolver hsolver::HSolver* phsol = nullptr; elecstate::ElecState* pelec = nullptr; - hamilt::Hamilt* p_hamilt = nullptr; + hamilt::Hamilt* p_hamilt = nullptr; ModulePW::PW_Basis_K* pw_wfc = nullptr; Charge_Extra CE; diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp index 060db7c0e6a..a454221d35c 100644 --- a/source/module_esolver/esolver_ks_pw.cpp +++ b/source/module_esolver/esolver_ks_pw.cpp @@ -59,7 +59,7 @@ namespace ModuleESolver //delete Hamilt if(this->p_hamilt != nullptr) { - delete (hamilt::HamiltPW*)this->p_hamilt; + delete (hamilt::HamiltPW*)this->p_hamilt; this->p_hamilt = nullptr; } } @@ -198,13 +198,13 @@ namespace ModuleESolver //delete Hamilt if not first scf if(this->p_hamilt != nullptr) { - delete (hamilt::HamiltPW*)this->p_hamilt; + delete (hamilt::HamiltPW*)this->p_hamilt; this->p_hamilt = nullptr; } //allocate HamiltPW if(this->p_hamilt == nullptr) { - this->p_hamilt = new hamilt::HamiltPW(); + this->p_hamilt = new hamilt::HamiltPW(); } //---------------------------------------------------------- diff --git a/source/module_hamilt/hamilt.h b/source/module_hamilt/hamilt.h index ea6badd5c61..bc07fbfef4e 100644 --- a/source/module_hamilt/hamilt.h +++ b/source/module_hamilt/hamilt.h @@ -11,6 +11,7 @@ namespace hamilt { +template class Hamilt { public: @@ -32,8 +33,9 @@ class Hamilt int non_first_scf=0; // first node operator, add operations from each operators - Operator>* ops = nullptr; - Operator* opsd = nullptr; + Operator, Device>* ops = nullptr; + Operator* opsd = nullptr; + }; } // namespace hamilt diff --git a/source/module_hamilt/hamilt_lcao.h b/source/module_hamilt/hamilt_lcao.h index 6e68ef2eda2..44aad02c1d2 100644 --- a/source/module_hamilt/hamilt_lcao.h +++ b/source/module_hamilt/hamilt_lcao.h @@ -15,7 +15,7 @@ namespace hamilt // template first for type of k space H matrix elements // template second for type of temporary matrix, gamma_only fix-gamma-matrix + S-gamma, multi-k fix-Real + S-Real -template class HamiltLCAO : public Hamilt +template class HamiltLCAO : public Hamilt { public: HamiltLCAO( diff --git a/source/module_hamilt/hamilt_pw.cpp b/source/module_hamilt/hamilt_pw.cpp index 89299c6d8ab..46748df35d7 100644 --- a/source/module_hamilt/hamilt_pw.cpp +++ b/source/module_hamilt/hamilt_pw.cpp @@ -3,8 +3,6 @@ #include "module_base/blas_connector.h" #include "module_base/global_function.h" #include "module_base/global_variable.h" -#include "module_base/timer.h" -#include "src_parallel/parallel_reduce.h" #include "src_pw/global.h" #include "ks_pw/veff_pw.h" @@ -15,7 +13,8 @@ namespace hamilt { -HamiltPW::HamiltPW() +template +HamiltPW::HamiltPW() { this->classname = "HamiltPW"; const double tpiba2 = GlobalC::ucell.tpiba2; @@ -26,7 +25,7 @@ HamiltPW::HamiltPW() if (GlobalV::T_IN_H) { // Operator* ekinetic = new Ekinetic> - Operator>* ekinetic = new Ekinetic>( + Operator, Device>* ekinetic = new Ekinetic>( tpiba2, gk2, GlobalC::wfcpw->nks, @@ -43,7 +42,7 @@ HamiltPW::HamiltPW() } if (GlobalV::VL_IN_H) { - Operator>* veff = new Veff>( + Operator, Device>* veff = new Veff>( isk, &(GlobalC::pot.vr_eff), GlobalC::wfcpw @@ -59,7 +58,7 @@ HamiltPW::HamiltPW() } if (GlobalV::VNL_IN_H) { - Operator>* nonlocal = new Nonlocal>( + Operator, Device>* nonlocal = new Nonlocal>( isk, &GlobalC::ppcell, &GlobalC::ucell @@ -73,7 +72,7 @@ HamiltPW::HamiltPW() this->ops->add(nonlocal); } } - Operator>* meta = new Meta>( + Operator, Device>* meta = new Meta>( tpiba, isk, &GlobalC::pot.vofk, @@ -89,7 +88,8 @@ HamiltPW::HamiltPW() } } -HamiltPW::~HamiltPW() +template +HamiltPW::~HamiltPW() { if(this->ops!= nullptr) { @@ -97,24 +97,93 @@ HamiltPW::~HamiltPW() } } -void HamiltPW::updateHk(const int ik) +template +void HamiltPW::updateHk(const int ik) { ModuleBase::TITLE("HamiltPW","updateHk"); - this->ops->init(ik); - - return; + ModuleBase::TITLE("HamiltPW","updateHk"); } -void HamiltPW::sPsi +template +void HamiltPW::sPsi ( const std::complex *psi, std::complex *spsi, size_t size ) const { - ModuleBase::GlobalFunc::COPYARRAY(psi, spsi, size); - return; + // ModuleBase::GlobalFunc::COPYARRAY(psi, spsi, size); + // denghui replaced at 2022.11.04 + syncmem_complex_op()(this->ctx, this->ctx, spsi, psi, size); +} + +template +template +HamiltPW::HamiltPW(const HamiltPW *hamilt) +{ + this->classname = hamilt->classname; + OperatorPW, Device_in> * node = + reinterpret_cast, Device_in> *>(hamilt->ops); + + while(node != nullptr) { + if (node->classname == "Ekinetic") { + Operator, Device>* ekinetic = + new Ekinetic>( + reinterpret_cast>*>(node)); + if(this->ops == nullptr) { + this->ops = ekinetic; + } + else { + this->ops->add(ekinetic); + } + // this->ops = reinterpret_cast, Device>*>(node); + } + else if (node->classname == "Nonlocal") { + Operator, Device>* nonlocal = + new Nonlocal>( + reinterpret_cast>*>(node)); + if(this->ops == nullptr) { + this->ops = nonlocal; + } + else { + this->ops->add(nonlocal); + } + } + else if (node->classname == "Veff") { + Operator, Device>* veff = + new Veff>( + reinterpret_cast>*>(node)); + if(this->ops == nullptr) { + this->ops = veff; + } + else { + this->ops->add(veff); + } + } + else if (node->classname == "Meta") { + Operator, Device>* meta = + new Meta>( + reinterpret_cast>*>(node)); + if(this->ops == nullptr) { + this->ops = meta; + } + else { + this->ops->add(meta); + } + } + else { + ModuleBase::WARNING_QUIT("HamiltPW", "Unrecognized Operator type!"); + } + node = reinterpret_cast, Device_in> *>(node->next_op); + } } +template class HamiltPW; +#if ((defined __CUDA) || (defined __ROCM)) +template class HamiltPW; +template HamiltPW::HamiltPW(const HamiltPW *hamilt); +template HamiltPW::HamiltPW(const HamiltPW *hamilt); +#endif + } // namespace hamilt \ No newline at end of file diff --git a/source/module_hamilt/hamilt_pw.h b/source/module_hamilt/hamilt_pw.h index 25a1c7c5179..efe301de5dd 100644 --- a/source/module_hamilt/hamilt_pw.h +++ b/source/module_hamilt/hamilt_pw.h @@ -6,10 +6,13 @@ namespace hamilt { -class HamiltPW : public Hamilt +template +class HamiltPW : public Hamilt { public: HamiltPW(); + template + explicit HamiltPW(const HamiltPW* hamilt); ~HamiltPW(); // for target K point, update consequence of hPsi() and matrix() @@ -19,6 +22,9 @@ class HamiltPW : public Hamilt virtual void sPsi(const std::complex *psi_in, std::complex *spsi, const size_t size) const override; private: + + Device *ctx = {}; + using syncmem_complex_op = psi::memory::synchronize_memory_op, Device, Device>; }; } // namespace hamilt diff --git a/source/module_hamilt/ks_pw/ekinetic_pw.cpp b/source/module_hamilt/ks_pw/ekinetic_pw.cpp index 6fff645c977..743aff29bc9 100644 --- a/source/module_hamilt/ks_pw/ekinetic_pw.cpp +++ b/source/module_hamilt/ks_pw/ekinetic_pw.cpp @@ -13,6 +13,7 @@ Ekinetic>::Ekinetic( const int gk2_row, const int gk2_col) { + this->classname = "Ekinetic"; this->cal_type = pw_ekinetic; this->tpiba2 = tpiba2_in; this->gk2_row = gk2_row; @@ -61,9 +62,32 @@ void Ekinetic>::act( ModuleBase::timer::tick("Operator", "EkineticPW"); } +// copy construct added by denghui at 20221105 +template +template +hamilt::Ekinetic>::Ekinetic(const Ekinetic> *ekinetic) { + this->classname = "Ekinetic"; + this->cal_type = pw_ekinetic; + this->ik = ekinetic->get_ik(); + this->tpiba2 = ekinetic->get_tpiba2(); + this->gk2_row = ekinetic->get_gk2_row(); + this->gk2_col = ekinetic->get_gk2_col(); + resize_memory_op()(this->ctx, this->gk2, this->gk2_row * this->gk2_col); + psi::memory::synchronize_memory_op()( + this->ctx, ekinetic->get_ctx(), + this->gk2, ekinetic->get_gk2(), + this->gk2_row * this->gk2_col); + + if( this->tpiba2 < 1e-10 || this->gk2 == nullptr) { + ModuleBase::WARNING_QUIT("EkineticPW", "Copy Constuctor of Operator::EkineticPW is failed, please check your code!"); + } +} + namespace hamilt{ template class Ekinetic>; #if ((defined __CUDA) || (defined __ROCM)) template class Ekinetic>; +template Ekinetic>::Ekinetic(const Ekinetic> *ekinetic); +template Ekinetic>::Ekinetic(const Ekinetic> *ekinetic); #endif } // namespace hamilt \ No newline at end of file diff --git a/source/module_hamilt/ks_pw/ekinetic_pw.h b/source/module_hamilt/ks_pw/ekinetic_pw.h index 12a5b5e341c..e83332a0dfa 100644 --- a/source/module_hamilt/ks_pw/ekinetic_pw.h +++ b/source/module_hamilt/ks_pw/ekinetic_pw.h @@ -29,6 +29,9 @@ class Ekinetic> : public OperatorPW const int gk2_row, const int gk2_col); + template + explicit Ekinetic(const Ekinetic>* ekinetic); + virtual ~Ekinetic(); virtual void act( @@ -37,6 +40,13 @@ class Ekinetic> : public OperatorPW const std::complex* tmpsi_in, std::complex* tmhpsi)const override; + // denghuilu added for copy construct at 20221105 + int get_gk2_row() const {return this->gk2_row;} + int get_gk2_col() const {return this->gk2_col;} + FPTYPE get_tpiba2() const {return this->tpiba2;} + const FPTYPE* get_gk2() const {return this->gk2;} + Device* get_ctx() const {return this->ctx;} + private: mutable int max_npw = 0; @@ -44,7 +54,6 @@ class Ekinetic> : public OperatorPW mutable int npol = 0; FPTYPE tpiba2 = 0.0; - #if ((defined __CUDA) || (defined __ROCM)) FPTYPE* gk2 = nullptr; #else diff --git a/source/module_hamilt/ks_pw/meta_pw.cpp b/source/module_hamilt/ks_pw/meta_pw.cpp index 459e290ce64..43914f95708 100644 --- a/source/module_hamilt/ks_pw/meta_pw.cpp +++ b/source/module_hamilt/ks_pw/meta_pw.cpp @@ -18,6 +18,7 @@ Meta>::Meta( ModulePW::PW_Basis_K* wfcpw_in ) { + this->classname = "Meta"; this->cal_type = pw_meta; this->isk = isk_in; this->tpiba = tpiba_in; @@ -85,9 +86,27 @@ void Meta>::act( ModuleBase::timer::tick("Operator", "MetaPW"); } +template +template +hamilt::Meta>::Meta(const Meta> *meta) { + this->classname = "Meta"; + this->cal_type = pw_meta; + this->ik = meta->get_ik(); + this->isk = meta->get_isk(); + this->tpiba = meta->get_tpiba(); + this->vk = meta->get_vk(); + this->wfcpw = meta->get_wfcpw(); + if(this->isk == nullptr || this->tpiba < 1e-10 || this->vk == nullptr || this->wfcpw == nullptr) + { + ModuleBase::WARNING_QUIT("MetaPW", "Constuctor of Operator::MetaPW is failed, please check your code!"); + } +} + namespace hamilt{ template class Meta>; #if ((defined __CUDA) || (defined __ROCM)) template class Meta>; +template Meta>::Meta(const Meta> *meta); +template Meta>::Meta(const Meta> *meta); #endif } // namespace hamilt \ No newline at end of file diff --git a/source/module_hamilt/ks_pw/meta_pw.h b/source/module_hamilt/ks_pw/meta_pw.h index 478309bf127..1bf26d990cd 100644 --- a/source/module_hamilt/ks_pw/meta_pw.h +++ b/source/module_hamilt/ks_pw/meta_pw.h @@ -23,6 +23,9 @@ class Meta> : public OperatorPW public: Meta(FPTYPE tpiba2_in, const int* isk_in, const ModuleBase::matrix* vk, ModulePW::PW_Basis_K* wfcpw); + template + explicit Meta(const Meta>* meta); + virtual ~Meta(){}; virtual void act( @@ -32,6 +35,12 @@ class Meta> : public OperatorPW std::complex* tmhpsi )const override; + // denghui added for copy constructor at 20221105 + FPTYPE get_tpiba() const {return this->tpiba;} + const int * get_isk() const {return this->isk;} + const ModuleBase::matrix* get_vk() const {return this->vk;} + ModulePW::PW_Basis_K* get_wfcpw() const {return this->wfcpw;} + private: mutable int max_npw = 0; diff --git a/source/module_hamilt/ks_pw/nonlocal_pw.cpp b/source/module_hamilt/ks_pw/nonlocal_pw.cpp index 2b598b930f8..994b0a67047 100644 --- a/source/module_hamilt/ks_pw/nonlocal_pw.cpp +++ b/source/module_hamilt/ks_pw/nonlocal_pw.cpp @@ -16,6 +16,7 @@ Nonlocal>::Nonlocal( const UnitCell* ucell_in ) { + this->classname = "Nonlocal"; this->cal_type = pw_nonlocal; this->isk = isk_in; this->ppcell = ppcell_in; @@ -331,9 +332,35 @@ void Nonlocal>::act ModuleBase::timer::tick("Operator", "NonlocalPW"); } +template +template +hamilt::Nonlocal>::Nonlocal(const Nonlocal> *nonlocal) +{ + this->classname = "Nonlocal"; + this->cal_type = pw_nonlocal; + this->ik = nonlocal->get_ik(); + this->isk = nonlocal->get_isk(); + this->ppcell = nonlocal->get_ppcell(); + this->ucell = nonlocal->get_ucell(); + if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { + this->deeq = this->ppcell->d_deeq; + resize_memory_op()(this->ctx, this->vkb, this->ppcell->vkb.size); + } + else { + this->deeq = this->ppcell->deeq.ptr; + this->vkb = this->ppcell->vkb.c; + } + if( this->isk == nullptr || this->ppcell == nullptr || this->ucell == nullptr) + { + ModuleBase::WARNING_QUIT("NonlocalPW", "Constuctor of Operator::NonlocalPW is failed, please check your code!"); + } +} + namespace hamilt{ template class Nonlocal>; #if ((defined __CUDA) || (defined __ROCM)) template class Nonlocal>; +template Nonlocal>::Nonlocal(const Nonlocal> *nonlocal); +template Nonlocal>::Nonlocal(const Nonlocal> *nonlocal); #endif } // namespace hamilt \ No newline at end of file diff --git a/source/module_hamilt/ks_pw/nonlocal_pw.h b/source/module_hamilt/ks_pw/nonlocal_pw.h index a79891762f0..5e645adb060 100644 --- a/source/module_hamilt/ks_pw/nonlocal_pw.h +++ b/source/module_hamilt/ks_pw/nonlocal_pw.h @@ -27,6 +27,9 @@ class Nonlocal> : public OperatorPW public: Nonlocal(const int* isk_in,const pseudopot_cell_vnl* ppcell_in,const UnitCell* ucell_in); + template + explicit Nonlocal(const Nonlocal>* nonlocal); + virtual ~Nonlocal(); virtual void init(const int ik_in)override; @@ -38,6 +41,10 @@ class Nonlocal> : public OperatorPW std::complex* tmhpsi )const override; + const int *get_isk() const {return this->isk;} + const pseudopot_cell_vnl *get_ppcell() const {return this->ppcell;} + const UnitCell_pseudo *get_ucell() const {return this->ucell;} + private: void add_nonlocal_pp(std::complex *hpsi_in, const std::complex *becp, const int m) const; diff --git a/source/module_hamilt/ks_pw/operator_pw.h b/source/module_hamilt/ks_pw/operator_pw.h index a589bffdd01..0ecae9a4d69 100644 --- a/source/module_hamilt/ks_pw/operator_pw.h +++ b/source/module_hamilt/ks_pw/operator_pw.h @@ -21,6 +21,8 @@ class OperatorPW : public Operator, Device> const int n_npwx, const std::complex* tmpsi_in, std::complex* tmhpsi)const; + + std::string classname = ""; }; }//end namespace hamilt diff --git a/source/module_hamilt/ks_pw/veff_pw.cpp b/source/module_hamilt/ks_pw/veff_pw.cpp index 9cb2856f8e8..86db829e695 100644 --- a/source/module_hamilt/ks_pw/veff_pw.cpp +++ b/source/module_hamilt/ks_pw/veff_pw.cpp @@ -13,6 +13,7 @@ Veff>::Veff( const ModuleBase::matrix* veff_in, ModulePW::PW_Basis_K* wfcpw_in) { + this->classname = "Veff"; this->cal_type = pw_veff; this->isk = isk_in; // this->veff = veff_in; @@ -31,6 +32,9 @@ Veff>::Veff( if (this->isk == nullptr || this->wfcpw == nullptr) { ModuleBase::WARNING_QUIT("VeffPW", "Constuctor of Operator::VeffPW is failed, please check your code!"); } + if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { + resize_memory_double_op()(this->ctx, this->d_veff, this->veff_col * this->veff_row); + } } template @@ -128,9 +132,36 @@ void Veff>::act( ModuleBase::timer::tick("Operator", "VeffPW"); } +template +template +hamilt::Veff>::Veff(const Veff> *veff) { + this->classname = "Veff"; + this->cal_type = pw_veff; + this->ik = veff->get_ik(); + this->isk = veff->get_isk(); + // this->veff = veff_in; + // TODO: add an GPU veff array + this->veff = veff->get_veff(); + this->veff_col = veff->get_veff_col(); + this->veff_row = veff->get_veff_row(); + this->wfcpw = veff->get_wfcpw(); + resize_memory_complex_op()(this->ctx, this->porter, this->wfcpw->nmaxgr); + if (this->npol != 1) { + resize_memory_complex_op()(this->ctx, this->porter1, this->wfcpw->nmaxgr); + } + if (this->isk == nullptr || this->veff == nullptr || this->wfcpw == nullptr) { + ModuleBase::WARNING_QUIT("VeffPW", "Constuctor of Operator::VeffPW is failed, please check your code!"); + } + if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { + resize_memory_double_op()(this->ctx, this->d_veff, this->veff_col * this->veff_row); + } +} + namespace hamilt{ template class Veff>; #if ((defined __CUDA) || (defined __ROCM)) template class Veff>; +template Veff>::Veff(const Veff> *veff); +template Veff>::Veff(const Veff> *veff); #endif } // namespace hamilt \ No newline at end of file diff --git a/source/module_hamilt/ks_pw/veff_pw.h b/source/module_hamilt/ks_pw/veff_pw.h index 4e9a4d88b78..2548f4ded92 100644 --- a/source/module_hamilt/ks_pw/veff_pw.h +++ b/source/module_hamilt/ks_pw/veff_pw.h @@ -24,6 +24,9 @@ class Veff> : public OperatorPW public: Veff(const int* isk_in,const ModuleBase::matrix* veff_in,ModulePW::PW_Basis_K* wfcpw_in); + template + explicit Veff(const Veff>* veff); + virtual ~Veff(); virtual void act ( @@ -33,6 +36,13 @@ class Veff> : public OperatorPW std::complex* tmhpsi )const override; + // denghui added for copy constructor at 20221105 + FPTYPE *get_veff() const {return this->veff;} + int get_veff_col() const {return this->veff_col;} + int get_veff_row() const {return this->veff_row;} + const int *get_isk() const {return isk;} + ModulePW::PW_Basis_K* get_wfcpw() const {return this->wfcpw;} + private: mutable int max_npw = 0; diff --git a/source/module_hamilt/operator.cpp b/source/module_hamilt/operator.cpp index 87e77a8fb73..b87cf317624 100644 --- a/source/module_hamilt/operator.cpp +++ b/source/module_hamilt/operator.cpp @@ -107,7 +107,9 @@ FPTYPE* Operator::get_hpsi(const hpsi_info& info) const hpsi_pointer = this->hpsi->get_pointer(); size_t total_hpsi_size = nbands_range * this->hpsi->get_nbasis(); - ModuleBase::GlobalFunc::ZEROS(hpsi_pointer, total_hpsi_size); + // ModuleBase::GlobalFunc::ZEROS(hpsi_pointer, total_hpsi_size); + // denghui replaced at 20221104 + set_memory_op()(this->ctx, hpsi_pointer, 0, total_hpsi_size); return hpsi_pointer; } diff --git a/source/module_hamilt/operator.h b/source/module_hamilt/operator.h index d1fbb935283..99e42d067e1 100644 --- a/source/module_hamilt/operator.h +++ b/source/module_hamilt/operator.h @@ -45,11 +45,14 @@ class Operator virtual void add(Operator* next); - #if ((defined __CUDA) || (defined __ROCM)) + virtual const int get_ik() const {return this->ik;} + +#if ((defined __CUDA) || (defined __ROCM)) typedef std::tuple*, const psi::Range, FPTYPE*> hpsi_info_gpu; - virtual hpsi_info_gpu hPsi_gpu(hpsi_info_gpu& input) const; + virtual hpsi_info_gpu hPsi_gpu(hpsi_info_gpu& input) const; #endif // ((defined __CUDA) || (defined __ROCM)) + Operator* next_op = nullptr; protected: int ik = 0; @@ -58,7 +61,6 @@ class Operator //calculation type, only different type can be in main chain table enum calculation_type cal_type; - Operator* next_op = nullptr; Operator* next_sub_op = nullptr; bool is_first_node = true; @@ -74,6 +76,9 @@ class Operator */ FPTYPE* get_hpsi(const hpsi_info& info)const; + Device *ctx = {}; + using set_memory_op = psi::memory::set_memory_op; + }; }//end namespace hamilt diff --git a/source/module_hsolver/diagh.h b/source/module_hsolver/diagh.h index eb7d53f21c7..f771baa1bae 100644 --- a/source/module_hsolver/diagh.h +++ b/source/module_hsolver/diagh.h @@ -9,6 +9,7 @@ namespace hsolver { +template class DiagH { public: @@ -16,10 +17,9 @@ class DiagH // virtual void init()=0; std::string method = "none"; - virtual void diag(hamilt::Hamilt *phm_in, psi::Psi> &psi, double *eigenvalue_in) = 0; + virtual void diag(hamilt::Hamilt *phm_in, psi::Psi, Device> &psi, FPTYPE *eigenvalue_in) = 0; - virtual void diag(hamilt::Hamilt *phm_in, psi::Psi &psi, double *eigenvalue_in) - { + virtual void diag(hamilt::Hamilt *phm_in, psi::Psi &psi, FPTYPE *eigenvalue_in) { return; } }; diff --git a/source/module_hsolver/diago_blas.cpp b/source/module_hsolver/diago_blas.cpp index 92201998cb6..4f5879df62c 100644 --- a/source/module_hsolver/diago_blas.cpp +++ b/source/module_hsolver/diago_blas.cpp @@ -21,7 +21,7 @@ typedef hamilt::MatrixBlock> matcd; namespace hsolver { -void DiagoBlas::diag(hamilt::Hamilt *phm_in, psi::Psi &psi, double *eigenvalue_in) +void DiagoBlas::diag(hamilt::Hamilt *phm_in, psi::Psi &psi, double *eigenvalue_in) { ModuleBase::TITLE("DiagoElpa", "diag"); matd h_mat, s_mat; @@ -33,7 +33,7 @@ void DiagoBlas::diag(hamilt::Hamilt *phm_in, psi::Psi &psi, double *eige BlasConnector::copy(GlobalV::NBANDS, eigen.data(), inc, eigenvalue_in, inc); } -void DiagoBlas::diag(hamilt::Hamilt *phm_in, psi::Psi> &psi, double *eigenvalue_in) +void DiagoBlas::diag(hamilt::Hamilt *phm_in, psi::Psi> &psi, double *eigenvalue_in) { ModuleBase::TITLE("DiagoElpa", "diag"); matcd h_mat, s_mat; diff --git a/source/module_hsolver/diago_blas.h b/source/module_hsolver/diago_blas.h index 2eaccd7e508..a5b027967a5 100644 --- a/source/module_hsolver/diago_blas.h +++ b/source/module_hsolver/diago_blas.h @@ -20,13 +20,13 @@ namespace hsolver { -class DiagoBlas : public DiagH +class DiagoBlas : public DiagH { public: - void diag(hamilt::Hamilt *phm_in, psi::Psi &psi, double *eigenvalue_in) override; + void diag(hamilt::Hamilt *phm_in, psi::Psi &psi, double *eigenvalue_in) override; - void diag(hamilt::Hamilt *phm_in, psi::Psi> &psi, double *eigenvalue_in) override; + void diag(hamilt::Hamilt *phm_in, psi::Psi> &psi, double *eigenvalue_in) override; private: void pdsygvx_diag(const int *const desc, diff --git a/source/module_hsolver/diago_cg.cpp b/source/module_hsolver/diago_cg.cpp index 5efe30ae8b9..9cd87390f6c 100644 --- a/source/module_hsolver/diago_cg.cpp +++ b/source/module_hsolver/diago_cg.cpp @@ -7,33 +7,41 @@ #include "module_base/timer.h" #include "src_parallel/parallel_reduce.h" #include "module_hsolver/include/math_kernel.h" +#include "module_hamilt/hamilt_pw.h" using namespace hsolver; template DiagoCG::DiagoCG(const FPTYPE* precondition_in) { + this->device = psi::device::get_device_type(this->ctx); this->precondition = precondition_in; test_cg = 0; reorder = false; this->device = psi::device::get_device_type(this->ctx); +#if defined(__CUDA) || defined(__ROCM) + hsolver::createBLAShandle(); +#endif } template DiagoCG::~DiagoCG() { // delete this->cg; // delete this->phi_m; - delete_memory_op()(this->ctx, this->sphi); - delete_memory_op()(this->ctx, this->hphi); - delete_memory_op()(this->ctx, this->scg); - delete_memory_op()(this->ctx, this->pphi); - delete_memory_op()(this->ctx, this->gradient); - delete_memory_op()(this->ctx, this->g0); - delete_memory_op()(this->ctx, this->lagrange); + delete_memory_complex_op()(this->ctx, this->sphi); + delete_memory_complex_op()(this->ctx, this->hphi); + delete_memory_complex_op()(this->ctx, this->scg); + delete_memory_complex_op()(this->ctx, this->pphi); + delete_memory_complex_op()(this->ctx, this->gradient); + delete_memory_complex_op()(this->ctx, this->g0); + delete_memory_complex_op()(this->ctx, this->lagrange); +#if defined(__CUDA) || defined(__ROCM) + hsolver::destoryBLAShandle(); +#endif } template -void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, psi::Psi, Device> &phi, FPTYPE *eigenvalue_in) +void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, psi::Psi, Device> &phi, FPTYPE *eigenvalue_in) { ModuleBase::TITLE("DiagoCG", "diag_once"); ModuleBase::timer::tick("DiagoCG", "diag_once"); @@ -47,8 +55,7 @@ void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, psi::Psin_band = phi.get_nbands(); this->eigenvalue = eigenvalue_in; - // haozhihan replace ZEROS - psi::memory::set_memory_op()(this->ctx, this->eigenvalue, 0, this->n_band); + set_memory_var_cpu_op()(this->cpu_ctx, this->eigenvalue, 0, this->n_band); // ModuleBase::GlobalFunc::ZEROS(this->eigenvalue, this->n_band); /// record for how many loops in cg convergence @@ -63,32 +70,32 @@ void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, psi::Psiphi_m = new psi::Psi, Device>(phi, 1, 1); // this->hphi.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->hphi, this->dmx); - set_memory_op()(this->ctx, this->hphi, 0, this->dmx); + resize_memory_complex_op()(this->ctx, this->hphi, this->dmx); + set_memory_complex_op()(this->ctx, this->hphi, 0, this->dmx); // this->sphi.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->sphi, this->dmx); - set_memory_op()(this->ctx, this->sphi, 0, this->dmx); + resize_memory_complex_op()(this->ctx, this->sphi, this->dmx); + set_memory_complex_op()(this->ctx, this->sphi, 0, this->dmx); this->cg = new psi::Psi, Device>(phi, 1, 1); // this->scg.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->scg, this->dmx); - set_memory_op()(this->ctx, this->scg, 0, this->dmx); + resize_memory_complex_op()(this->ctx, this->scg, this->dmx); + set_memory_complex_op()(this->ctx, this->scg, 0, this->dmx); // this->pphi.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->pphi, this->dmx); - set_memory_op()(this->ctx, this->pphi, 0, this->dmx); + resize_memory_complex_op()(this->ctx, this->pphi, this->dmx); + set_memory_complex_op()(this->ctx, this->pphi, 0, this->dmx); //in band_by_band CG method, only the first band in phi_m would be calculated psi::Range cg_hpsi_range(0); // this->gradient.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->gradient, this->dmx); - set_memory_op()(this->ctx, this->gradient, 0, this->dmx); + resize_memory_complex_op()(this->ctx, this->gradient, this->dmx); + set_memory_complex_op()(this->ctx, this->gradient, 0, this->dmx); // this->g0.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->g0, this->dmx); - set_memory_op()(this->ctx, this->g0, 0, this->dmx); + resize_memory_complex_op()(this->ctx, this->g0, this->dmx); + set_memory_complex_op()(this->ctx, this->g0, 0, this->dmx); // this->lagrange.resize(this->n_band, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->lagrange, this->n_band); - set_memory_op()(this->ctx, this->lagrange, 0, this->n_band); + resize_memory_complex_op()(this->ctx, this->lagrange, this->n_band); + set_memory_complex_op()(this->ctx, this->lagrange, 0, this->n_band); for (int m = 0; m < this->n_band; m++) { @@ -124,7 +131,7 @@ void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, psi::Psi::PW_DIAG_NMAX; iter++) + for (iter = 0; iter < DiagoIterAssist::PW_DIAG_NMAX; iter++) { this->calculate_gradient(); this->orthogonal_gradient(phm_in, phi, m); @@ -136,8 +143,9 @@ void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, psi::PsisPsi(this->cg->get_pointer(), this->scg, (size_t)this->dim); converged = this->update_psi(cg_norm, theta, this->eigenvalue[m]); - if (converged) + if (converged) { break; + } } // end iter std::complex* psi_temp = &(phi(m, 0)); @@ -157,13 +165,13 @@ void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, psi::Psi 0 && reorder) { ModuleBase::GlobalFunc::NOTE("reorder bands!"); - if (eigenvalue[m] - eigenvalue[m - 1] < -2.0 * DiagoIterAssist::PW_DIAG_THR) + if (eigenvalue[m] - eigenvalue[m - 1] < -2.0 * DiagoIterAssist::PW_DIAG_THR) { // if the last calculated eigenvalue is not the largest... int i = 0; for (i = m - 2; i >= 0; i--) { - if (eigenvalue[m] - eigenvalue[i] > 2.0 * DiagoIterAssist::PW_DIAG_THR) + if (eigenvalue[m] - eigenvalue[i] > 2.0 * DiagoIterAssist::PW_DIAG_THR) break; } i++; @@ -200,180 +208,20 @@ void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, psi::Psin_band; - DiagoIterAssist::avg_iter += avg; + DiagoIterAssist::avg_iter += avg; delete this->phi_m; delete this->cg; ModuleBase::timer::tick("DiagoCG", "diag_once"); - return; } // end subroutine ccgdiagg -#if ((defined __CUDA) || (defined __ROCM)) -template<> -void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, psi::Psi, psi::DEVICE_GPU> &phi, double *eigenvalue_in) -{ - ModuleBase::TITLE("DiagoCG", "diag_once"); - ModuleBase::timer::tick("DiagoCG", "diag_once"); - - /// out : record for states of convergence - this->notconv = 0; - - /// initialize variables - this->dim = phi.get_current_nbas(); - this->dmx = phi.get_nbasis(); - this->n_band = phi.get_nbands(); - this->eigenvalue = eigenvalue_in; - ModuleBase::GlobalFunc::ZEROS(this->eigenvalue, this->n_band); - - /// record for how many loops in cg convergence - double avg = 0.0; - - //------------------------------------------------------------------- - // "poor man" iterative diagonalization of a complex hermitian matrix - // through preconditioned conjugate gradient algorithm - // Band-by-band algorithm with minimal use of memory - // Calls hPhi and sPhi to calculate H|phi> and S|phi> - // Works for generalized eigenvalue problem (US pseudopotentials) as well - //------------------------------------------------------------------- - this->phi_m = new psi::Psi, psi::DEVICE_GPU>(phi, 1, 1); - // this->hphi.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->hphi, this->dmx); - set_memory_op()(this->ctx, this->hphi, 0, this->dmx); - // this->sphi.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->sphi, this->dmx); - set_memory_op()(this->ctx, this->sphi, 0, this->dmx); - - this->cg = new psi::Psi, psi::DEVICE_GPU>(phi, 1, 1); - // this->scg.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->scg, this->dmx); - set_memory_op()(this->ctx, this->scg, 0, this->dmx); - // this->pphi.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->pphi, this->dmx); - set_memory_op()(this->ctx, this->pphi, 0, this->dmx); - - //in band_by_band CG method, only the first band in phi_m would be calculated - psi::Range cg_hpsi_range(0); - - // this->gradient.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->gradient, this->dmx); - set_memory_op()(this->ctx, this->gradient, 0, this->dmx); - // this->g0.resize(this->dmx, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->g0, this->dmx); - set_memory_op()(this->ctx, this->g0, 0, this->dmx); - // this->lagrange.resize(this->n_band, ModuleBase::ZERO); - resize_memory_op()(this->ctx, this->lagrange, this->n_band); - set_memory_op()(this->ctx, this->lagrange, 0, this->dmx); - - for (int m = 0; m < this->n_band; m++) - { - if (test_cg > 2) - GlobalV::ofs_running << "Diagonal Band : " << m << std::endl; - //copy psi_in into internal psi, m=0 has been done in Constructor - if(m>0) - { - const std::complex* psi_m_in = &(phi(m, 0)); - auto pphi_m = this->phi_m->get_pointer(); - ModuleBase::GlobalFunc::COPYARRAY(psi_m_in, pphi_m, this->dim); - } - phm_in->sPsi(this->phi_m->get_pointer(), this->sphi, (size_t)this->dim); // sphi = S|psi(m)> - this->schmit_orth(m, phi); - phm_in->sPsi(this->phi_m->get_pointer(), this->sphi, (size_t)this->dim); // sphi = S|psi(m)> - - //do hPsi, actually the result of hpsi stored in Operator, - //the necessary of copying operation should be checked later - hpsi_info_gpu cg_hpsi_in(this->phi_m, cg_hpsi_range, this->hphi); - phm_in->ops->hPsi_gpu(cg_hpsi_in); - - this->eigenvalue[m] = zdot_real_op()(this->ctx, this->dim, this->phi_m->get_pointer(), this->hphi, this->phi_m->get_device()); - - int iter = 0; - double gg_last = 0.0; - double cg_norm = 0.0; - double theta = 0.0; - bool converged = false; - for (iter = 0; iter < DiagoIterAssist::PW_DIAG_NMAX; iter++) - { - this->calculate_gradient(); - this->orthogonal_gradient(phm_in, phi, m); - this->calculate_gamma_cg(iter, gg_last, cg_norm, theta); - - hpsi_info_gpu cg_hpsi_in(this->cg, cg_hpsi_range, this->pphi); - phm_in->ops->hPsi_gpu(cg_hpsi_in); - - phm_in->sPsi(this->cg->get_pointer(), this->scg, (size_t)this->dim); - converged = this->update_psi(cg_norm, theta, this->eigenvalue[m]); - - if (converged) - break; - } // end iter - - std::complex* psi_temp = &(phi(m, 0)); - ModuleBase::GlobalFunc::COPYARRAY(this->phi_m->get_pointer(), psi_temp, this->dim); - - if (!converged) - { - ++this->notconv; - } - avg += static_cast(iter) + 1.00; - - // reorder eigenvalues if they are not in the right order - // (this CAN and WILL happen in not-so-special cases) - - if (m > 0 && reorder) - { - ModuleBase::GlobalFunc::NOTE("reorder bands!"); - if (eigenvalue[m] - eigenvalue[m - 1] < -2.0 * DiagoIterAssist::PW_DIAG_THR) - { - // if the last calculated eigenvalue is not the largest... - int i = 0; - for (i = m - 2; i >= 0; i--) - { - if (eigenvalue[m] - eigenvalue[i] > 2.0 * DiagoIterAssist::PW_DIAG_THR) - break; - } - i++; - - // last calculated eigenvalue should be in the i-th position: reorder - double e0 = eigenvalue[m]; - ModuleBase::GlobalFunc::COPYARRAY(psi_temp, this->pphi, this->dim); - - for (int j = m; j >= i + 1; j--) - { - eigenvalue[j] = eigenvalue[j - 1]; - std::complex* phi_j = &phi(j, 0); - std::complex* phi_j1 = &phi(j-1, 0); - ModuleBase::GlobalFunc::COPYARRAY(phi_j1, phi_j, this->dim); - } - - eigenvalue[i] = e0; - // dcopy(pphi, phi, i); - std::complex* phi_pointer = &phi(i, 0); - ModuleBase::GlobalFunc::COPYARRAY(this->pphi, phi_pointer, this->dim); - // this procedure should be good if only a few inversions occur, - // extremely inefficient if eigenvectors are often in bad order - // (but this should not happen) - } // endif - } // end reorder - - } // end m - - avg /= this->n_band; - DiagoIterAssist::avg_iter += avg; - - delete this->phi_m; - delete this->cg; - - ModuleBase::timer::tick("DiagoCG", "diag_once"); - return; -} // end subroutine ccgdiagg -#endif // ((defined __CUDA) || (defined __ROCM)) - template void DiagoCG::calculate_gradient() { - if (this->test_cg == 1) + if (this->test_cg == 1) { ModuleBase::TITLE("DiagoCG", "calculate_gradient"); + } // ModuleBase::timer::tick("DiagoCG","grad"); // for (int i = 0; i < this->dim; i++) @@ -384,9 +232,14 @@ void DiagoCG::calculate_gradient() // this->pphi[i] = this->sphi[i] / this->precondition[i]; // } // haozhihan replace this 2022-10-6 - vector_div_vector_op()(this->ctx, this->dim, this->gradient, this->hphi, this->precondition); - vector_div_vector_op()(this->ctx, this->dim, this->pphi, this->sphi, this->precondition); - + if (this->device == psi::GpuDevice) { + vector_div_vector_op()(this->ctx, this->dim, this->gradient, this->hphi, this->d_precondition); + vector_div_vector_op()(this->ctx, this->dim, this->pphi, this->sphi, this->d_precondition); + } + else { + vector_div_vector_op()(this->ctx, this->dim, this->gradient, this->hphi, this->precondition); + vector_div_vector_op()(this->ctx, this->dim, this->pphi, this->sphi, this->precondition); + } // Update lambda ! @@ -411,14 +264,14 @@ void DiagoCG::calculate_gradient() // ModuleBase::timer::tick("DiagoCG","grad"); - return; } template -void DiagoCG::orthogonal_gradient(hamilt::Hamilt *phm_in, const psi::Psi> &eigenfunction, const int m) +void DiagoCG::orthogonal_gradient(hamilt::Hamilt *phm_in, const psi::Psi, Device> &eigenfunction, const int m) { - if (test_cg == 1) + if (test_cg == 1) { ModuleBase::TITLE("DiagoCG", "orthogonal_gradient"); + } // ModuleBase::timer::tick("DiagoCG","orth_grad"); phm_in->sPsi(this->gradient, this->scg, (size_t)this->dim); @@ -535,8 +388,9 @@ void DiagoCG::orthogonal_gradient(hamilt::Hamilt *phm_in, const template void DiagoCG::calculate_gamma_cg(const int iter, FPTYPE &gg_last, const FPTYPE &cg_norm, const FPTYPE &theta) { - if (test_cg == 1) + if (test_cg == 1) { ModuleBase::TITLE("DiagoCG", "calculate_gamma_cg"); + } // ModuleBase::timer::tick("DiagoCG","gamma_cg"); auto pcg = this->cg->get_pointer(); auto pphi_m = this->phi_m->get_pointer(); @@ -560,7 +414,12 @@ void DiagoCG::calculate_gamma_cg(const int iter, FPTYPE &gg_last // this->g0[i] = this->precondition[i] * this->scg[i]; // } // haozhihan replace this 2022-10-6 - vector_mul_vector_op()(this->ctx, this->dim, this->g0, this->scg, this->precondition); + if (this->device == psi::GpuDevice) { + vector_mul_vector_op()(this->ctx, this->dim, this->g0, this->scg, this->d_precondition); + } + else { + vector_mul_vector_op()(this->ctx, this->dim, this->g0, this->scg, this->precondition); + } // (3) Update gg_now! // gg_now = < g|P|scg > = < g|g0 > @@ -607,7 +466,6 @@ void DiagoCG::calculate_gamma_cg(const int iter, FPTYPE &gg_last }*/ } // ModuleBase::timer::tick("DiagoCG","gamma_cg"); - return; } template @@ -663,7 +521,7 @@ bool DiagoCG::update_psi(FPTYPE &cg_norm, FPTYPE &theta, FPTYPE // std::cout << "\n overlap2 = " << this->ddot(dim, phi_m, phi_m); - if (abs(eigenvalue - e0) < DiagoIterAssist::PW_DIAG_THR) + if (abs(eigenvalue - e0) < DiagoIterAssist::PW_DIAG_THR) { // ModuleBase::timer::tick("DiagoCG","update"); return 1; @@ -700,7 +558,8 @@ void DiagoCG::schmit_orth( assert(m >= 0); assert(psi.get_nbands() >= m); - std::vector> lagrange_so(m + 1, ModuleBase::ZERO); + std::complex * lagrange_so = nullptr; + resize_memory_complex_op()(this->ctx, lagrange_so, m + 1); //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // haozhihan replace 2022-10-6 @@ -715,7 +574,7 @@ void DiagoCG::schmit_orth( this->sphi, inc, &ModuleBase::ZERO, - lagrange_so.data(), + lagrange_so, inc); //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // qianrui replace 2021-3-15 @@ -744,9 +603,11 @@ void DiagoCG::schmit_orth( //>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> // be careful , here reduce m+1 - Parallel_Reduce::reduce_complex_double_pool(lagrange_so.data(), m + 1); + Parallel_Reduce::reduce_complex_double_pool(lagrange_so, m + 1); - FPTYPE psi_norm = lagrange_so[m].real(); + std::complex var = {0, 0}; + syncmem_complex_d2h_op()(this->cpu_ctx, this->ctx, &var, lagrange_so + m, 1); + FPTYPE psi_norm = var.real(); //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // haozhihan replace 2022-10-6 @@ -757,7 +618,7 @@ void DiagoCG::schmit_orth( &ModuleBase::NEG_ONE, psi.get_pointer(), this->dmx, - lagrange_so.data(), + lagrange_so, inc, &ModuleBase::ONE, this->phi_m->get_pointer(), @@ -777,7 +638,7 @@ void DiagoCG::schmit_orth( // this->phi_m->get_pointer(), // &inc); - psi_norm -= hsolver::zdot_real_op()(this->ctx, m, lagrange_so.data(), lagrange_so.data(), false); + psi_norm -= hsolver::zdot_real_op()(this->ctx, m, lagrange_so, lagrange_so, false); //====================================================================== /*for (int j = 0; j < m; j++) { @@ -816,37 +677,53 @@ void DiagoCG::schmit_orth( // } // ModuleBase::timer::tick("DiagoCG","schmit_orth"); - return; + delete_memory_complex_op()(this->ctx, lagrange_so); } template -void DiagoCG::diag(hamilt::Hamilt *phm_in, psi::Psi, Device> &psi, FPTYPE *eigenvalue_in) +void DiagoCG::diag(hamilt::Hamilt *phm_in, psi::Psi, Device> &psi, FPTYPE *eigenvalue_in) { /// record the times of trying iterative diagonalization int ntry = 0; this->notconv = 0; + if (this->device == psi::GpuDevice) { + resize_memory_var_op()(this->ctx, this->d_precondition, psi.get_nbasis()); + syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, this->d_precondition, this->precondition, psi.get_nbasis()); + } do { - if(DiagoIterAssist::need_subspace || ntry > 0) + if(DiagoIterAssist::need_subspace || ntry > 0) { - DiagoIterAssist::diagH_subspace(phm_in, psi, psi, eigenvalue_in); + #if defined(__CUDA) || defined(__ROCM) + psi::Psi, psi::DEVICE_CPU> cpu_psi = psi; + // hamilt::Hamilt* h_phm_in = new hamilt::HamiltPW(reinterpret_cast*>(phm_in)); + hamilt::Hamilt* h_phm_in = + new hamilt::HamiltPW( + reinterpret_cast*>(phm_in)); + DiagoIterAssist::diagH_subspace(h_phm_in, cpu_psi, cpu_psi, eigenvalue_in); + psi::memory::synchronize_memory_op, Device, psi::DEVICE_CPU>()( + psi.get_device(), + cpu_psi.get_device(), + psi.get_pointer() - psi.get_psi_bias(), + cpu_psi.get_pointer() - cpu_psi.get_psi_bias(), + psi.size()); + #else + DiagoIterAssist::diagH_subspace(phm_in, psi, psi, eigenvalue_in); + #endif } - DiagoIterAssist::avg_iter += 1.0; + DiagoIterAssist::avg_iter += 1.0; this->reorder = true; this->diag_mock(phm_in, psi, eigenvalue_in); ++ntry; - } while (DiagoIterAssist::test_exit_cond(ntry, this->notconv)); + } while (DiagoIterAssist::test_exit_cond(ntry, this->notconv)); - if (notconv > max(5, psi.get_nbands() / 4)) - { + if (notconv > max(5, psi.get_nbands() / 4)) { std::cout << "\n notconv = " << this->notconv; std::cout << "\n DiagoCG::diag', too many bands are not converged! \n"; } - - return; } namespace hsolver{ diff --git a/source/module_hsolver/diago_cg.h b/source/module_hsolver/diago_cg.h index cf636030c29..4a04a48a03b 100644 --- a/source/module_hsolver/diago_cg.h +++ b/source/module_hsolver/diago_cg.h @@ -14,7 +14,7 @@ namespace hsolver { template -class DiagoCG : public DiagH +class DiagoCG : public DiagH { public: // Constructor need: @@ -26,7 +26,7 @@ class DiagoCG : public DiagH // virtual void init(){}; // refactor hpsi_info // this is the override function diag() for CG method - void diag(hamilt::Hamilt *phm_in, psi::Psi, Device> &psi, FPTYPE *eigenvalue_in); + void diag(hamilt::Hamilt *phm_in, psi::Psi, Device> &psi, FPTYPE *eigenvalue_in) override; private: /// static variables, used for passing control variables @@ -48,7 +48,7 @@ class DiagoCG : public DiagH const FPTYPE *precondition = nullptr; /// eigenvalue results FPTYPE *eigenvalue = nullptr; - + FPTYPE *d_precondition = nullptr; /// temp vector for new psi for one band, size dim psi::Psi, Device>* phi_m = nullptr; @@ -76,10 +76,11 @@ class DiagoCG : public DiagH /// device type of psi psi::AbacusDevice_t device = {}; Device * ctx = {}; + psi::DEVICE_CPU *cpu_ctx = {}; void calculate_gradient(); - void orthogonal_gradient(hamilt::Hamilt *phm_in, const psi::Psi> &eigenfunction, const int m); + void orthogonal_gradient(hamilt::Hamilt *phm_in, const psi::Psi, Device> &eigenfunction, const int m); void calculate_gamma_cg(const int iter, FPTYPE &gg_last, const FPTYPE &cg0, const FPTYPE &theta); @@ -88,16 +89,19 @@ class DiagoCG : public DiagH void schmit_orth(const int &m, const psi::Psi, Device> &psi); // used in diag() for template replace Hamilt with Hamilt_PW - void diag_mock(hamilt::Hamilt *phm_in, psi::Psi, Device> &phi, FPTYPE *eigenvalue_in); + void diag_mock(hamilt::Hamilt *phm_in, psi::Psi, Device> &phi, FPTYPE *eigenvalue_in); + using hpsi_info = typename hamilt::Operator, Device>::hpsi_info; using zdot_real_op = hsolver::zdot_real_op; - using set_memory_op = psi::memory::set_memory_op, Device>; - using delete_memory_op = psi::memory::delete_memory_op, Device>; - using resize_memory_op = psi::memory::resize_memory_op, Device>; - using hpsi_info = typename hamilt::Operator, psi::DEVICE_CPU>::hpsi_info; -#if ((defined __CUDA) || (defined __ROCM)) - using hpsi_info_gpu = typename hamilt::Operator, psi::DEVICE_GPU>::hpsi_info_gpu; -#endif + + using set_memory_complex_op = psi::memory::set_memory_op, Device>; + using delete_memory_complex_op = psi::memory::delete_memory_op, Device>; + using resize_memory_complex_op = psi::memory::resize_memory_op, Device>; + using syncmem_complex_d2h_op = psi::memory::synchronize_memory_op, psi::DEVICE_CPU, Device>; + + using set_memory_var_cpu_op = psi::memory::set_memory_op; + using resize_memory_var_op = psi::memory::resize_memory_op; + using syncmem_var_h2d_op = psi::memory::synchronize_memory_op; }; } // namespace hsolver diff --git a/source/module_hsolver/diago_david.cpp b/source/module_hsolver/diago_david.cpp index d88ee636a08..5ef39a43c6f 100644 --- a/source/module_hsolver/diago_david.cpp +++ b/source/module_hsolver/diago_david.cpp @@ -27,7 +27,7 @@ DiagoDavid::DiagoDavid(const FPTYPE* precondition_in) } template -void DiagoDavid::cal_grad(hamilt::Hamilt* phm_in, +void DiagoDavid::cal_grad(hamilt::Hamilt* phm_in, const int &npw, const int &nbase, // current dimension of the reduced basis const int ¬conv, @@ -197,7 +197,7 @@ void DiagoDavid::cal_grad(hamilt::Hamilt* phm_in, #if ((defined __CUDA) || (defined __ROCM)) template<> -void DiagoDavid::cal_grad(hamilt::Hamilt* phm_in, +void DiagoDavid::cal_grad(hamilt::Hamilt* phm_in, const int &npw, const int &nbase, // current dimension of the reduced basis const int ¬conv, @@ -366,7 +366,7 @@ void DiagoDavid::cal_grad(hamilt::Hamilt* phm_in, template<> -void DiagoDavid::diag_mock(hamilt::Hamilt* phm_in, psi::Psi, psi::DEVICE_GPU> &psi, double* eigenvalue_in) +void DiagoDavid::diag_mock(hamilt::Hamilt* phm_in, psi::Psi, psi::DEVICE_GPU> &psi, double* eigenvalue_in) { if (test_david == 1) ModuleBase::TITLE("DiagoDavid", "diag_mock"); @@ -563,7 +563,7 @@ void DiagoDavid::diag_mock(hamilt::Hamilt* phm_in, psi: #endif // ((defined __CUDA) || (defined __ROCM)) template -void DiagoDavid::diag_mock(hamilt::Hamilt* phm_in, psi::Psi, Device> &psi, FPTYPE *eigenvalue_in) +void DiagoDavid::diag_mock(hamilt::Hamilt* phm_in, psi::Psi, Device> &psi, FPTYPE *eigenvalue_in) { if (test_david == 1) ModuleBase::TITLE("DiagoDavid", "diag_mock"); @@ -1350,7 +1350,7 @@ void DiagoDavid::planSchmitOrth( } template -void DiagoDavid::diag(hamilt::Hamilt *phm_in, psi::Psi, Device> &psi, FPTYPE* eigenvalue_in) +void DiagoDavid::diag(hamilt::Hamilt *phm_in, psi::Psi, Device> &psi, FPTYPE* eigenvalue_in) { /// record the times of trying iterative diagonalization int ntry = 0; diff --git a/source/module_hsolver/diago_david.h b/source/module_hsolver/diago_david.h index 696aa99408e..892fadbb1ef 100644 --- a/source/module_hsolver/diago_david.h +++ b/source/module_hsolver/diago_david.h @@ -17,13 +17,13 @@ namespace hsolver { template -class DiagoDavid : public DiagH +class DiagoDavid : public DiagH { public: DiagoDavid(const FPTYPE* precondition_in); // this is the override function diag() for CG method - void diag(hamilt::Hamilt* phm_in, psi::Psi, Device>& phi, FPTYPE* eigenvalue_in); + void diag(hamilt::Hamilt* phm_in, psi::Psi, Device>& phi, FPTYPE* eigenvalue_in); static int PW_DIAG_NDIM; @@ -37,7 +37,7 @@ class DiagoDavid : public DiagH psi::AbacusDevice_t device = {}; psi::DEVICE_CPU * ctx = {}; - void cal_grad(hamilt::Hamilt* phm_in, + void cal_grad(hamilt::Hamilt* phm_in, const int& npw, const int& nbase, const int& notconv, @@ -99,7 +99,7 @@ class DiagoDavid : public DiagH FPTYPE* e, ModuleBase::ComplexMatrix& vc); - void diag_mock(hamilt::Hamilt* phm_in, psi::Psi, Device>& psi, FPTYPE* eigenvalue_in); + void diag_mock(hamilt::Hamilt* phm_in, psi::Psi, Device>& psi, FPTYPE* eigenvalue_in); const FPTYPE* precondition = nullptr; diff --git a/source/module_hsolver/diago_elpa.cpp b/source/module_hsolver/diago_elpa.cpp index 69ac995e3db..a47c363fe2f 100644 --- a/source/module_hsolver/diago_elpa.cpp +++ b/source/module_hsolver/diago_elpa.cpp @@ -17,7 +17,7 @@ typedef hamilt::MatrixBlock> matcd; namespace hsolver { int DiagoElpa::DecomposedState = 0; -void DiagoElpa::diag(hamilt::Hamilt *phm_in, psi::Psi> &psi, double *eigenvalue_in) +void DiagoElpa::diag(hamilt::Hamilt *phm_in, psi::Psi> &psi, double *eigenvalue_in) { ModuleBase::TITLE("DiagoElpa", "diag"); #ifdef __MPI @@ -42,7 +42,7 @@ void DiagoElpa::diag(hamilt::Hamilt *phm_in, psi::Psi> &psi #endif } -void DiagoElpa::diag(hamilt::Hamilt *phm_in, psi::Psi &psi, double *eigenvalue_in) +void DiagoElpa::diag(hamilt::Hamilt *phm_in, psi::Psi &psi, double *eigenvalue_in) { ModuleBase::TITLE("DiagoElpa", "diag"); #ifdef __MPI diff --git a/source/module_hsolver/diago_elpa.h b/source/module_hsolver/diago_elpa.h index 0820bc62006..7bad7dbf0d7 100644 --- a/source/module_hsolver/diago_elpa.h +++ b/source/module_hsolver/diago_elpa.h @@ -7,13 +7,13 @@ namespace hsolver { -class DiagoElpa : public DiagH +class DiagoElpa : public DiagH { public: - void diag(hamilt::Hamilt* phm_in, psi::Psi& psi, double* eigenvalue_in) override; + void diag(hamilt::Hamilt* phm_in, psi::Psi& psi, double* eigenvalue_in) override; - void diag(hamilt::Hamilt* phm_in, psi::Psi>& psi, double* eigenvalue_in) override; + void diag(hamilt::Hamilt* phm_in, psi::Psi>& psi, double* eigenvalue_in) override; static int DecomposedState; diff --git a/source/module_hsolver/diago_iter_assist.cpp b/source/module_hsolver/diago_iter_assist.cpp index 61ac16ac250..be44920b1a1 100644 --- a/source/module_hsolver/diago_iter_assist.cpp +++ b/source/module_hsolver/diago_iter_assist.cpp @@ -29,7 +29,7 @@ bool DiagoIterAssist::need_subspace = false; //---------------------------------------------------------------------- template void DiagoIterAssist::diagH_subspace( - hamilt::Hamilt* pHamilt, + hamilt::Hamilt* pHamilt, const psi::Psi, Device> &psi, psi::Psi, Device> &evc, FPTYPE *en, @@ -174,7 +174,7 @@ void DiagoIterAssist::diagH_subspace( template void DiagoIterAssist::diagH_subspace_init( - hamilt::Hamilt* pHamilt, + hamilt::Hamilt* pHamilt, const ModuleBase::ComplexMatrix &psi, psi::Psi, Device> &evc, FPTYPE *en) @@ -335,12 +335,12 @@ void DiagoIterAssist::diagH_subspace_init( // Produces on output n_band eigenvectors (n_band <= nstart) in evc. //---------------------------------------------------------------------- template<> -void DiagoIterAssist::diagH_subspace(hamilt::Hamilt* pHamilt, const psi::Psi, psi::DEVICE_GPU> &psi, psi::Psi, psi::DEVICE_GPU> &evc, double *en, int n_band) { +void DiagoIterAssist::diagH_subspace(hamilt::Hamilt* pHamilt, const psi::Psi, psi::DEVICE_GPU> &psi, psi::Psi, psi::DEVICE_GPU> &evc, double *en, int n_band) { ModuleBase::WARNING_QUIT("DiagoIterAssist::diagH_subspace","GPU's implementation is not supported currently!"); } template<> -void DiagoIterAssist::diagH_subspace_init(hamilt::Hamilt* pHamilt, const ModuleBase::ComplexMatrix &psi, psi::Psi, psi::DEVICE_GPU> &evc, double *en) { +void DiagoIterAssist::diagH_subspace_init(hamilt::Hamilt* pHamilt, const ModuleBase::ComplexMatrix &psi, psi::Psi, psi::DEVICE_GPU> &evc, double *en) { ModuleBase::WARNING_QUIT("DiagoIterAssist::diagH_subspace_init","GPU's implementation is not supported currently!"); } #endif diff --git a/source/module_hsolver/diago_iter_assist.h b/source/module_hsolver/diago_iter_assist.h index f0f3c9b7d39..79bc4baa30f 100644 --- a/source/module_hsolver/diago_iter_assist.h +++ b/source/module_hsolver/diago_iter_assist.h @@ -21,14 +21,14 @@ class DiagoIterAssist // for CG diagonalization only static void diagH_subspace( - hamilt::Hamilt* pHamilt, + hamilt::Hamilt* pHamilt, const psi::Psi, Device> &psi, psi::Psi, Device> &evc, FPTYPE *en, int n_band = 0); // for initializing wave function , this is a template function static void diagH_subspace_init( - hamilt::Hamilt* pHamilt, + hamilt::Hamilt* pHamilt, const ModuleBase::ComplexMatrix &psi, psi::Psi, Device> &evc, FPTYPE *en); @@ -45,9 +45,6 @@ class DiagoIterAssist static bool test_exit_cond(const int &ntry, const int ¬conv); using hpsi_info = typename hamilt::Operator, psi::DEVICE_CPU>::hpsi_info; -#if ((defined __CUDA) || (defined __ROCM)) - using hpsi_info_gpu = typename hamilt::Operator, psi::DEVICE_GPU>::hpsi_info_gpu; -#endif }; } // namespace hsolver diff --git a/source/module_hsolver/diago_lapack.cpp b/source/module_hsolver/diago_lapack.cpp index 3757c172cb9..489421777e6 100644 --- a/source/module_hsolver/diago_lapack.cpp +++ b/source/module_hsolver/diago_lapack.cpp @@ -15,7 +15,7 @@ typedef hamilt::MatrixBlock> matcd; namespace hsolver { -void DiagoLapack::diag(hamilt::Hamilt *phm_in, psi::Psi> &psi, double *eigenvalue_in) +void DiagoLapack::diag(hamilt::Hamilt *phm_in, psi::Psi> &psi, double *eigenvalue_in) { ModuleBase::TITLE("DiagoLapack", "diag"); assert(GlobalV::NPROC == 1); @@ -76,7 +76,7 @@ void DiagoLapack::diag(hamilt::Hamilt *phm_in, psi::Psi> &p } } -void DiagoLapack::diag(hamilt::Hamilt *phm_in, psi::Psi &psi, double *eigenvalue_in) +void DiagoLapack::diag(hamilt::Hamilt *phm_in, psi::Psi &psi, double *eigenvalue_in) { ModuleBase::TITLE("DiagoLapack", "diag"); } diff --git a/source/module_hsolver/diago_lapack.h b/source/module_hsolver/diago_lapack.h index 1dbe35de8bc..2d303be6998 100644 --- a/source/module_hsolver/diago_lapack.h +++ b/source/module_hsolver/diago_lapack.h @@ -7,13 +7,13 @@ namespace hsolver { -class DiagoLapack : public DiagH +class DiagoLapack : public DiagH { public: - void diag(hamilt::Hamilt* phm_in, psi::Psi& psi, double* eigenvalue_in) override; + void diag(hamilt::Hamilt* phm_in, psi::Psi& psi, double* eigenvalue_in) override; - void diag(hamilt::Hamilt* phm_in, psi::Psi>& psi, double* eigenvalue_in) override; + void diag(hamilt::Hamilt* phm_in, psi::Psi>& psi, double* eigenvalue_in) override; private: }; diff --git a/source/module_hsolver/hsolver.h b/source/module_hsolver/hsolver.h index b38e517c778..9ba048c7cae 100644 --- a/source/module_hsolver/hsolver.h +++ b/source/module_hsolver/hsolver.h @@ -32,7 +32,7 @@ class HSolver // solve Hamiltonian to electronic density in ElecState virtual void solve ( - hamilt::Hamilt* phm, + hamilt::Hamilt* phm, psi::Psi>& ppsi, elecstate::ElecState* pes, const std::string method, @@ -43,7 +43,7 @@ class HSolver } virtual void solve ( - hamilt::Hamilt* phm, + hamilt::Hamilt* phm, psi::Psi& ppsi, elecstate::ElecState* pes, const std::string method, @@ -55,7 +55,7 @@ class HSolver virtual void solve ( - hamilt::Hamilt* phm, + hamilt::Hamilt* phm, psi::Psi>& ppsi, elecstate::ElecState* pes, Stochastic_WF& stowf, @@ -94,7 +94,8 @@ class HSolver }; protected: - DiagH* pdiagh = nullptr; // for single Hamiltonian matrix diagonal solver + DiagH* pdiagh = nullptr; // for single Hamiltonian matrix diagonal solver + DiagH* gpu_diagh = nullptr; }; diff --git a/source/module_hsolver/hsolver_lcao.cpp b/source/module_hsolver/hsolver_lcao.cpp index 1ca3b9e39f6..d9b4e931737 100644 --- a/source/module_hsolver/hsolver_lcao.cpp +++ b/source/module_hsolver/hsolver_lcao.cpp @@ -9,7 +9,7 @@ namespace hsolver { template -void HSolverLCAO::solveTemplate(hamilt::Hamilt* pHamilt, +void HSolverLCAO::solveTemplate(hamilt::Hamilt* pHamilt, psi::Psi& psi, elecstate::ElecState* pes, const std::string method_in, @@ -128,7 +128,7 @@ void HSolverLCAO::solveTemplate(hamilt::Hamilt* pHamilt, int HSolverLCAO::out_mat_hs = 0; int HSolverLCAO::out_mat_hsR = 0; -void HSolverLCAO::solve(hamilt::Hamilt* pHamilt, +void HSolverLCAO::solve(hamilt::Hamilt* pHamilt, psi::Psi>& psi, elecstate::ElecState* pes, const std::string method_in, @@ -136,7 +136,7 @@ void HSolverLCAO::solve(hamilt::Hamilt* pHamilt, { this->solveTemplate(pHamilt, psi, pes, method, skip_charge); } -void HSolverLCAO::solve(hamilt::Hamilt* pHamilt, +void HSolverLCAO::solve(hamilt::Hamilt* pHamilt, psi::Psi& psi, elecstate::ElecState* pes, const std::string method_in, @@ -145,7 +145,7 @@ void HSolverLCAO::solve(hamilt::Hamilt* pHamilt, this->solveTemplate(pHamilt, psi, pes, method, skip_charge); } -void HSolverLCAO::hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi>& psi, double* eigenvalue) +void HSolverLCAO::hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi>& psi, double* eigenvalue) { ModuleBase::TITLE("HSolverLCAO", "hamiltSolvePsiK"); ModuleBase::timer::tick("HSolverLCAO", "hamiltSolvePsiK"); @@ -153,7 +153,7 @@ void HSolverLCAO::hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi& psi, double* eigenvalue) +void HSolverLCAO::hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi& psi, double* eigenvalue) { ModuleBase::TITLE("HSolverLCAO", "hamiltSolvePsiK"); ModuleBase::timer::tick("HSolverLCAO", "hamiltSolvePsiK"); diff --git a/source/module_hsolver/hsolver_lcao.h b/source/module_hsolver/hsolver_lcao.h index 4b041321560..e88fad62a1f 100644 --- a/source/module_hsolver/hsolver_lcao.h +++ b/source/module_hsolver/hsolver_lcao.h @@ -22,18 +22,18 @@ class HSolverLCAO : public HSolver void update(//Input &in ) override;*/ - void solve(hamilt::Hamilt* pHamilt, psi::Psi>& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) override; + void solve(hamilt::Hamilt* pHamilt, psi::Psi>& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) override; - void solve(hamilt::Hamilt* pHamilt, psi::Psi& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) override; + void solve(hamilt::Hamilt* pHamilt, psi::Psi& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) override; static int out_mat_hs; // mohan add 2010-09-02 static int out_mat_hsR; // LiuXh add 2019-07-16 private: - void hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi>& psi, double* eigenvalue); - void hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi& psi, double* eigenvalue); + void hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi>& psi, double* eigenvalue); + void hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi& psi, double* eigenvalue); - template void solveTemplate(hamilt::Hamilt* pHamilt, psi::Psi& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge); + template void solveTemplate(hamilt::Hamilt* pHamilt, psi::Psi& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge); /*void solveTemplate( hamilt::Hamilt* pHamilt, psi::Psi>& psi, diff --git a/source/module_hsolver/hsolver_pw.cpp b/source/module_hsolver/hsolver_pw.cpp index 205cd5baeea..7b75d6bb05d 100644 --- a/source/module_hsolver/hsolver_pw.cpp +++ b/source/module_hsolver/hsolver_pw.cpp @@ -4,6 +4,7 @@ #include "diago_david.h" #include "module_base/tool_quit.h" #include "module_base/timer.h" +#include "module_hamilt/hamilt_pw.h" #include "module_elecstate/elecstate_pw.h" #include "src_pw/global.h" #include @@ -44,6 +45,11 @@ void HSolverPW::initDiagh() { pdiagh = new DiagoCG(precondition.data()); pdiagh->method = this->method; + // temperary added for debugging! + #if defined(__CUDA) || defined(__ROCM) + gpu_diagh = new DiagoCG(precondition.data()); + gpu_diagh->method = this->method; + #endif } } else if (this->method == "dav") @@ -70,7 +76,7 @@ void HSolverPW::initDiagh() } } -void HSolverPW::solve(hamilt::Hamilt* pHamilt, psi::Psi>& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) +void HSolverPW::solve(hamilt::Hamilt* pHamilt, psi::Psi>& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) { ModuleBase::TITLE("HSolverPW", "solve"); ModuleBase::timer::tick("HSolverPW", "solve"); @@ -119,6 +125,10 @@ void HSolverPW::endDiagh() { delete (DiagoCG*)pdiagh; pdiagh = nullptr; + #if defined(__CUDA) || defined(__ROCM) + gpu_diagh = nullptr; + delete (DiagoCG*)gpu_diagh; + #endif } if(this->method == "dav") { @@ -132,7 +142,7 @@ void HSolverPW::endDiagh() } } -void HSolverPW::updatePsiK(hamilt::Hamilt* pHamilt, psi::Psi>& psi, const int ik) +void HSolverPW::updatePsiK(hamilt::Hamilt* pHamilt, psi::Psi>& psi, const int ik) { psi.fix_k(ik); if(!this->initialed_psi) @@ -151,9 +161,43 @@ void HSolverPW::updatePsiK(hamilt::Hamilt* pHamilt, psi::Psi>& psi, double* eigenvalue) +void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi>& psi, double* eigenvalue) { + /// a huge victory here! + /// psi::Psi, psi::DEVICE_GPU> gpu_psi = psi; + /// psi::Psi, psi::DEVICE_CPU> cpu_psi = gpu_psi; + /// pdiagh->diag(hm, psi, eigenvalue); + /// psi::memory::synchronize_memory_op, psi::DEVICE_CPU, psi::DEVICE_CPU>()( + /// psi.get_device(), + /// cpu_psi.get_device(), + /// psi.get_pointer() - psi.get_psi_bias(), + /// cpu_psi.get_pointer() - cpu_psi.get_psi_bias(), + /// psi.size()); + + + /// hamilt::Hamilt* h_phm_in = + /// new hamilt::HamiltPW( + /// reinterpret_cast*>(d_phm_in)); + /// + /// pdiagh->diag(h_phm_in, psi, eigenvalue); + /// new era + /// if this works, then everything done! + +#if defined(__CUDA) || defined(__ROCM) + psi::Psi, psi::DEVICE_GPU> gpu_psi = psi; + hamilt::Hamilt* d_phm_in = + new hamilt::HamiltPW( + reinterpret_cast*>(hm)); + gpu_diagh->diag(d_phm_in, gpu_psi, eigenvalue); + psi::memory::synchronize_memory_op, psi::DEVICE_CPU, psi::DEVICE_GPU>()( + psi.get_device(), + gpu_psi.get_device(), + psi.get_pointer() - psi.get_psi_bias(), + gpu_psi.get_pointer() - gpu_psi.get_psi_bias(), + psi.size()); +#else pdiagh->diag(hm, psi, eigenvalue); +#endif } void HSolverPW::update_precondition(std::vector &h_diag, const int ik, const int npw) diff --git a/source/module_hsolver/hsolver_pw.h b/source/module_hsolver/hsolver_pw.h index 51aaf971e42..61e250518c7 100644 --- a/source/module_hsolver/hsolver_pw.h +++ b/source/module_hsolver/hsolver_pw.h @@ -19,7 +19,7 @@ class HSolverPW : public HSolver void update(//Input &in ) override;*/ - void solve(hamilt::Hamilt* pHamilt, psi::Psi>& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) override; + void solve(hamilt::Hamilt* pHamilt, psi::Psi>& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) override; virtual double cal_hsolerror() override; virtual double set_diagethr(const int istep, const int iter, const double drho) override; @@ -27,9 +27,9 @@ class HSolverPW : public HSolver protected: void initDiagh(); void endDiagh(); - void hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi>& psi, double* eigenvalue); + void hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi>& psi, double* eigenvalue); - void updatePsiK(hamilt::Hamilt* pHamilt, psi::Psi>& psi, const int ik); + void updatePsiK(hamilt::Hamilt* pHamilt, psi::Psi>& psi, const int ik); ModulePW::PW_Basis_K* wfc_basis = nullptr; diff --git a/source/module_hsolver/hsolver_pw_sdft.cpp b/source/module_hsolver/hsolver_pw_sdft.cpp index 970f94c3df4..26725b541ee 100644 --- a/source/module_hsolver/hsolver_pw_sdft.cpp +++ b/source/module_hsolver/hsolver_pw_sdft.cpp @@ -10,7 +10,7 @@ #include "src_pw/global.h" namespace hsolver { - void HSolverPW_SDFT::solve(hamilt::Hamilt* pHamilt, + void HSolverPW_SDFT::solve(hamilt::Hamilt* pHamilt, psi::Psi>& psi, elecstate::ElecState* pes, Stochastic_WF& stowf, diff --git a/source/module_hsolver/hsolver_pw_sdft.h b/source/module_hsolver/hsolver_pw_sdft.h index 908ce4b7b5f..333d51e47f1 100644 --- a/source/module_hsolver/hsolver_pw_sdft.h +++ b/source/module_hsolver/hsolver_pw_sdft.h @@ -12,7 +12,7 @@ namespace hsolver this->classname = "HSolverPW_SDFT"; stoiter.init(wfc_basis_in->npwk_max, stowf.nchip, method_sto, stowf); } - virtual void solve(hamilt::Hamilt* pHamilt, + virtual void solve(hamilt::Hamilt* pHamilt, psi::Psi>& psi, elecstate::ElecState* pes, Stochastic_WF& stowf, diff --git a/source/module_hsolver/src/cuda/math_kernel.cu b/source/module_hsolver/src/cuda/math_kernel.cu index c87ba84bfc3..24ea1829a0b 100644 --- a/source/module_hsolver/src/cuda/math_kernel.cu +++ b/source/module_hsolver/src/cuda/math_kernel.cu @@ -9,15 +9,20 @@ namespace hsolver { -static cublasHandle_t diag_handle; +static cublasHandle_t diag_handle = nullptr; void createBLAShandle(){ - cublasCreate(&diag_handle); + if (diag_handle == nullptr) { + cublasCreate(&diag_handle); + } } void destoryBLAShandle(){ - cublasDestroy(diag_handle); + if (diag_handle != nullptr) { + cublasDestroy(diag_handle); + diag_handle = nullptr; + } } // Define the CUDA kernel: diff --git a/source/module_hsolver/test/diago_cg_test.cpp b/source/module_hsolver/test/diago_cg_test.cpp index 4538de20a63..d11286d2846 100644 --- a/source/module_hsolver/test/diago_cg_test.cpp +++ b/source/module_hsolver/test/diago_cg_test.cpp @@ -95,8 +95,8 @@ class DiagoCGPrepare //====================================================================== double *en = new double[npw]; int ik = 1; - hamilt::Hamilt* ha; - ha =new hamilt::HamiltPW(); + hamilt::Hamilt* ha; + ha =new hamilt::HamiltPW(); int* ngk = new int [1]; //psi::Psi> psi(ngk,ik,nband,npw); psi::Psi> psi; diff --git a/source/module_hsolver/test/diago_david_test.cpp b/source/module_hsolver/test/diago_david_test.cpp index 2d8cf624691..ad7dd88a150 100644 --- a/source/module_hsolver/test/diago_david_test.cpp +++ b/source/module_hsolver/test/diago_david_test.cpp @@ -80,8 +80,8 @@ class DiagoDavPrepare //do Diago_David::diag() double* en = new double[npw]; - hamilt::Hamilt *phm; - phm = new hamilt::HamiltPW(); + hamilt::Hamilt *phm; + phm = new hamilt::HamiltPW(); hsolver::DiagoDavid dav(precondition); hsolver::DiagoDavid::PW_DIAG_NDIM = order; hsolver::DiagoIterAssist::PW_DIAG_NMAX = maxiter; diff --git a/source/module_hsolver/test/diago_lcao_test.cpp b/source/module_hsolver/test/diago_lcao_test.cpp index 068213e215f..da50dd6a1fc 100644 --- a/source/module_hsolver/test/diago_lcao_test.cpp +++ b/source/module_hsolver/test/diago_lcao_test.cpp @@ -24,7 +24,7 @@ * self-realized functions in src_pdiag/test/diago_elpa_utils.h */ -template class HamiltTEST : public hamilt::Hamilt +template class HamiltTEST : public hamilt::Hamilt { public: int desc[9]; @@ -82,7 +82,7 @@ template class DiagoPrepare std::vector s; std::vector h_local; std::vector s_local; - hsolver::DiagH* dh=0; + hsolver::DiagH* dh=0; psi::Psi psi; std::vector e_solver; std::vector e_lapack; diff --git a/source/module_hsolver/test/diago_mock.h b/source/module_hsolver/test/diago_mock.h index 28ac25f0aec..0a6be96403a 100644 --- a/source/module_hsolver/test/diago_mock.h +++ b/source/module_hsolver/test/diago_mock.h @@ -232,7 +232,7 @@ class HPsi }; //totally same as the original function -void hamilt::HamiltPW::sPsi +template<> void hamilt::HamiltPW::sPsi ( const std::complex *psi, std::complex *spsi, @@ -284,17 +284,17 @@ class OperatorMock : public hamilt::OperatorPW } }; -void hamilt::HamiltPW::updateHk(const int ik) +template<> void hamilt::HamiltPW::updateHk(const int ik) { return; } -hamilt::HamiltPW::HamiltPW() +template<> hamilt::HamiltPW::HamiltPW() { this->ops = new OperatorMock; } -hamilt::HamiltPW::~HamiltPW() +template<> hamilt::HamiltPW::~HamiltPW() { delete this->ops; } diff --git a/source/module_psi/psi.cpp b/source/module_psi/psi.cpp index 12c5a549c62..eff7e97b56f 100644 --- a/source/module_psi/psi.cpp +++ b/source/module_psi/psi.cpp @@ -6,7 +6,7 @@ #include -using namespace psi; +namespace psi { Range::Range(const size_t range_in) { @@ -121,9 +121,8 @@ Psi::Psi( } template -template Psi::Psi( - const Psi& psi_in) + const Psi& psi_in) { this->ngk = psi_in.get_ngk_pointer(); this->npol = psi_in.npol; @@ -132,17 +131,45 @@ Psi::Psi( this->nbasis = psi_in.get_nbasis(); this->current_k = psi_in.get_current_k(); this->current_b = psi_in.get_current_b(); - this->current_nbasis = psi_in.get_current_nbas(); 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->device = device::get_device_type(this->ctx); this->resize(psi_in.get_nk(), psi_in.get_nbands(), psi_in.get_nbasis()); - memory::synchronize_memory_op()( + memory::synchronize_memory_op()( this->ctx, psi_in.get_device(), this->psi, - psi_in.get_pointer(), + 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(); +} + +template +template +Psi::Psi( + const Psi& psi_in) +{ + this->ngk = psi_in.get_ngk_pointer(); + this->npol = psi_in.npol; + this->nk = psi_in.get_nk(); + this->nbands = psi_in.get_nbands(); + this->nbasis = psi_in.get_nbasis(); + 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->device = device::get_device_type(this->ctx); + this->resize(psi_in.get_nk(), psi_in.get_nbands(), psi_in.get_nbasis()); + memory::synchronize_memory_op()( + this->ctx, + psi_in.get_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(); } @@ -275,11 +302,11 @@ T& Psi::operator()(const int ik, const int ibands, const int ibasis) assert(ik >= 0 && ik < this->nk); assert(ibands >= 0 && ibands < this->nbands); assert(ibasis >= 0 && ibasis < this->nbasis); - if (this->device != CpuDevice) { - #if ((defined __CUDA) || (defined __ROCM)) - ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - #endif - } + // if (this->device != CpuDevice) { + // #if ((defined __CUDA) || (defined __ROCM)) + // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); + // #endif + // } return this->psi[(ik * this->nbands + ibands) * this->nbasis + ibasis]; } @@ -289,11 +316,11 @@ const T& Psi::operator()(const int ik, const int ibands, const int ib assert(ik >= 0 && ik < this->nk); assert(ibands >= 0 && ibands < this->nbands); assert(ibasis >= 0 && ibasis < this->nbasis); - if (this->device != CpuDevice) { - #if ((defined __CUDA) || (defined __ROCM)) - ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - #endif - } + // if (this->device != CpuDevice) { + // #if ((defined __CUDA) || (defined __ROCM)) + // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); + // #endif + // } return this->psi[(ik * this->nbands + ibands) * this->nbasis + ibasis]; } @@ -303,11 +330,11 @@ T& Psi::operator()(const int ibands, const int ibasis) assert(this->current_b == 0); assert(ibands >= 0 && ibands < this->nbands); assert(ibasis >= 0 && ibasis < this->nbasis); - if (this->device != CpuDevice) { - #if ((defined __CUDA) || (defined __ROCM)) - ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - #endif - } + // if (this->device != CpuDevice) { + // #if ((defined __CUDA) || (defined __ROCM)) + // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); + // #endif + // } return this->psi_current[ibands * this->nbasis + ibasis]; } @@ -317,11 +344,11 @@ const T& Psi::operator()(const int ibands, const int ibasis) const assert(this->current_b==0); assert(ibands >= 0 && ibands < this->nbands); assert(ibasis >= 0 && ibasis < this->nbasis); - if (this->device != CpuDevice) { - #if ((defined __CUDA) || (defined __ROCM)) - ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - #endif - } + // if (this->device != CpuDevice) { + // #if ((defined __CUDA) || (defined __ROCM)) + // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); + // #endif + // } return this->psi_current[ibands * this->nbasis + ibasis]; } @@ -329,11 +356,11 @@ template T& Psi::operator()(const int ibasis) { assert(ibasis >= 0 && ibasis < this->nbasis); - if (this->device != CpuDevice) { - #if ((defined __CUDA) || (defined __ROCM)) - ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - #endif - } + // if (this->device != CpuDevice) { + // #if ((defined __CUDA) || (defined __ROCM)) + // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); + // #endif + // } return this->psi_current[this->current_b * this->nbasis + ibasis]; } @@ -341,11 +368,11 @@ template const T& Psi::operator()(const int ibasis) const { assert(ibasis >= 0 && ibasis < this->nbasis); - if (this->device != CpuDevice) { - #if ((defined __CUDA) || (defined __ROCM)) - ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - #endif - } + // if (this->device != CpuDevice) { + // #if ((defined __CUDA) || (defined __ROCM)) + // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); + // #endif + // } return this->psi_current[this->current_b * this->nbasis + ibasis]; } @@ -418,19 +445,14 @@ void initialize(Psi &psi) } -namespace psi { -template class Psi, DEVICE_CPU>; template class Psi; -template Psi, DEVICE_CPU>::Psi, DEVICE_CPU>(const Psi, DEVICE_CPU>&); -template Psi::Psi(const Psi&); +template class Psi, DEVICE_CPU>; #if ((defined __CUDA) || (defined __ROCM)) -template class Psi, DEVICE_GPU>; template class Psi; -template Psi, DEVICE_GPU>::Psi, DEVICE_GPU>(const Psi, DEVICE_GPU>&); -template Psi::Psi(const Psi&); -template Psi, DEVICE_GPU>::Psi, DEVICE_CPU>(const Psi, DEVICE_CPU>&); -template Psi::Psi(const Psi&); -template Psi, DEVICE_CPU>::Psi, DEVICE_GPU>(const Psi, DEVICE_GPU>&); -template Psi::Psi(const Psi&); +template class Psi, DEVICE_GPU>; +template Psi::Psi(const Psi&); +template Psi::Psi(const Psi&); +template Psi, DEVICE_CPU>::Psi(const Psi, DEVICE_GPU>&); +template Psi, DEVICE_GPU>::Psi(const Psi, DEVICE_CPU>&); #endif } // namespace psi \ No newline at end of file diff --git a/source/module_psi/psi.h b/source/module_psi/psi.h index 188bf264bc9..4f77f4fdfa3 100644 --- a/source/module_psi/psi.h +++ b/source/module_psi/psi.h @@ -47,6 +47,8 @@ class Psi //in this case, fix_k can not be used Psi(T* psi_pointer, const Psi& psi_in, const int nk_in, int nband_in=0); //Constructor 6: initialize a new psi from the given psi_in + Psi(const Psi& psi_in); + //Constructor 7: initialize a new psi from the given psi_in with a different class template //in this case, psi_in may have a different device type. template Psi(const Psi& psi_in); //Destructor for deleting the psi array manually diff --git a/source/src_lcao/ELEC_evolve.cpp b/source/src_lcao/ELEC_evolve.cpp index 62991777958..68fdcfde509 100644 --- a/source/src_lcao/ELEC_evolve.cpp +++ b/source/src_lcao/ELEC_evolve.cpp @@ -27,7 +27,7 @@ int ELEC_evolve::td_dipoleout; // this routine only serves for TDDFT using LCAO basis set void ELEC_evolve::evolve_psi(const int& istep, - hamilt::Hamilt* phm, + hamilt::Hamilt* phm, Local_Orbital_wfc& lowf, psi::Psi>* psi, psi::Psi>* psi_laststep, diff --git a/source/src_lcao/ELEC_evolve.h b/source/src_lcao/ELEC_evolve.h index 8fc70cb53af..c68568f31b6 100644 --- a/source/src_lcao/ELEC_evolve.h +++ b/source/src_lcao/ELEC_evolve.h @@ -44,7 +44,7 @@ class ELEC_evolve private: static void evolve_psi(const int& istep, - hamilt::Hamilt* phm, + hamilt::Hamilt* phm, Local_Orbital_wfc& lowf, psi::Psi>* psi, psi::Psi>* psi_laststep, diff --git a/source/src_lcao/LCAO_evolve.cpp b/source/src_lcao/LCAO_evolve.cpp index 690ff0406b9..d9b593b8d41 100644 --- a/source/src_lcao/LCAO_evolve.cpp +++ b/source/src_lcao/LCAO_evolve.cpp @@ -23,7 +23,7 @@ inline int globalIndex(int localindex, int nblk, int nprocs, int myproc) } void Evolve_LCAO_Matrix::evolve_complex_matrix(const int& ik, - hamilt::Hamilt* p_hamilt, + hamilt::Hamilt* p_hamilt, psi::Psi>* psi_k, psi::Psi>* psi_k_laststep, double* ekb) const @@ -60,7 +60,7 @@ void Evolve_LCAO_Matrix::evolve_complex_matrix(const int& ik, } void Evolve_LCAO_Matrix::using_LAPACK_complex(const int& ik, - hamilt::Hamilt* p_hamilt, + hamilt::Hamilt* p_hamilt, std::complex* psi_k, std::complex* psi_k_laststep, double* ekb) const @@ -346,7 +346,7 @@ void Evolve_LCAO_Matrix::using_LAPACK_complex(const int& ik, #ifdef __MPI void Evolve_LCAO_Matrix::using_ScaLAPACK_complex(const int& ik, - hamilt::Hamilt* p_hamilt, + hamilt::Hamilt* p_hamilt, std::complex* psi_k, std::complex* psi_k_laststep, double* ekb) const diff --git a/source/src_lcao/LCAO_evolve.h b/source/src_lcao/LCAO_evolve.h index e2f00845916..962bf4c12fa 100644 --- a/source/src_lcao/LCAO_evolve.h +++ b/source/src_lcao/LCAO_evolve.h @@ -19,7 +19,7 @@ class Evolve_LCAO_Matrix ~Evolve_LCAO_Matrix(); void evolve_complex_matrix(const int& ik, - hamilt::Hamilt* p_hamilt, + hamilt::Hamilt* p_hamilt, psi::Psi>* psi_k, psi::Psi>* psi_k_laststep, double* ekb) const; @@ -29,13 +29,13 @@ class Evolve_LCAO_Matrix const Parallel_Orbitals* ParaV; void using_LAPACK_complex(const int& ik, - hamilt::Hamilt* p_hamilt, + hamilt::Hamilt* p_hamilt, std::complex* psi_k, std::complex* psi_k_laststep, double* ekb) const; #ifdef __MPI void using_ScaLAPACK_complex(const int& ik, - hamilt::Hamilt* p_hamilt, + hamilt::Hamilt* p_hamilt, std::complex* psi_k, std::complex* psi_k_laststep, double* ekb) const; diff --git a/source/src_pw/sto_iter.cpp b/source/src_pw/sto_iter.cpp index ed6e29e88e4..ec1425ad86d 100644 --- a/source/src_pw/sto_iter.cpp +++ b/source/src_pw/sto_iter.cpp @@ -426,7 +426,7 @@ void Stochastic_Iter::calHsqrtchi(Stochastic_WF& stowf) } } -void Stochastic_Iter::sum_stoband(Stochastic_WF& stowf, elecstate::ElecState* pes,hamilt::Hamilt* pHamilt) +void Stochastic_Iter::sum_stoband(Stochastic_WF& stowf, elecstate::ElecState* pes,hamilt::Hamilt* pHamilt) { ModuleBase::TITLE("Stochastic_Iter","sum_stoband"); ModuleBase::timer::tick("Stochastic_Iter","sum_stoband"); diff --git a/source/src_pw/sto_iter.h b/source/src_pw/sto_iter.h index 65f79a4e48d..4c0e1ab6d98 100644 --- a/source/src_pw/sto_iter.h +++ b/source/src_pw/sto_iter.h @@ -27,7 +27,7 @@ class Stochastic_Iter void init(const int, int* nchip_in, const int method_in, Stochastic_WF& stowf); - void sum_stoband(Stochastic_WF& stowf, elecstate::ElecState* pes,hamilt::Hamilt* pHamilt); + void sum_stoband(Stochastic_WF& stowf, elecstate::ElecState* pes,hamilt::Hamilt* pHamilt); double calne(elecstate::ElecState* pes); diff --git a/source/src_pw/wavefunc.cpp b/source/src_pw/wavefunc.cpp index 9a653a4d033..6b5540d3bbb 100644 --- a/source/src_pw/wavefunc.cpp +++ b/source/src_pw/wavefunc.cpp @@ -228,7 +228,7 @@ void wavefunc::diago_PAO_in_pw_k(const int &ik, psi::Psi> & return; } -void wavefunc::diago_PAO_in_pw_k2(const int &ik, psi::Psi> &wvf, hamilt::Hamilt* phm_in) +void wavefunc::diago_PAO_in_pw_k2(const int &ik, psi::Psi> &wvf, hamilt::Hamilt* phm_in) { ModuleBase::TITLE("wavefunc","diago_PAO_in_pw_k2"); // (6) Prepare for atmoic orbitals or random orbitals diff --git a/source/src_pw/wavefunc.h b/source/src_pw/wavefunc.h index 66cb4f9a258..bee638e84dc 100644 --- a/source/src_pw/wavefunc.h +++ b/source/src_pw/wavefunc.h @@ -44,7 +44,7 @@ class wavefunc : public WF_atomic // used if k dependent staff is ready. void prepare_k(void); - void diago_PAO_in_pw_k2(const int &ik, psi::Psi> &wvf, hamilt::Hamilt* phm_in = nullptr); + void diago_PAO_in_pw_k2(const int &ik, psi::Psi> &wvf, hamilt::Hamilt* phm_in = nullptr); void diago_PAO_in_pw_k2(const int &ik, ModuleBase::ComplexMatrix &wvf); int get_R(int ix, int iy, int iz); // pengfei 2016-11-23 From 934496b15030344093a56c9a4efdc42fa20a7d6c Mon Sep 17 00:00:00 2001 From: denghuilu Date: Sun, 6 Nov 2022 11:18:51 +0800 Subject: [PATCH 2/5] track changes --- source/module_hsolver/diago_cg.cpp | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/source/module_hsolver/diago_cg.cpp b/source/module_hsolver/diago_cg.cpp index 9cd87390f6c..21fc2a575cd 100644 --- a/source/module_hsolver/diago_cg.cpp +++ b/source/module_hsolver/diago_cg.cpp @@ -222,8 +222,6 @@ void DiagoCG::calculate_gradient() if (this->test_cg == 1) { ModuleBase::TITLE("DiagoCG", "calculate_gradient"); } - // ModuleBase::timer::tick("DiagoCG","grad"); - // for (int i = 0; i < this->dim; i++) // { // //(2) PH|psi> @@ -231,7 +229,8 @@ void DiagoCG::calculate_gradient() // //(3) PS|psi> // this->pphi[i] = this->sphi[i] / this->precondition[i]; // } - // haozhihan replace this 2022-10-6 + // denghui replace this at 20221106 + // TODO: use GPU precondition to initialize CG class if (this->device == psi::GpuDevice) { vector_div_vector_op()(this->ctx, this->dim, this->gradient, this->hphi, this->d_precondition); vector_div_vector_op()(this->ctx, this->dim, this->pphi, this->sphi, this->d_precondition); @@ -241,7 +240,6 @@ void DiagoCG::calculate_gradient() vector_div_vector_op()(this->ctx, this->dim, this->pphi, this->sphi, this->precondition); } - // Update lambda ! // (4) const FPTYPE eh = hsolver::zdot_real_op()(this->ctx, this->dim, this->sphi, this->gradient); @@ -261,9 +259,6 @@ void DiagoCG::calculate_gradient() // } // haozhihan replace this 2022-10-6 constantvector_addORsub_constantVector_op()(this->ctx, this->dim, this->gradient, this->gradient, 1.0, this->pphi, (-lambda)); - - - // ModuleBase::timer::tick("DiagoCG","grad"); } template @@ -380,9 +375,6 @@ void DiagoCG::orthogonal_gradient(hamilt::Hamilt } }*/ //>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - - // ModuleBase::timer::tick("DiagoCG","orth_grad"); - return; } template @@ -391,7 +383,6 @@ void DiagoCG::calculate_gamma_cg(const int iter, FPTYPE &gg_last if (test_cg == 1) { ModuleBase::TITLE("DiagoCG", "calculate_gamma_cg"); } - // ModuleBase::timer::tick("DiagoCG","gamma_cg"); auto pcg = this->cg->get_pointer(); auto pphi_m = this->phi_m->get_pointer(); FPTYPE gg_inter; @@ -413,7 +404,8 @@ void DiagoCG::calculate_gamma_cg(const int iter, FPTYPE &gg_last // { // this->g0[i] = this->precondition[i] * this->scg[i]; // } - // haozhihan replace this 2022-10-6 + // denghui replace this 20221106 + // TODO: use GPU precondition instead if (this->device == psi::GpuDevice) { vector_mul_vector_op()(this->ctx, this->dim, this->g0, this->scg, this->d_precondition); } @@ -431,7 +423,6 @@ void DiagoCG::calculate_gamma_cg(const int iter, FPTYPE &gg_last gg_last = gg_now; // (50) cg direction first value : |g> // |cg> = |g> - // haozhihan replace COPY_ARRAY psi::memory::synchronize_memory_op, Device, Device>()(this->ctx, this->ctx, pcg, this->gradient, this->dim); // ModuleBase::GlobalFunc::COPYARRAY(this->gradient, pcg, this->dim); @@ -465,7 +456,6 @@ void DiagoCG::calculate_gamma_cg(const int iter, FPTYPE &gg_last pcg[i] -= norma * pphi_m[i]; }*/ } - // ModuleBase::timer::tick("DiagoCG","gamma_cg"); } template @@ -473,7 +463,6 @@ bool DiagoCG::update_psi(FPTYPE &cg_norm, FPTYPE &theta, FPTYPE { if (test_cg == 1) ModuleBase::TITLE("DiagoCG", "update_psi"); - // ModuleBase::timer::tick("DiagoCG","update"); cg_norm = sqrt(hsolver::zdot_real_op()(this->ctx, this->dim, this->cg->get_pointer(), this->scg)); if (cg_norm < 1.0e-10) From d8a0eadd62ffee4b9392296789fd63d77aeb7dac Mon Sep 17 00:00:00 2001 From: denghuilu Date: Sat, 12 Nov 2022 18:34:18 +0800 Subject: [PATCH 3/5] fix memory leak --- source/module_hamilt/ks_pw/nonlocal_pw.h | 2 +- source/module_hamilt/ks_pw/veff_pw.cpp | 1 + source/module_hamilt/ks_pw/veff_pw.h | 1 + source/module_hsolver/diago_cg.cpp | 1 + source/module_hsolver/hsolver_pw.cpp | 3 ++- 5 files changed, 6 insertions(+), 2 deletions(-) diff --git a/source/module_hamilt/ks_pw/nonlocal_pw.h b/source/module_hamilt/ks_pw/nonlocal_pw.h index 5e645adb060..820cabd6d40 100644 --- a/source/module_hamilt/ks_pw/nonlocal_pw.h +++ b/source/module_hamilt/ks_pw/nonlocal_pw.h @@ -43,7 +43,7 @@ class Nonlocal> : public OperatorPW const int *get_isk() const {return this->isk;} const pseudopot_cell_vnl *get_ppcell() const {return this->ppcell;} - const UnitCell_pseudo *get_ucell() const {return this->ucell;} + const UnitCell *get_ucell() const {return this->ucell;} private: void add_nonlocal_pp(std::complex *hpsi_in, const std::complex *becp, const int m) const; diff --git a/source/module_hamilt/ks_pw/veff_pw.cpp b/source/module_hamilt/ks_pw/veff_pw.cpp index 86db829e695..113a02c770e 100644 --- a/source/module_hamilt/ks_pw/veff_pw.cpp +++ b/source/module_hamilt/ks_pw/veff_pw.cpp @@ -145,6 +145,7 @@ hamilt::Veff>::Veff(const Veffveff_col = veff->get_veff_col(); this->veff_row = veff->get_veff_row(); this->wfcpw = veff->get_wfcpw(); + this->npol = veff->get_npol(); resize_memory_complex_op()(this->ctx, this->porter, this->wfcpw->nmaxgr); if (this->npol != 1) { resize_memory_complex_op()(this->ctx, this->porter1, this->wfcpw->nmaxgr); diff --git a/source/module_hamilt/ks_pw/veff_pw.h b/source/module_hamilt/ks_pw/veff_pw.h index 2548f4ded92..83fa12b1ce5 100644 --- a/source/module_hamilt/ks_pw/veff_pw.h +++ b/source/module_hamilt/ks_pw/veff_pw.h @@ -40,6 +40,7 @@ class Veff> : public OperatorPW FPTYPE *get_veff() const {return this->veff;} int get_veff_col() const {return this->veff_col;} int get_veff_row() const {return this->veff_row;} + int get_npol() const {return this->npol;} const int *get_isk() const {return isk;} ModulePW::PW_Basis_K* get_wfcpw() const {return this->wfcpw;} diff --git a/source/module_hsolver/diago_cg.cpp b/source/module_hsolver/diago_cg.cpp index 21fc2a575cd..1f2a39698ca 100644 --- a/source/module_hsolver/diago_cg.cpp +++ b/source/module_hsolver/diago_cg.cpp @@ -696,6 +696,7 @@ void DiagoCG::diag(hamilt::Hamilt *phm_in, psi:: psi.get_pointer() - psi.get_psi_bias(), cpu_psi.get_pointer() - cpu_psi.get_psi_bias(), psi.size()); + delete reinterpret_cast*>(h_phm_in); #else DiagoIterAssist::diagH_subspace(phm_in, psi, psi, eigenvalue_in); #endif diff --git a/source/module_hsolver/hsolver_pw.cpp b/source/module_hsolver/hsolver_pw.cpp index 7b75d6bb05d..ec5b3147238 100644 --- a/source/module_hsolver/hsolver_pw.cpp +++ b/source/module_hsolver/hsolver_pw.cpp @@ -126,8 +126,8 @@ void HSolverPW::endDiagh() delete (DiagoCG*)pdiagh; pdiagh = nullptr; #if defined(__CUDA) || defined(__ROCM) - gpu_diagh = nullptr; delete (DiagoCG*)gpu_diagh; + gpu_diagh = nullptr; #endif } if(this->method == "dav") @@ -195,6 +195,7 @@ void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi*>(d_phm_in); #else pdiagh->diag(hm, psi, eigenvalue); #endif From a35a095b9d9cc67eb0888954933fde390df3d3f6 Mon Sep 17 00:00:00 2001 From: denghuilu Date: Sat, 12 Nov 2022 21:33:26 +0800 Subject: [PATCH 4/5] ready for merge --- source/module_hamilt/ks_pw/veff_pw.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/source/module_hamilt/ks_pw/veff_pw.cpp b/source/module_hamilt/ks_pw/veff_pw.cpp index 113a02c770e..8a2ad03166b 100644 --- a/source/module_hamilt/ks_pw/veff_pw.cpp +++ b/source/module_hamilt/ks_pw/veff_pw.cpp @@ -146,6 +146,7 @@ hamilt::Veff>::Veff(const Veffveff_row = veff->get_veff_row(); this->wfcpw = veff->get_wfcpw(); this->npol = veff->get_npol(); + this->device = psi::device::get_device_type(this->ctx); resize_memory_complex_op()(this->ctx, this->porter, this->wfcpw->nmaxgr); if (this->npol != 1) { resize_memory_complex_op()(this->ctx, this->porter1, this->wfcpw->nmaxgr); From a3404f5ea320569c397b038df55799bf399ac9b6 Mon Sep 17 00:00:00 2001 From: denghuilu Date: Mon, 14 Nov 2022 14:26:23 +0800 Subject: [PATCH 5/5] address comments --- source/module_hsolver/include/math_kernel.h | 44 +++++++++++++++++++ source/module_hsolver/src/cuda/math_kernel.cu | 29 ++++++------ source/module_psi/psi.cpp | 30 ------------- 3 files changed, 58 insertions(+), 45 deletions(-) diff --git a/source/module_hsolver/include/math_kernel.h b/source/module_hsolver/include/math_kernel.h index f81c1ac0c93..d343a0f3dcf 100644 --- a/source/module_hsolver/include/math_kernel.h +++ b/source/module_hsolver/include/math_kernel.h @@ -7,6 +7,50 @@ #include "module_psi/psi.h" #include "src_parallel/parallel_reduce.h" + +#if defined(__CUDA) || defined(__UT_USE_CUDA) +#include +#include "cublas_v2.h" + +#define cublasErrcheck(res) { cudaAssert((res), __FILE__, __LINE__); } + +static const char *_cudaGetErrorEnum(cublasStatus_t error) { + switch (error) { + case CUBLAS_STATUS_SUCCESS: + return "CUBLAS_STATUS_SUCCESS"; + + case CUBLAS_STATUS_NOT_INITIALIZED: + return "CUBLAS_STATUS_NOT_INITIALIZED"; + + case CUBLAS_STATUS_ALLOC_FAILED: + return "CUBLAS_STATUS_ALLOC_FAILED"; + + case CUBLAS_STATUS_INVALID_VALUE: + return "CUBLAS_STATUS_INVALID_VALUE"; + + case CUBLAS_STATUS_ARCH_MISMATCH: + return "CUBLAS_STATUS_ARCH_MISMATCH"; + + case CUBLAS_STATUS_MAPPING_ERROR: + return "CUBLAS_STATUS_MAPPING_ERROR"; + + case CUBLAS_STATUS_EXECUTION_FAILED: + return "CUBLAS_STATUS_EXECUTION_FAILED"; + + case CUBLAS_STATUS_INTERNAL_ERROR: + return "CUBLAS_STATUS_INTERNAL_ERROR"; + } + return ""; +} + +inline void cudaAssert(cublasStatus_t code, const char *file, int line, bool abort=true) { + if (code != CUBLAS_STATUS_SUCCESS) { + fprintf(stderr,"cuBLAS Assert: %s %s %d\n", _cudaGetErrorEnum(code), file, line); + if (abort) exit(code); + } +} +#endif //__CUDA || __UT_USE_CUDA + namespace hsolver { diff --git a/source/module_hsolver/src/cuda/math_kernel.cu b/source/module_hsolver/src/cuda/math_kernel.cu index 24ea1829a0b..3f8c04214c9 100644 --- a/source/module_hsolver/src/cuda/math_kernel.cu +++ b/source/module_hsolver/src/cuda/math_kernel.cu @@ -4,7 +4,6 @@ #include #include #include -#include "cublas_v2.h" namespace hsolver { @@ -14,13 +13,13 @@ static cublasHandle_t diag_handle = nullptr; void createBLAShandle(){ if (diag_handle == nullptr) { - cublasCreate(&diag_handle); + cublasErrcheck(cublasCreate(&diag_handle)); } } void destoryBLAShandle(){ if (diag_handle != nullptr) { - cublasDestroy(diag_handle); + cublasErrcheck(cublasDestroy(diag_handle)); diag_handle = nullptr; } } @@ -190,7 +189,7 @@ void axpy_op::operator()( std::complex *Y, const int& incY) { - cublasCaxpy(diag_handle, N, (float2*)alpha, (float2*)X, incX, (float2*)Y, incY); + cublasErrcheck(cublasCaxpy(diag_handle, N, (float2*)alpha, (float2*)X, incX, (float2*)Y, incY)); } template <> @@ -203,7 +202,7 @@ void axpy_op::operator()( std::complex *Y, const int& incY) { - cublasZaxpy(diag_handle, N, (double2*)alpha, (double2*)X, incX, (double2*)Y, incY); + cublasErrcheck(cublasZaxpy(diag_handle, N, (double2*)alpha, (double2*)X, incX, (double2*)Y, incY)); } template <> @@ -230,8 +229,8 @@ void gemv_op::operator()( } else if (trans == 'C'){ cutrans = CUBLAS_OP_C; - } - cublasCgemv(diag_handle, cutrans, m, n, (float2*)alpha, (float2*)A, lda, (float2*)X, incx, (float2*)beta, (float2*)Y, incx); + } + cublasErrcheck(cublasCgemv(diag_handle, cutrans, m, n, (float2*)alpha, (float2*)A, lda, (float2*)X, incx, (float2*)beta, (float2*)Y, incx)); } template <> @@ -258,8 +257,8 @@ void gemv_op::operator()( } else if (trans == 'C'){ cutrans = CUBLAS_OP_C; - } - cublasZgemv(diag_handle, cutrans, m, n, (double2*)alpha, (double2*)A, lda, (double2*)X, incx, (double2*)beta, (double2*)Y, incx); + } + cublasErrcheck(cublasZgemv(diag_handle, cutrans, m, n, (double2*)alpha, (double2*)A, lda, (double2*)X, incx, (double2*)beta, (double2*)Y, incx)); } @@ -271,7 +270,7 @@ void scal_op::operator()(const psi::DEVICE_GPU* d, std::complex* X, const int& incx) { - cublasZscal(diag_handle, N, (double2*)alpha, (double2*)X, incx); + cublasErrcheck(cublasZscal(diag_handle, N, (double2*)alpha, (double2*)X, incx)); } template <> @@ -282,7 +281,7 @@ void scal_op::operator()(const psi::DEVICE_GPU* d, const int& incx) { - cublasCscal(diag_handle, N, (float2*)alpha, (float2*)X, incx); + cublasErrcheck(cublasCscal(diag_handle, N, (float2*)alpha, (float2*)X, incx)); } @@ -323,8 +322,8 @@ void gemm_op::operator()(const psi::DEVICE_GPU* d, } else if (transb == 'C'){ cutransB = CUBLAS_OP_C; - } - cublasCgemm(diag_handle, cutransA, cutransB, m, n ,k, (float2*)alpha, (float2*)a , lda, (float2*)b, ldb, (float2*)beta, (float2*)c, ldc); + } + cublasErrcheck(cublasCgemm(diag_handle, cutransA, cutransB, m, n ,k, (float2*)alpha, (float2*)a , lda, (float2*)b, ldb, (float2*)beta, (float2*)c, ldc)); } template <> @@ -364,8 +363,8 @@ void gemm_op::operator()(const psi::DEVICE_GPU* d, } else if (transb == 'C'){ cutransB = CUBLAS_OP_C; - } - cublasZgemm(diag_handle, cutransA, cutransB, m, n ,k, (double2*)alpha, (double2*)a , lda, (double2*)b, ldb, (double2*)beta, (double2*)c, ldc); + } + cublasErrcheck(cublasZgemm(diag_handle, cutransA, cutransB, m, n ,k, (double2*)alpha, (double2*)a , lda, (double2*)b, ldb, (double2*)beta, (double2*)c, ldc)); } diff --git a/source/module_psi/psi.cpp b/source/module_psi/psi.cpp index eff7e97b56f..bc0bbcde7fe 100644 --- a/source/module_psi/psi.cpp +++ b/source/module_psi/psi.cpp @@ -302,11 +302,6 @@ T& Psi::operator()(const int ik, const int ibands, const int ibasis) assert(ik >= 0 && ik < this->nk); assert(ibands >= 0 && ibands < this->nbands); assert(ibasis >= 0 && ibasis < this->nbasis); - // if (this->device != CpuDevice) { - // #if ((defined __CUDA) || (defined __ROCM)) - // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - // #endif - // } return this->psi[(ik * this->nbands + ibands) * this->nbasis + ibasis]; } @@ -316,11 +311,6 @@ const T& Psi::operator()(const int ik, const int ibands, const int ib assert(ik >= 0 && ik < this->nk); assert(ibands >= 0 && ibands < this->nbands); assert(ibasis >= 0 && ibasis < this->nbasis); - // if (this->device != CpuDevice) { - // #if ((defined __CUDA) || (defined __ROCM)) - // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - // #endif - // } return this->psi[(ik * this->nbands + ibands) * this->nbasis + ibasis]; } @@ -330,11 +320,6 @@ T& Psi::operator()(const int ibands, const int ibasis) assert(this->current_b == 0); assert(ibands >= 0 && ibands < this->nbands); assert(ibasis >= 0 && ibasis < this->nbasis); - // if (this->device != CpuDevice) { - // #if ((defined __CUDA) || (defined __ROCM)) - // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - // #endif - // } return this->psi_current[ibands * this->nbasis + ibasis]; } @@ -344,11 +329,6 @@ const T& Psi::operator()(const int ibands, const int ibasis) const assert(this->current_b==0); assert(ibands >= 0 && ibands < this->nbands); assert(ibasis >= 0 && ibasis < this->nbasis); - // if (this->device != CpuDevice) { - // #if ((defined __CUDA) || (defined __ROCM)) - // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - // #endif - // } return this->psi_current[ibands * this->nbasis + ibasis]; } @@ -356,11 +336,6 @@ template T& Psi::operator()(const int ibasis) { assert(ibasis >= 0 && ibasis < this->nbasis); - // if (this->device != CpuDevice) { - // #if ((defined __CUDA) || (defined __ROCM)) - // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - // #endif - // } return this->psi_current[this->current_b * this->nbasis + ibasis]; } @@ -368,11 +343,6 @@ template const T& Psi::operator()(const int ibasis) const { assert(ibasis >= 0 && ibasis < this->nbasis); - // if (this->device != CpuDevice) { - // #if ((defined __CUDA) || (defined __ROCM)) - // ModuleBase::WARNING_QUIT("Psi operator ()", "GPU psi cannot fetch value by an overloaded operator ()!"); - // #endif - // } return this->psi_current[this->current_b * this->nbasis + ibasis]; }