diff --git a/.clang-format b/.clang-format index 8970df99c29..5ea3e064416 100644 --- a/.clang-format +++ b/.clang-format @@ -1,49 +1,63 @@ --- Language: Cpp -BasedOnStyle: Microsoft +BasedOnStyle: LLVM -AlwaysBreakTemplateDeclarations: Yes +ColumnLimit: 120 -AllowAllArgumentsOnNextLine: false -AllowAllParametersOfDeclarationOnNextLine: false +#AlwaysBreakTemplateDeclarations: Yes +# +#AllowAllArgumentsOnNextLine: false +#AllowAllParametersOfDeclarationOnNextLine: false BinPackArguments: false BinPackParameters: false - -BreakBeforeBinaryOperators: All -BreakBeforeTernaryOperators: true - -Cpp11BracedListStyle: true -FixNamespaceComments: true - -# InsertBraces: true # only for clang-format version15.0.0 or later - -# About include -IncludeBlocks: Regroup -IncludeIsMainRegex: '([-_](test|unittest))?$' -IncludeIsMainSourceRegex: '' - +PackArguments: + BinPack: UseBreakAfter + BreakAfter: 3 +PackParameters: + BinPack: UseBreakAfter + BreakAfter: 3 +# +#BreakBeforeBinaryOperators: None +#BreakBeforeTernaryOperators: true +BreakBeforeBraces: Allman +BreakAfterOpenBracketFunction: true +AlignAfterOpenBracket: DontAlign +# +#Cpp11BracedListStyle: true +#FixNamespaceComments: true +# +InsertBraces: true # only for clang-format version15.0.0 or later +# +## About include +#IncludeBlocks: Regroup +#IncludeIsMainRegex: '([-_](test|unittest))?$' +#IncludeIsMainSourceRegex: '' +# IndentWidth: 4 -IndentWrappedFunctionNames: true - -# About Point -DerivePointerAlignment: false -PointerAlignment: Left - -ReflowComments: true -SortIncludes: true -SortUsingDeclarations: true - -# About space -SpaceAfterCStyleCast: false -SpaceAfterLogicalNot: false -SpaceBeforeRangeBasedForLoopColon: false -SpaceInEmptyParentheses: false -SpacesInAngles: false -SpacesInCStyleCastParentheses: false -SpacesInContainerLiterals: false -SpacesInParentheses: false -SpacesInSquareBrackets: false +#IndentWrappedFunctionNames: true +# +## About Point +#DerivePointerAlignment: false +#PointerAlignment: Left +PenaltyBreakBeforeMemberAccess: 1000 +PenaltyBreakScopeResolution: 1000 +# +#ReflowComments: true +#SortIncludes: true +#SortUsingDeclarations: true +# +## About space +#SpaceAfterCStyleCast: false +#SpaceAfterLogicalNot: false +#SpaceBeforeRangeBasedForLoopColon: false +#SpaceInEmptyParentheses: false +#SpacesInAngles: false +#SpacesInCStyleCastParentheses: false +#SpacesInContainerLiterals: false +#SpacesInParentheses: false +#SpacesInSquareBrackets: false +# Standard: c++11 TabWidth: 4 UseTab: Never diff --git a/source/source_hsolver/diag_comm_info.h b/source/source_hsolver/diag_comm_info.h index 767bb0bf1d5..c9a9e99d2fb 100644 --- a/source/source_hsolver/diag_comm_info.h +++ b/source/source_hsolver/diag_comm_info.h @@ -15,12 +15,10 @@ struct diag_comm_info const int nproc; #ifndef __MPI - diag_comm_info(const int rank_in, const int nproc_in) : rank(rank_in), nproc(nproc_in) - { - } + diag_comm_info(const int rank_in, const int nproc_in) : rank(rank_in), nproc(nproc_in) {} #else const MPI_Comm comm; - diag_comm_info(const MPI_Comm& comm_in, const int rank_in, const int nproc_in) + diag_comm_info(const MPI_Comm &comm_in, const int rank_in, const int nproc_in) : comm(comm_in), rank(rank_in), nproc(nproc_in) { } diff --git a/source/source_hsolver/diag_const_nums.cpp b/source/source_hsolver/diag_const_nums.cpp index 53949581130..7b589fd8e5a 100644 --- a/source/source_hsolver/diag_const_nums.cpp +++ b/source/source_hsolver/diag_const_nums.cpp @@ -8,8 +8,7 @@ template class const_nums>; template class const_nums>; // Specialize templates to support double types -template <> -const_nums::const_nums() +template <> const_nums::const_nums() { base_device::memory::resize_memory_op()(this->zero, 1); this->zero[0] = 0.0; @@ -20,8 +19,7 @@ const_nums::const_nums() } // Specialize templates to support double types -template <> -const_nums::const_nums() +template <> const_nums::const_nums() { base_device::memory::resize_memory_op()(this->zero, 1); this->zero[0] = 0.0; @@ -32,8 +30,7 @@ const_nums::const_nums() } // Specialized templates to support std:: std::complextypes -template <> -const_nums>::const_nums() +template <> const_nums>::const_nums() { base_device::memory::resize_memory_op, base_device::DEVICE_CPU>()(this->zero, 1); this->zero[0] = std::complex(0.0, 0.0); @@ -44,8 +41,7 @@ const_nums>::const_nums() } // Specialized templates to support std:: complextypes -template <> -const_nums>::const_nums() +template <> const_nums>::const_nums() { base_device::memory::resize_memory_op, base_device::DEVICE_CPU>()(this->zero, 1); this->zero[0] = std::complex(0.0, 0.0); diff --git a/source/source_hsolver/diag_const_nums.h b/source/source_hsolver/diag_const_nums.h index 05520e4e84a..dee5b72026d 100644 --- a/source/source_hsolver/diag_const_nums.h +++ b/source/source_hsolver/diag_const_nums.h @@ -2,14 +2,13 @@ #define DIAG_CONST_NUMS #include "source_base/module_device/memory_op.h" -template -struct const_nums +template struct const_nums { const_nums(); - base_device::DEVICE_CPU* cpu_ctx = {}; - T* zero = nullptr; - T* one = nullptr; - T* neg_one = nullptr; + base_device::DEVICE_CPU *cpu_ctx = {}; + T *zero = nullptr; + T *one = nullptr; + T *neg_one = nullptr; }; #endif \ No newline at end of file diff --git a/source/source_hsolver/diag_hs_para.cpp b/source/source_hsolver/diag_hs_para.cpp index 13f1ea1a619..800734ee9d4 100644 --- a/source/source_hsolver/diag_hs_para.cpp +++ b/source/source_hsolver/diag_hs_para.cpp @@ -15,53 +15,79 @@ namespace hsolver { #ifdef __ELPA -void elpa_diag(MPI_Comm comm, - const int nband, - std::complex* h_local, - std::complex* s_local, - double* ekb, - std::complex* wfc_2d, - Parallel_2D& para2d_local) +void elpa_diag( + MPI_Comm comm, + const int nband, + std::complex *h_local, + std::complex *s_local, + double *ekb, + std::complex *wfc_2d, + Parallel_2D ¶2d_local) { int DecomposedState = 0; - ELPA_Solver es(false, comm, nband, para2d_local.get_row_size(), para2d_local.get_col_size(), para2d_local.desc); - es.generalized_eigenvector(h_local, s_local, DecomposedState, ekb, wfc_2d); + ELPA_Solver es( + false, + comm, + nband, + para2d_local.get_row_size(), + para2d_local.get_col_size(), + para2d_local.desc); + es.generalized_eigenvector( + h_local, + s_local, + DecomposedState, + ekb, + wfc_2d); es.exit(); } -void elpa_diag(MPI_Comm comm, - const int nband, - double* h_local, - double* s_local, - double* ekb, - double* wfc_2d, - Parallel_2D& para2d_local) +void elpa_diag( + MPI_Comm comm, + const int nband, + double *h_local, + double *s_local, + double *ekb, + double *wfc_2d, + Parallel_2D ¶2d_local) { int DecomposedState = 0; - ELPA_Solver es(true, comm, nband, para2d_local.get_row_size(), para2d_local.get_col_size(), para2d_local.desc); - es.generalized_eigenvector(h_local, s_local, DecomposedState, ekb, wfc_2d); + ELPA_Solver es( + true, + comm, + nband, + para2d_local.get_row_size(), + para2d_local.get_col_size(), + para2d_local.desc); + es.generalized_eigenvector( + h_local, + s_local, + DecomposedState, + ekb, + wfc_2d); es.exit(); } -void elpa_diag(MPI_Comm comm, - const int nband, - std::complex* h_local, - std::complex* s_local, - float* ekb, - std::complex* wfc_2d, - Parallel_2D& para2d_local) +void elpa_diag( + MPI_Comm comm, + const int nband, + std::complex *h_local, + std::complex *s_local, + float *ekb, + std::complex *wfc_2d, + Parallel_2D ¶2d_local) { std::cout << "Error: ELPA do not support single precision. " << std::endl; exit(1); } -void elpa_diag(MPI_Comm comm, - const int nband, - float* h_local, - float* s_local, - float* ekb, - float* wfc_2d, - Parallel_2D& para2d_local) +void elpa_diag( + MPI_Comm comm, + const int nband, + float *h_local, + float *s_local, + float *ekb, + float *wfc_2d, + Parallel_2D ¶2d_local) { std::cout << "Error: ELPA do not support single precision. " << std::endl; exit(1); @@ -72,21 +98,26 @@ void elpa_diag(MPI_Comm comm, #ifdef __MPI template -void diago_hs_para(T* h, - T* s, - const int lda, - const int nband, - typename GetTypeReal::type* const ekb, - T* const wfc, - const MPI_Comm& comm, - const int diag_subspace, - const int block_size) +void diago_hs_para( + T *h, + T *s, + const int lda, + const int nband, + typename GetTypeReal::type *const ekb, + T *const wfc, + const MPI_Comm &comm, + const int diag_subspace, + const int block_size) { int myrank = 0; MPI_Comm_rank(comm, &myrank); Parallel_2D para2d_global; Parallel_2D para2d_local; - para2d_global.init(lda, lda, lda, comm); + para2d_global.init( + lda, + lda, + lda, + comm); int max_nb = block_size; if (block_size == 0) @@ -115,19 +146,52 @@ void diago_hs_para(T* h, } } - para2d_local.init(lda, lda, max_nb, comm); + para2d_local.init( + lda, + lda, + max_nb, + comm); std::vector h_local(para2d_local.get_col_size() * para2d_local.get_row_size()); std::vector s_local(para2d_local.get_col_size() * para2d_local.get_row_size()); std::vector wfc_2d(para2d_local.get_col_size() * para2d_local.get_row_size()); // distribute h and s to 2D - Cpxgemr2d(lda, lda, h, 1, 1, para2d_global.desc, h_local.data(), 1, 1, para2d_local.desc, para2d_local.blacs_ctxt); - Cpxgemr2d(lda, lda, s, 1, 1, para2d_global.desc, s_local.data(), 1, 1, para2d_local.desc, para2d_local.blacs_ctxt); + Cpxgemr2d( + lda, + lda, + h, + 1, + 1, + para2d_global.desc, + h_local.data(), + 1, + 1, + para2d_local.desc, + para2d_local.blacs_ctxt); + Cpxgemr2d( + lda, + lda, + s, + 1, + 1, + para2d_global.desc, + s_local.data(), + 1, + 1, + para2d_local.desc, + para2d_local.blacs_ctxt); if (diag_subspace == 1) { #ifdef __ELPA - elpa_diag(comm, nband, h_local.data(), s_local.data(), ekb, wfc_2d.data(), para2d_local); + elpa_diag( + comm, + nband, + h_local.data(), + s_local.data(), + ekb, + wfc_2d.data(), + para2d_local); #else std::cout << "ERROR: try to use ELPA to solve the generalized eigenvalue problem, but ELPA is not compiled. " << std::endl; @@ -136,14 +200,15 @@ void diago_hs_para(T* h, } else if (diag_subspace == 2) { - hsolver::pxxxgvx_diag(para2d_local.desc, - para2d_local.get_row_size(), - para2d_local.get_col_size(), - nband, - h_local.data(), - s_local.data(), - ekb, - wfc_2d.data()); + hsolver::pxxxgvx_diag( + para2d_local.desc, + para2d_local.get_row_size(), + para2d_local.get_col_size(), + nband, + h_local.data(), + s_local.data(), + ekb, + wfc_2d.data()); } else { @@ -153,49 +218,64 @@ void diago_hs_para(T* h, } // gather wfc - Cpxgemr2d(lda, lda, wfc_2d.data(), 1, 1, para2d_local.desc, wfc, 1, 1, para2d_global.desc, para2d_local.blacs_ctxt); + Cpxgemr2d( + lda, + lda, + wfc_2d.data(), + 1, + 1, + para2d_local.desc, + wfc, + 1, + 1, + para2d_global.desc, + para2d_local.blacs_ctxt); } // template instantiation -template void diago_hs_para(double* h, - double* s, - const int lda, - const int nband, - typename GetTypeReal::type* const ekb, - double* const wfc, - const MPI_Comm& comm, - const int diag_subspace, - const int block_size); - -template void diago_hs_para>(std::complex* h, - std::complex* s, - const int lda, - const int nband, - typename GetTypeReal>::type* const ekb, - std::complex* const wfc, - const MPI_Comm& comm, - const int diag_subspace, - const int block_size); - -template void diago_hs_para(float* h, - float* s, - const int lda, - const int nband, - typename GetTypeReal::type* const ekb, - float* const wfc, - const MPI_Comm& comm, - const int diag_subspace, - const int block_size); - -template void diago_hs_para>(std::complex* h, - std::complex* s, - const int lda, - const int nband, - typename GetTypeReal>::type* const ekb, - std::complex* const wfc, - const MPI_Comm& comm, - const int diag_subspace, - const int block_size); +template void diago_hs_para( + double *h, + double *s, + const int lda, + const int nband, + typename GetTypeReal::type *const ekb, + double *const wfc, + const MPI_Comm &comm, + const int diag_subspace, + const int block_size); + +template void diago_hs_para>( + std::complex *h, + std::complex *s, + const int lda, + const int nband, + typename GetTypeReal>::type *const ekb, + std::complex *const wfc, + const MPI_Comm &comm, + const int diag_subspace, + const int block_size); + +template void diago_hs_para( + float *h, + float *s, + const int lda, + const int nband, + typename GetTypeReal::type *const ekb, + float *const wfc, + const MPI_Comm &comm, + const int diag_subspace, + const int block_size); + +template void diago_hs_para>( + std::complex *h, + std::complex *s, + const int lda, + const int nband, + typename GetTypeReal>::type *const ekb, + std::complex *const wfc, + const MPI_Comm &comm, + const int diag_subspace, + const int block_size); #endif diff --git a/source/source_hsolver/diag_hs_para.h b/source/source_hsolver/diag_hs_para.h index 7b65cefe688..741c300494a 100644 --- a/source/source_hsolver/diag_hs_para.h +++ b/source/source_hsolver/diag_hs_para.h @@ -1,5 +1,5 @@ -#include "source_base/parallel_2d.h" #include "source_base/macros.h" +#include "source_base/parallel_2d.h" #ifdef __MPI #include @@ -8,12 +8,11 @@ namespace hsolver { - #ifdef __MPI /** * @brief Parallel do the generalized eigenvalue problem - * + * * @tparam T double or std::complex or float or complex * @param H the hermitian matrix H. * @param S the overlap matrix S. @@ -23,24 +22,25 @@ namespace hsolver * @param wfc to store the eigenvectors * @param comm the communicator * @param diag_subspace the method to solve the generalized eigenvalue problem - * @param block_size the block size in 2d block cyclic distribution if use elpa or scalapack. - * + * @param block_size the block size in 2d block cyclic distribution if use elpa or scalapack. + * * @note 1. h and s should be full matrix in rank 0 of the communicator, and the other ranks is not concerned. * @note 2. wfc is complete in rank 0, and not store in other ranks. * @note 3. diag_subspace should be 1: by elpa, 2: by scalapack - * @note 4. block_size should be 0 or a positive integer. If it is 0, then will use a value as large as possible that is allowed + * @note 4. block_size should be 0 or a positive integer. If it is 0, then will use a value as large as possible that is + * allowed */ template -void diago_hs_para(T* h, - T* s, - const int lda, - const int nband, - typename GetTypeReal::type* const ekb, - T* const wfc, - const MPI_Comm& comm, - const int diag_subspace, - const int block_size = 0); +void diago_hs_para( + T *h, + T *s, + const int lda, + const int nband, + typename GetTypeReal::type *const ekb, + T *const wfc, + const MPI_Comm &comm, + const int diag_subspace, + const int block_size = 0); #endif -} // namespace hsolver - \ No newline at end of file +} // namespace hsolver diff --git a/source/source_hsolver/diago_bpcg.cpp b/source/source_hsolver/diago_bpcg.cpp index 363dd5429cb..e631a6c88ba 100644 --- a/source/source_hsolver/diago_bpcg.cpp +++ b/source/source_hsolver/diago_bpcg.cpp @@ -1,80 +1,115 @@ #include "source_hsolver/diago_bpcg.h" #include "diago_iter_assist.h" +#include "para_lin_tf.h" #include "source_base/global_function.h" #include "source_base/kernels/math_kernel_op.h" #include "source_base/parallel_comm.h" // different MPI worlds #include "source_hsolver/kernels/bpcg_kernel_op.h" -#include "para_lin_tf.h" #include #include #include #include -namespace hsolver { +namespace hsolver +{ -template -DiagoBPCG::DiagoBPCG(const Real* precondition_in) +template DiagoBPCG::DiagoBPCG(const Real *precondition_in) { - this->r_type = ct::DataTypeToEnum::value; - this->t_type = ct::DataTypeToEnum::value; - this->device_type = ct::DeviceTypeToEnum::value; + this->r_type = ct::DataTypeToEnum::value; + this->t_type = ct::DataTypeToEnum::value; + this->device_type = ct::DeviceTypeToEnum::value; - this->h_prec = std::move(ct::TensorMap((void *) precondition_in, r_type, device_type, {this->n_basis})); + this->h_prec = std::move( + ct::TensorMap( + (void *)precondition_in, + r_type, + device_type, {this->n_basis})); this->one = &one_; this->zero = &zero_; this->neg_one = &neg_one_; } -template -DiagoBPCG::~DiagoBPCG() { +template DiagoBPCG::~DiagoBPCG() +{ // Note, we do not need to free the h_prec and psi pointer as they are refs to the outside data } -template -void DiagoBPCG::init_iter(const int nband, const int nband_l, const int nbasis, const int ndim) { +template +void DiagoBPCG::init_iter( + const int nband, + const int nband_l, + const int nbasis, + const int ndim) +{ // Specify the problem size n_basis, n_band, while lda is n_basis - this->n_band = nband; - this->n_band_l = nband_l; - this->n_basis = nbasis; - this->n_dim = ndim; + this->n_band = nband; + this->n_band_l = nband_l; + this->n_basis = nbasis; + this->n_dim = ndim; // All column major tensors - this->beta = std::move(ct::Tensor(r_type, device_type, {this->n_band_l})); - this->eigen = std::move(ct::Tensor(r_type, device_type, {this->n_band})); - this->err_st = std::move(ct::Tensor(r_type, device_type, {this->n_band_l})); + this->beta = std::move(ct::Tensor(r_type, device_type, {this->n_band_l})); + this->eigen = std::move(ct::Tensor(r_type, device_type, {this->n_band})); + this->err_st = std::move(ct::Tensor(r_type, device_type, {this->n_band_l})); - this->hsub = std::move(ct::Tensor(t_type, device_type, {this->n_band, this->n_band})); + this->hsub = std::move(ct::Tensor(t_type, device_type, {this->n_band, this->n_band})); - this->hpsi = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); - this->spsi = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); - this->work = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); - this->hgrad = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); - this->sgrad = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); - this->grad_old = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); + this->hpsi = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); + this->spsi = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); + this->work = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); + this->hgrad = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); + this->sgrad = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); + this->grad_old = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); - this->prec = std::move(ct::Tensor(r_type, device_type, {this->n_basis})); + this->prec = std::move(ct::Tensor(r_type, device_type, {this->n_basis})); - this->grad = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); + this->grad = std::move(ct::Tensor(t_type, device_type, {this->n_band_l, this->n_basis})); #ifdef __MPI - this->pmmcn.set_dimension(BP_WORLD, POOL_WORLD, n_band_l, n_basis, n_band_l, n_basis, n_dim, n_band); - this->plintrans.set_dimension(n_dim, nband_l, n_band_l, n_basis, BP_WORLD, false); + this->pmmcn.set_dimension( + BP_WORLD, + POOL_WORLD, + n_band_l, + n_basis, + n_band_l, + n_basis, + n_dim, + n_band); + this->plintrans.set_dimension( + n_dim, + nband_l, + n_band_l, + n_basis, + BP_WORLD, + false); #else - this->pmmcn.set_dimension(n_band_l, n_basis, n_band_l, n_basis, n_dim, n_band); - this->plintrans.set_dimension(n_dim, nband_l, n_band_l, n_basis, false); + this->pmmcn.set_dimension( + n_band_l, + n_basis, + n_band_l, + n_basis, + n_dim, + n_band); + this->plintrans.set_dimension( + n_dim, + nband_l, + n_band_l, + n_basis, + false); #endif } -template -bool DiagoBPCG::test_error(const ct::Tensor& err_in, const std::vector& ethr_band) +template +bool DiagoBPCG::test_error(const ct::Tensor &err_in, const std::vector ðr_band) { - Real* _err_st = err_in.data(); + Real *_err_st = err_in.data(); bool not_conv = false; std::vector tmp_cpu; - if (err_in.device_type() == ct::DeviceType::GpuDevice) { + if (err_in.device_type() == ct::DeviceType::GpuDevice) + { // ct::Tensor h_err_in = err_in.to_device(); // _err_st = h_err_in.data(); // qianrui change it, because it can not pass the valgrind test @@ -82,186 +117,244 @@ bool DiagoBPCG::test_error(const ct::Tensor& err_in, const std::vecto _err_st = tmp_cpu.data(); syncmem_var_d2h_op()(_err_st, err_in.data(), this->n_band_l); } - for (int ii = 0; ii < this->n_band_l; ii++) { - if (_err_st[ii] > ethr_band[ii]) { + for (int ii = 0; ii < this->n_band_l; ii++) + { + if (_err_st[ii] > ethr_band[ii]) + { not_conv = true; } } #ifdef __MPI - MPI_Allreduce(MPI_IN_PLACE, ¬_conv, 1, MPI_C_BOOL, MPI_LOR, BP_WORLD); + MPI_Allreduce( + MPI_IN_PLACE, + ¬_conv, + 1, + MPI_C_BOOL, + MPI_LOR, + BP_WORLD); #endif return not_conv; } // Finally, the last one! -template +template void DiagoBPCG::line_minimize( - ct::Tensor& grad_in, - ct::Tensor& hgrad_in, - ct::Tensor& sgrad_in, - ct::Tensor& psi_out, - ct::Tensor& hpsi_out, - ct::Tensor& spsi_out) + ct::Tensor &grad_in, + ct::Tensor &hgrad_in, + ct::Tensor &sgrad_in, + ct::Tensor &psi_out, + ct::Tensor &hpsi_out, + ct::Tensor &spsi_out) { - line_minimize_with_block_op()(grad_in.data(), - hgrad_in.data(), - sgrad_in.data(), - psi_out.data(), - hpsi_out.data(), - spsi_out.data(), - this->n_dim, - this->n_basis, - this->n_band_l); + line_minimize_with_block_op()( + grad_in.data(), + hgrad_in.data(), + sgrad_in.data(), + psi_out.data(), + hpsi_out.data(), + spsi_out.data(), + this->n_dim, + this->n_basis, + this->n_band_l); } - // Finally, the last two! -template +template void DiagoBPCG::orth_cholesky( - ct::Tensor& workspace_in, - ct::Tensor& psi_out, - ct::Tensor& hpsi_out, - ct::Tensor& spsi_out, - ct::Tensor& hsub_out) + ct::Tensor &workspace_in, + ct::Tensor &psi_out, + ct::Tensor &hpsi_out, + ct::Tensor &spsi_out, + ct::Tensor &hsub_out) { // gemm: hsub_out(n_band x n_band) = psi_out^H(n_band x n_basis) * spsi_out(n_basis x n_band) - this->pmmcn.multiply(1.0, psi_out.data(), spsi_out.data(), 0.0, hsub_out.data()); + this->pmmcn.multiply( + 1.0, + psi_out.data(), + spsi_out.data(), + 0.0, + hsub_out.data()); // set hsub matrix to lower format; - ct::kernels::set_matrix()( - 'L', hsub_out.data(), this->n_band); + ct::kernels::set_matrix()('L', hsub_out.data(), this->n_band); ct::kernels::lapack_potrf()( - 'U', this->n_band, hsub_out.data(), this->n_band); + 'U', + this->n_band, + hsub_out.data(), + this->n_band); ct::kernels::lapack_trtri()( - 'U', 'N', this->n_band, hsub_out.data(), this->n_band); + 'U', + 'N', + this->n_band, + hsub_out.data(), + this->n_band); this->rotate_wf(hsub_out, psi_out, workspace_in); this->rotate_wf(hsub_out, hpsi_out, workspace_in); this->rotate_wf(hsub_out, spsi_out, workspace_in); } -template +template void DiagoBPCG::calc_grad_with_block( - const ct::Tensor& prec_in, - ct::Tensor& err_out, - ct::Tensor& beta_out, - ct::Tensor& psi_in, - ct::Tensor& hpsi_in, - ct::Tensor& spsi_in, - ct::Tensor& grad_out, - ct::Tensor& grad_old_out) + const ct::Tensor &prec_in, + ct::Tensor &err_out, + ct::Tensor &beta_out, + ct::Tensor &psi_in, + ct::Tensor &hpsi_in, + ct::Tensor &spsi_in, + ct::Tensor &grad_out, + ct::Tensor &grad_old_out) { - calc_grad_with_block_op()(prec_in.data(), - err_out.data(), - beta_out.data(), - psi_in.data(), - hpsi_in.data(), - spsi_in.data(), - grad_out.data(), - grad_old_out.data(), - this->n_dim, - this->n_basis, - this->n_band_l); + calc_grad_with_block_op()( + prec_in.data(), + err_out.data(), + beta_out.data(), + psi_in.data(), + hpsi_in.data(), + spsi_in.data(), + grad_out.data(), + grad_old_out.data(), + this->n_dim, + this->n_basis, + this->n_band_l); } -template -void DiagoBPCG::calc_prec() +template void DiagoBPCG::calc_prec() { syncmem_var_h2d_op()(this->prec.template data(), this->h_prec.template data(), this->n_basis); } -template +template void DiagoBPCG::orth_projection( - const ct::Tensor& psi_in, - const ct::Tensor& spsi_in, - ct::Tensor& hsub_in, - ct::Tensor& grad_out, - ct::Tensor& sgrad_out) + const ct::Tensor &psi_in, + const ct::Tensor &spsi_in, + ct::Tensor &hsub_in, + ct::Tensor &grad_out, + ct::Tensor &sgrad_out) { // gemm: hsub_in(n_band x n_band) = psi_in^H(n_band x n_basis) * sgrad_out(n_basis x n_band) - this->pmmcn.multiply(1.0, psi_in.data(), sgrad_out.data(), 0.0, hsub_in.data()); + this->pmmcn.multiply( + 1.0, + psi_in.data(), + sgrad_out.data(), + 0.0, + hsub_in.data()); // grad_out(n_basis x n_band) = 1.0 * grad_out(n_basis x n_band) - psi_in(n_basis x n_band) * hsub_in(n_band x // n_band) - this->plintrans.act(-1.0, psi_in.data(), hsub_in.data(), 1.0, grad_out.data()); - this->plintrans.act(-1.0, spsi_in.data(), hsub_in.data(), 1.0, sgrad_out.data()); + this->plintrans.act( + -1.0, + psi_in.data(), + hsub_in.data(), + 1.0, + grad_out.data()); + this->plintrans.act( + -1.0, + spsi_in.data(), + hsub_in.data(), + 1.0, + sgrad_out.data()); return; } -template -void DiagoBPCG::rotate_wf( - const ct::Tensor& hsub_in, - ct::Tensor& psi_out, - ct::Tensor& workspace_in) +template +void DiagoBPCG::rotate_wf(const ct::Tensor &hsub_in, ct::Tensor &psi_out, ct::Tensor &workspace_in) { // gemm: workspace_in(n_basis x n_band) = psi_out(n_basis x n_band) * hsub_in(n_band x n_band) - this->plintrans.act(1.0, psi_out.data(), hsub_in.data(), 0.0, workspace_in.data()); - syncmem_complex_2d_op()(psi_out.template data(), - this->n_basis, - workspace_in.template data(), - this->n_basis, - this->n_dim, - this->n_band_l); + this->plintrans.act( + 1.0, + psi_out.data(), + hsub_in.data(), + 0.0, + workspace_in.data()); + syncmem_complex_2d_op()( + psi_out.template data(), + this->n_basis, + workspace_in.template data(), + this->n_basis, + this->n_dim, + this->n_band_l); return; } -template -void DiagoBPCG::calc_hpsi_with_block( - const HSOperator& op, - const T *psi_in, - ct::Tensor& hpsi_out) +template +void DiagoBPCG::calc_hpsi_with_block(const HSOperator &op, const T *psi_in, ct::Tensor &hpsi_out) { // calculate all-band hpsi - op.hpsi(psi_in, hpsi_out.data(), this->n_basis, this->n_band_l); + op.hpsi( + psi_in, + hpsi_out.data(), + this->n_basis, + this->n_band_l); } -template -void DiagoBPCG::calc_spsi_with_block( - const HSOperator& op, - const T* psi_in, - ct::Tensor& spsi_out) +template +void DiagoBPCG::calc_spsi_with_block(const HSOperator &op, const T *psi_in, ct::Tensor &spsi_out) { - op.spsi(psi_in, spsi_out.data(), this->n_basis, this->n_band_l); + op.spsi( + psi_in, + spsi_out.data(), + this->n_basis, + this->n_band_l); } -template +template void DiagoBPCG::diag_hsub( - const ct::Tensor& psi_in, - const ct::Tensor& hpsi_in, - ct::Tensor& hsub_out, - ct::Tensor& eigenvalue_out) + const ct::Tensor &psi_in, + const ct::Tensor &hpsi_in, + ct::Tensor &hsub_out, + ct::Tensor &eigenvalue_out) { // gemm: hsub_out(n_band x n_band) = hpsi_in^T(n_band x n_basis) * psi_in(n_basis x n_band) - this->pmmcn.multiply(1.0, hpsi_in.data(), psi_in.data(), 0.0, hsub_out.data()); - - // ct::kernels::lapack_heevd()('V', 'U', hsub_out.data(), this->n_band, eigenvalue_out.data()); - ct::kernels::lapack_heevd()(this->n_band, hsub_out.data(), this->n_band, eigenvalue_out.data()); + this->pmmcn.multiply( + 1.0, + hpsi_in.data(), + psi_in.data(), + 0.0, + hsub_out.data()); + + // ct::kernels::lapack_heevd()('V', 'U', hsub_out.data(), this->n_band, + // eigenvalue_out.data()); + ct::kernels::lapack_heevd()( + this->n_band, + hsub_out.data(), + this->n_band, + eigenvalue_out.data()); return; } -template +template void DiagoBPCG::calc_hsub_with_block( - const HSOperator& op, - T *psi_in, - ct::Tensor& psi_out, - ct::Tensor& hpsi_out, - ct::Tensor& spsi_out, - ct::Tensor& hsub_out, - ct::Tensor& workspace_in, - ct::Tensor& eigenvalue_out) + const HSOperator &op, + T *psi_in, + ct::Tensor &psi_out, + ct::Tensor &hpsi_out, + ct::Tensor &spsi_out, + ct::Tensor &hsub_out, + ct::Tensor &workspace_in, + ct::Tensor &eigenvalue_out) { // Apply the H operator to psi and obtain the hpsi matrix. this->calc_hpsi_with_block(op, psi_in, hpsi_out); this->calc_spsi_with_block(op, psi_in, spsi_out); // Transform the generalized problem to an S-orthonormal subspace. - this->orth_cholesky(workspace_in, psi_out, hpsi_out, spsi_out, hsub_out); + this->orth_cholesky( + workspace_in, + psi_out, + hpsi_out, + spsi_out, + hsub_out); // Diagonalization of the subspace matrix. - this->diag_hsub(psi_out,hpsi_out, hsub_out, eigenvalue_out); + this->diag_hsub( + psi_out, + hpsi_out, + hsub_out, + eigenvalue_out); // inplace matmul to get the initial guessed wavefunction psi. // psi_out[n_basis, n_band] = psi_out[n_basis, n_band] x hsub_out[n_band, n_band] @@ -273,16 +366,20 @@ void DiagoBPCG::calc_hsub_with_block( return; } -template +template void DiagoBPCG::calc_hsub_with_block_exit( - ct::Tensor& psi_out, - ct::Tensor& hpsi_out, - ct::Tensor& hsub_out, - ct::Tensor& workspace_in, - ct::Tensor& eigenvalue_out) + ct::Tensor &psi_out, + ct::Tensor &hpsi_out, + ct::Tensor &hsub_out, + ct::Tensor &workspace_in, + ct::Tensor &eigenvalue_out) { // Diagonalization of the subspace matrix. - this->diag_hsub(psi_out, hpsi_out, hsub_out, eigenvalue_out); + this->diag_hsub( + psi_out, + hpsi_out, + hsub_out, + eigenvalue_out); // inplace matmul to get the initial guessed wavefunction psi. // psi_out[n_basis, n_band] = psi_out[n_basis, n_band] x hsub_out[n_band, n_band] @@ -292,36 +389,40 @@ void DiagoBPCG::calc_hsub_with_block_exit( } template -void DiagoBPCG::diag(const HSOperator& op, - T* psi_in, - Real* eigenvalue_in, - const std::vector& ethr_band) +void DiagoBPCG::diag( + const HSOperator &op, + T *psi_in, + Real *eigenvalue_in, + const std::vector ðr_band) { const int current_scf_iter = hsolver::DiagoIterAssist::SCF_ITER; // Get the pointer of the input psi - this->psi = std::move(ct::TensorMap(psi_in /*psi_in.get_pointer()*/, t_type, device_type, {this->n_band_l, this->n_basis})); + this->psi = std::move( + ct::TensorMap( + psi_in /*psi_in.get_pointer()*/, + t_type, + device_type, {this->n_band_l, this->n_basis})); // Update the precondition array this->calc_prec(); // Improving the initial guess of the wave function psi through a subspace diagonalization. - this->calc_hsub_with_block(op, - psi_in, - this->psi, - this->hpsi, - this->spsi, - this->hsub, - this->work, - this->eigen); + this->calc_hsub_with_block( + op, + psi_in, + this->psi, + this->hpsi, + this->spsi, + this->hsub, + this->work, + this->eigen); setmem_complex_op()(this->grad_old.template data(), 0, this->n_basis * this->n_band_l); setmem_var_op()(this->beta.template data(), std::numeric_limits::infinity(), this->n_band_l); int ntry = 0; - int max_iter = current_scf_iter > 1 ? - this->nline : - this->nline * 6; + int max_iter = current_scf_iter > 1 ? this->nline : this->nline * 6; do { ++ntry; @@ -331,8 +432,15 @@ void DiagoBPCG::diag(const HSOperator& op, // 3. calculate the gradient by hpsi - epsilo * psi // 4. gradient mix with the previous gradient // 5. Do precondition - this->calc_grad_with_block(this->prec, this->err_st, this->beta, - this->psi, this->hpsi, this->spsi, this->grad, this->grad_old); + this->calc_grad_with_block( + this->prec, + this->err_st, + this->beta, + this->psi, + this->hpsi, + this->spsi, + this->grad, + this->grad_old); // Apply S before projecting the search directions in the generalized metric. this->calc_spsi_with_block(op, this->grad.template data(), this->sgrad); @@ -340,7 +448,12 @@ void DiagoBPCG::diag(const HSOperator& op, // Orthogonalize column vectors g_i in matrix grad to column vectors p_j in matrix psi // for all 'j less or equal to i'. // Note: hsub and work are only used to store intermediate variables of gemm operator. - this->orth_projection(this->psi, this->spsi, this->hsub, this->grad, this->sgrad); + this->orth_projection( + this->psi, + this->spsi, + this->hsub, + this->grad, + this->sgrad); // this->grad_old = this->grad; syncmem_complex_op()(this->grad_old.template data(), this->grad.template data(), n_basis * n_band_l); @@ -352,24 +465,42 @@ void DiagoBPCG::diag(const HSOperator& op, // 1. normalize grad // 2. calculate theta // 3. update psi as well as hpsi - this->line_minimize(this->grad, this->hgrad, this->sgrad, this->psi, this->hpsi, this->spsi); + this->line_minimize( + this->grad, + this->hgrad, + this->sgrad, + this->psi, + this->hpsi, + this->spsi); // orthogonal psi by cholesky method - this->orth_cholesky(this->work, this->psi, this->hpsi, this->spsi, this->hsub); - - if (current_scf_iter == 1 && ntry % this->nline == 0) { - this->calc_hsub_with_block(op, - psi_in, - this->psi, - this->hpsi, - this->spsi, - this->hsub, - this->work, - this->eigen); + this->orth_cholesky( + this->work, + this->psi, + this->hpsi, + this->spsi, + this->hsub); + + if (current_scf_iter == 1 && ntry % this->nline == 0) + { + this->calc_hsub_with_block( + op, + psi_in, + this->psi, + this->hpsi, + this->spsi, + this->hsub, + this->work, + this->eigen); } } while (ntry < max_iter && this->test_error(this->err_st, ethr_band)); - this->calc_hsub_with_block_exit(this->psi, this->hpsi, this->hsub, this->work, this->eigen); + this->calc_hsub_with_block_exit( + this->psi, + this->hpsi, + this->hsub, + this->work, + this->eigen); int start_nband = 0; #ifdef __MPI diff --git a/source/source_hsolver/diago_bpcg.h b/source/source_hsolver/diago_bpcg.h index 66489e795c5..5625b2b5142 100644 --- a/source/source_hsolver/diago_bpcg.h +++ b/source/source_hsolver/diago_bpcg.h @@ -13,7 +13,8 @@ #include #include -namespace hsolver { +namespace hsolver +{ /** * @class DiagoBPCG @@ -21,22 +22,21 @@ namespace hsolver { * @tparam T The floating-point type used for calculations. * @tparam Device The device used for calculations (e.g., cpu or gpu). */ -template , typename Device = base_device::DEVICE_CPU> -class DiagoBPCG +template , typename Device = base_device::DEVICE_CPU> class DiagoBPCG { private: - // Note GetTypeReal::type will - // return T if T is real type(float, double), + // Note GetTypeReal::type will + // return T if T is real type(float, double), // otherwise return the real type of T(complex, std::complex) using Real = typename GetTypeReal::type; - // Column major psi in this class + // Column major psi in this class public: /** * @brief Constructor for DiagoBPCG class. * * @param precondition precondition data passed by the "Hamilt_PW" class. */ - explicit DiagoBPCG(const Real* precondition); + explicit DiagoBPCG(const Real *precondition); /** * @brief Destructor for DiagoBPCG class. @@ -54,7 +54,11 @@ class DiagoBPCG * @param nbasis The number of basis functions. Leading dimension of psi. * @param ndim The number of valid dimension of psi. */ - void init_iter(const int nband, const int nband_l, const int nbasis, const int ndim); + void init_iter( + const int nband, + const int nband_l, + const int nbasis, + const int ndim); /** * @brief Diagonalize the Hamiltonian using the BPCG method. @@ -65,10 +69,11 @@ class DiagoBPCG * @param psi_in Pointer to input wavefunction psi matrix with [dim: n_basis x n_band, column major]. * @param eigenvalue_in Pointer to the eigen array with [dim: n_band, column major]. */ - void diag(const HSOperator& op, - T* psi_in, - Real* eigenvalue_in, - const std::vector& ethr_band); + void diag( + const HSOperator &op, + T *psi_in, + Real *eigenvalue_in, + const std::vector ðr_band); private: /// the number of bands of all processes @@ -85,9 +90,8 @@ class DiagoBPCG ModuleBase::PGemmCN pmmcn; PLinearTransform plintrans; - - ct::DataType r_type = ct::DataType::DT_INVALID; - ct::DataType t_type = ct::DataType::DT_INVALID; + ct::DataType r_type = ct::DataType::DT_INVALID; + ct::DataType t_type = ct::DataType::DT_INVALID; ct::DeviceType device_type = ct::DeviceType::UnKnown; ct::Tensor prec = {}, h_prec = {}; @@ -105,7 +109,7 @@ class DiagoBPCG ct::Tensor psi = {}; ct::Tensor hpsi = {}; ct::Tensor spsi = {}; - + ct::Tensor hsub = {}; /// H|psi> - epsilo * S|psi>, grad of the generalized problem. @@ -120,8 +124,8 @@ class DiagoBPCG // These are for hsolver gemm_op use /// ctx is nothing but the devices used in gemm_op (Device * ctx = nullptr;), - Device * ctx = {}; - // Pointer to objects of 1 and 0 for gemm + Device *ctx = {}; + // Pointer to objects of 1 and 0 for gemm const T *one = nullptr, *zero = nullptr, *neg_one = nullptr; const T one_ = static_cast(1.0), zero_ = static_cast(0.0), neg_one_ = static_cast(-1.0); @@ -159,10 +163,7 @@ class DiagoBPCG * @param psi_in The input wavefunction psi. * @param hpsi_out Pointer to the array where the resulting hpsi matrix will be stored. */ - void calc_hpsi_with_block( - const HSOperator& op, - const T *psi_in, - ct::Tensor& hpsi_out); + void calc_hpsi_with_block(const HSOperator &op, const T *psi_in, ct::Tensor &hpsi_out); /** * @brief Apply the overlap operator to a wavefunction block. @@ -171,7 +172,7 @@ class DiagoBPCG * @param psi_in The input wavefunction block. * @param spsi_out The resulting S|psi> block. */ - void calc_spsi_with_block(const HSOperator& op, const T* psi_in, ct::Tensor& spsi_out); + void calc_spsi_with_block(const HSOperator &op, const T *psi_in, ct::Tensor &spsi_out); /** * @brief Diagonalization of the subspace matrix. @@ -188,10 +189,10 @@ class DiagoBPCG * @param eigenvalue_out Computed eigen array with [dim: n_band] */ void diag_hsub( - const ct::Tensor& psi_in, - const ct::Tensor& hpsi_in, - ct::Tensor& hsub_out, - ct::Tensor& eigenvalue_out); + const ct::Tensor &psi_in, + const ct::Tensor &hpsi_in, + ct::Tensor &hsub_out, + ct::Tensor &eigenvalue_out); /** * @brief Inplace matrix multiplication to obtain the initial guessed wavefunction. @@ -203,10 +204,7 @@ class DiagoBPCG * @param hsub_in Subspace matrix input, dim [n_basis, n_band] with column major. * @param psi_out output wavefunction matrix with dim [n_basis, n_band], column major. */ - void rotate_wf( - const ct::Tensor& hsub_in, - ct::Tensor& psi_out, - ct::Tensor& workspace_in); + void rotate_wf(const ct::Tensor &hsub_in, ct::Tensor &psi_out, ct::Tensor &workspace_in); /** * @brief Calculate the gradient for all bands used in CG method. @@ -234,11 +232,14 @@ class DiagoBPCG * 5. Do precondition */ void calc_grad_with_block( - const ct::Tensor& prec_in, - ct::Tensor& err_out, - ct::Tensor& beta_out, - ct::Tensor& psi_in, ct::Tensor& hpsi_in, ct::Tensor& spsi_in, - ct::Tensor& grad_out, ct::Tensor& grad_old_out); + const ct::Tensor &prec_in, + ct::Tensor &err_out, + ct::Tensor &beta_out, + ct::Tensor &psi_in, + ct::Tensor &hpsi_in, + ct::Tensor &spsi_in, + ct::Tensor &grad_out, + ct::Tensor &grad_old_out); /** * @@ -257,12 +258,15 @@ class DiagoBPCG * @param eigenvalue_out Computed eigen. */ void calc_hsub_with_block( - const HSOperator& op, + const HSOperator &op, T *psi_in, - ct::Tensor& psi_out, ct::Tensor& hpsi_out, ct::Tensor& spsi_out, - ct::Tensor& hsub_out, ct::Tensor& workspace_in, - ct::Tensor& eigenvalue_out); - + ct::Tensor &psi_out, + ct::Tensor &hpsi_out, + ct::Tensor &spsi_out, + ct::Tensor &hsub_out, + ct::Tensor &workspace_in, + ct::Tensor &eigenvalue_out); + /** * * @brief Apply the Hamiltonian operator to psi and obtain the hpsi matrix. @@ -280,11 +284,11 @@ class DiagoBPCG * @param eigenvalue_out Computed eigen. */ void calc_hsub_with_block_exit( - ct::Tensor& psi_out, - ct::Tensor& hpsi_out, - ct::Tensor& hsub_out, - ct::Tensor& workspace_in, - ct::Tensor& eigenvalue_out); + ct::Tensor &psi_out, + ct::Tensor &hpsi_out, + ct::Tensor &hsub_out, + ct::Tensor &workspace_in, + ct::Tensor &eigenvalue_out); /** * @brief Orthogonalize column vectors in grad to column vectors in psi. @@ -297,11 +301,11 @@ class DiagoBPCG * @note This function is a member of the DiagoBPCG class. */ void orth_projection( - const ct::Tensor& psi_in, - const ct::Tensor& spsi_in, - ct::Tensor& hsub_in, - ct::Tensor& grad_out, - ct::Tensor& sgrad_out); + const ct::Tensor &psi_in, + const ct::Tensor &spsi_in, + ct::Tensor &hsub_in, + ct::Tensor &grad_out, + ct::Tensor &sgrad_out); /** * @@ -317,12 +321,12 @@ class DiagoBPCG * 3. Update psi as well as hpsi. */ void line_minimize( - ct::Tensor& grad_in, - ct::Tensor& hgrad_in, - ct::Tensor& sgrad_in, - ct::Tensor& psi_out, - ct::Tensor& hpsi_out, - ct::Tensor& spsi_out); + ct::Tensor &grad_in, + ct::Tensor &hgrad_in, + ct::Tensor &sgrad_in, + ct::Tensor &psi_out, + ct::Tensor &hpsi_out, + ct::Tensor &spsi_out); /** * @brief S-orthogonalize and normalize the column vectors in psi_out using Cholesky decomposition. @@ -333,11 +337,11 @@ class DiagoBPCG * @param hsub_out Input Hamiltonian product array. [dim: n_band x n_band, column major, lda = n_band]. */ void orth_cholesky( - ct::Tensor& workspace_in, - ct::Tensor& psi_out, - ct::Tensor& hpsi_out, - ct::Tensor& spsi_out, - ct::Tensor& hsub_out); + ct::Tensor &workspace_in, + ct::Tensor &psi_out, + ct::Tensor &hpsi_out, + ct::Tensor &spsi_out, + ct::Tensor &hsub_out); /** * @brief Checks if the error satisfies the given threshold. @@ -346,7 +350,7 @@ class DiagoBPCG * @param thr_in The threshold. * @return Returns true if all error values are less than or equal to the threshold, false otherwise. */ - bool test_error(const ct::Tensor& err_in, const std::vector& ethr_band); + bool test_error(const ct::Tensor &err_in, const std::vector ðr_band); using ct_Device = typename ct::PsiToContainer::type; using setmem_var_op = ct::kernels::set_memory; @@ -365,7 +369,6 @@ class DiagoBPCG // defined in source_base/module_device/types.h // different from ct_Device! using gemm_op = ModuleBase::gemm_op; - }; } // namespace hsolver diff --git a/source/source_hsolver/diago_cg.cpp b/source/source_hsolver/diago_cg.cpp index 95071ca3a78..60f52cba129 100644 --- a/source/source_hsolver/diago_cg.cpp +++ b/source/source_hsolver/diago_cg.cpp @@ -1,4 +1,5 @@ -#include +#include "source_hsolver/diago_cg.h" + #include "ATen/core/tensor_map.h" #include "ATen/core/tensor_utils.h" #include "ATen/kernels/lapack.h" @@ -6,18 +7,19 @@ #include "ATen/ops/einsum_op.h" #include "ATen/ops/linalg_op.h" #include "source_base/constants.h" +#include "source_base/global_function.h" // ModuleBase::GlobalFunc::NOTE #include "source_base/memory_recorder.h" #include "source_base/parallel_reduce.h" #include "source_base/timer.h" -#include "source_base/tool_title.h" // ModuleBase::TITLE -#include "source_base/global_function.h" // ModuleBase::GlobalFunc::NOTE -#include "source_hsolver/diago_cg.h" +#include "source_base/tool_title.h" // ModuleBase::TITLE #include "source_hsolver/diago_iter_assist.h" +#include + using namespace hsolver; template -DiagoCG::DiagoCG(const std::string& basis_type, const std::string& calculation) +DiagoCG::DiagoCG(const std::string &basis_type, const std::string &calculation) #ifdef __MPI : diag_comm_(MPI_COMM_SELF, 0, 1) #else @@ -32,12 +34,13 @@ DiagoCG::DiagoCG(const std::string& basis_type, const std::string& ca } template -DiagoCG::DiagoCG(const std::string& basis_type, - const std::string& calculation, - const bool& need_subspace, - const diag_comm_info& diag_comm, - const Real& pw_diag_thr, - const int& pw_diag_nmax) +DiagoCG::DiagoCG( + const std::string &basis_type, + const std::string &calculation, + const bool &need_subspace, + const diag_comm_info &diag_comm, + const Real &pw_diag_thr, + const int &pw_diag_nmax) : diag_comm_(diag_comm) { basis_type_ = basis_type; @@ -50,8 +53,7 @@ DiagoCG::DiagoCG(const std::string& basis_type, this->neg_one_ = new T(static_cast(-1.0)); } -template -DiagoCG::~DiagoCG() +template DiagoCG::~DiagoCG() { delete this->one_; delete this->zero_; @@ -59,10 +61,11 @@ DiagoCG::~DiagoCG() } template -void DiagoCG::diag_once(const ct::Tensor& prec_in, - ct::Tensor& psi, - ct::Tensor& eigen, - const std::vector& ethr_band) +void DiagoCG::diag_once( + const ct::Tensor &prec_in, + ct::Tensor &psi, + ct::Tensor &eigen, + const std::vector ðr_band) { ModuleBase::TITLE("DiagoCG", "diag_once"); ModuleBase::timer::start("DiagoCG", "diag_once"); @@ -83,34 +86,34 @@ void DiagoCG::diag_once(const ct::Tensor& prec_in, // Works for generalized eigenvalue problem (US pseudopotentials) as well //------------------------------------------------------------------- // phi_m = new psi::Psi(phi, 1, 1); - auto phi_m - = std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); + auto phi_m = + std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); // hphi.resize(this->n_basis_max_, ModuleBase::ZERO); - auto hphi - = std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); + auto hphi = + std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); // sphi.resize(this->n_basis_max_, ModuleBase::ZERO); - auto sphi - = std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); + auto sphi = + std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); // pphi.resize(this->n_basis_max_, ModuleBase::ZERO); - auto pphi - = std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); + auto pphi = + std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); // cg = new psi::Psi(phi, 1, 1); - auto cg - = std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); + auto cg = + std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); // scg.resize(this->n_basis_max_, ModuleBase::ZERO); - auto scg - = std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); + auto scg = + std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); // grad.resize(this->n_basis_max_, ModuleBase::ZERO); - auto grad - = std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); + auto grad = + std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); // g0.resize(this->n_basis_max_, ModuleBase::ZERO); - auto g0 - = std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); + auto g0 = + std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_basis_})); // lagrange.resize(this->n_band, ModuleBase::ZERO); - auto lagrange - = std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_band_})); + auto lagrange = + std::move(ct::Tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {this->n_band_})); auto prec = prec_in; if (prec.NumElements() == 0) @@ -127,10 +130,26 @@ void DiagoCG::diag_once(const ct::Tensor& prec_in, { phi_m.sync(psi[m]); // copy psi_in into internal psi, m=0 has been done in Constructor - this->op_->spsi(phi_m.data(), sphi.data(), this->n_basis_, 1); // sphi = S|psi(m)> - this->schmit_orth(m, psi, sphi, phi_m); - this->op_->spsi(phi_m.data(), sphi.data(), this->n_basis_, 1); // sphi = S|psi(m)> - this->op_->hpsi(phi_m.data(), hphi.data(), this->n_basis_, 1); // hphi = H|psi(m)> + this->op_->spsi( + phi_m.data(), + sphi.data(), + this->n_basis_, + 1); // sphi = S|psi(m)> + this->schmit_orth( + m, + psi, + sphi, + phi_m); + this->op_->spsi( + phi_m.data(), + sphi.data(), + this->n_basis_, + 1); // sphi = S|psi(m)> + this->op_->hpsi( + phi_m.data(), + hphi.data(), + this->n_basis_, + 1); // hphi = H|psi(m)> eigen_pack[m] = dot_real_op()(this->n_basis_, phi_m.data(), hphi.data()); @@ -142,32 +161,52 @@ void DiagoCG::diag_once(const ct::Tensor& prec_in, do { - this->calc_grad(prec, grad, hphi, sphi, pphi); - this->orth_grad(psi, m, grad, scg, lagrange); - this->calc_gamma_cg(iter, // const int& - cg_norm, - theta, // const Real& - prec, - scg, - grad, - phi_m, // const Tensor& - gg_last, // Real& - g0, - cg); // Tensor& - - this->op_->hpsi(cg.data(), pphi.data(), this->n_basis_, 1); - this->op_->spsi(cg.data(), scg.data(), this->n_basis_, 1); - - converged = this->update_psi(pphi, - cg, - scg, // const Tensor& - ethr_band[m], - cg_norm, - theta, - eigen_pack[m], // Real& - phi_m, - sphi, - hphi); // Tensor& + this->calc_grad( + prec, + grad, + hphi, + sphi, + pphi); + this->orth_grad( + psi, + m, + grad, + scg, + lagrange); + this->calc_gamma_cg( + iter, // const int& + cg_norm, + theta, // const Real& + prec, + scg, + grad, + phi_m, // const Tensor& + gg_last, // Real& + g0, + cg); // Tensor& + + this->op_->hpsi( + cg.data(), + pphi.data(), + this->n_basis_, + 1); + this->op_->spsi( + cg.data(), + scg.data(), + this->n_basis_, + 1); + + converged = this->update_psi( + pphi, + cg, + scg, // const Tensor& + ethr_band[m], + cg_norm, + theta, + eigen_pack[m], // Real& + phi_m, + sphi, + hphi); // Tensor& } while (!converged && ++iter < pw_diag_nmax_); @@ -190,7 +229,8 @@ void DiagoCG::diag_once(const ct::Tensor& prec_in, int ii = 0; for (ii = m - 2; ii >= 0; ii--) { - if (eigen_pack[m] - eigen_pack[ii] > 2.0 * pw_diag_thr_){ + if (eigen_pack[m] - eigen_pack[ii] > 2.0 * pw_diag_thr_) + { break; } } @@ -208,8 +248,8 @@ void DiagoCG::diag_once(const ct::Tensor& prec_in, eigen_pack[ii] = e0; psi[ii].sync(pphi); } // endif - } // end reorder - } // end m + } // end reorder + } // end m avg /= this->n_band_; avg_iter_ += avg; @@ -218,11 +258,12 @@ void DiagoCG::diag_once(const ct::Tensor& prec_in, } // end subroutine ccgdiagg template -void DiagoCG::calc_grad(const ct::Tensor& prec, - ct::Tensor& grad, - ct::Tensor& hphi, - ct::Tensor& sphi, - ct::Tensor& pphi) +void DiagoCG::calc_grad( + const ct::Tensor &prec, + ct::Tensor &grad, + ct::Tensor &hphi, + ct::Tensor &sphi, + ct::Tensor &pphi) { // for (int i = 0; i < this->n_basis_; i++) // { @@ -233,8 +274,16 @@ void DiagoCG::calc_grad(const ct::Tensor& prec, // } // denghui replace this at 20221106 // TODO: use GPU precondition to initialize CG class - ModuleBase::vector_div_vector_op()(this->n_basis_, grad.data(), hphi.data(), prec.data()); - ModuleBase::vector_div_vector_op()(this->n_basis_, pphi.data(), sphi.data(), prec.data()); + ModuleBase::vector_div_vector_op()( + this->n_basis_, + grad.data(), + hphi.data(), + prec.data()); + ModuleBase::vector_div_vector_op()( + this->n_basis_, + pphi.data(), + sphi.data(), + prec.data()); // Update lambda ! // (4) @@ -254,80 +303,89 @@ void DiagoCG::calc_grad(const ct::Tensor& prec, // grad.data()[i] -= lambda * this->pphi[i]; // } // haozhihan replace this 2022-10-6 - ModuleBase::vector_add_vector_op()(this->n_basis_, - grad.data(), - grad.data(), - 1.0, - pphi.data(), - (-lambda)); + ModuleBase::vector_add_vector_op()( + this->n_basis_, + grad.data(), + grad.data(), + 1.0, + pphi.data(), (-lambda)); } template -void DiagoCG::orth_grad(const ct::Tensor& psi, - const int& m, - ct::Tensor& grad, - ct::Tensor& scg, - ct::Tensor& lagrange) +void DiagoCG::orth_grad( + const ct::Tensor &psi, + const int &m, + ct::Tensor &grad, + ct::Tensor &scg, + ct::Tensor &lagrange) { const int matrix_ld = std::max(this->n_basis_, 1); if (this->n_basis_ == 0) { lagrange.zero(); } - this->op_->spsi(grad.data(), scg.data(), this->n_basis_, 1); // scg = S|grad> - ModuleBase::gemv_op()('C', - this->n_basis_, - m, - this->one_, - psi.data(), - matrix_ld, - scg.data(), - 1, - this->zero_, - lagrange.data(), - 1); + this->op_->spsi( + grad.data(), + scg.data(), + this->n_basis_, + 1); // scg = S|grad> + ModuleBase::gemv_op()( + 'C', + this->n_basis_, + m, + this->one_, + psi.data(), + matrix_ld, + scg.data(), + 1, + this->zero_, + lagrange.data(), + 1); Parallel_Reduce::reduce_pool(lagrange.data(), m); // (3) orthogonal |g> and |scg> to all states (0~m-1) //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // haozhihan replace 2022-10-07 - ModuleBase::gemv_op()('N', - this->n_basis_, - m, - this->neg_one_, - psi.data(), - matrix_ld, - lagrange.data(), - 1, - this->one_, - grad.data(), - 1); - - ModuleBase::gemv_op()('N', - this->n_basis_, - m, - this->neg_one_, - psi.data(), - matrix_ld, - lagrange.data(), - 1, - this->one_, - scg.data(), - 1); + ModuleBase::gemv_op()( + 'N', + this->n_basis_, + m, + this->neg_one_, + psi.data(), + matrix_ld, + lagrange.data(), + 1, + this->one_, + grad.data(), + 1); + + ModuleBase::gemv_op()( + 'N', + this->n_basis_, + m, + this->neg_one_, + psi.data(), + matrix_ld, + lagrange.data(), + 1, + this->one_, + scg.data(), + 1); } template -void DiagoCG::calc_gamma_cg(const int& iter, - const Real& cg_norm, - const Real& theta, - const ct::Tensor& prec, - const ct::Tensor& scg, - const ct::Tensor& grad, - const ct::Tensor& phi_m, - Real& gg_last, - ct::Tensor& g0, - ct::Tensor& cg) +void DiagoCG::calc_gamma_cg( + const int &iter, + const Real &cg_norm, + const Real &theta, + const ct::Tensor &prec, + const ct::Tensor &scg, + const ct::Tensor &grad, + const ct::Tensor &phi_m, + Real &gg_last, + ct::Tensor &g0, + ct::Tensor &cg) { Real gg_inter; if (iter > 0) @@ -335,8 +393,7 @@ void DiagoCG::calc_gamma_cg(const int& iter, // (1) Update gg_inter! // gg_inter = // Attention : the 'g' in g0 is getted last time - gg_inter - = ModuleBase::dot_real_op()(this->n_basis_, grad.data(), g0.data()); // b means before + gg_inter = ModuleBase::dot_real_op()(this->n_basis_, grad.data(), g0.data()); // b means before } // (2) Update for g0! @@ -350,7 +407,12 @@ void DiagoCG::calc_gamma_cg(const int& iter, // } // denghui replace this 20221106 // TODO: use GPU precondition instead - ModuleBase::vector_mul_vector_op()(this->n_basis_, g0.data(), scg.data(), prec.data(), false); + ModuleBase::vector_mul_vector_op()( + this->n_basis_, + g0.data(), + scg.data(), + prec.data(), + false); // (3) Update gg_now! // gg_now = < g|P|scg > = < g|g0 > @@ -378,12 +440,13 @@ void DiagoCG::calc_gamma_cg(const int& iter, // pcg[i] = gamma * pcg[i] + grad.data()[i]; // } // haozhihan replace this 2022-10-6 - ModuleBase::vector_add_vector_op()(this->n_basis_, - cg.data(), - cg.data(), - gamma, - grad.data(), - 1.0); + ModuleBase::vector_add_vector_op()( + this->n_basis_, + cg.data(), + cg.data(), + gamma, + grad.data(), + 1.0); const Real norma = gamma * cg_norm * sin(theta); T znorma = static_cast(norma * -1); @@ -395,32 +458,40 @@ void DiagoCG::calc_gamma_cg(const int& iter, { pcg[i] -= norma * pphi_m[i]; }*/ - ModuleBase::axpy_op()(this->n_basis_, &znorma, phi_m.data(), 1, cg.data(), 1); + ModuleBase::axpy_op()( + this->n_basis_, + &znorma, + phi_m.data(), + 1, + cg.data(), + 1); } } template -bool DiagoCG::update_psi(const ct::Tensor& pphi, - const ct::Tensor& cg, - const ct::Tensor& scg, - const double& ethreshold, - Real& cg_norm, - Real& theta, - Real& eigen, - ct::Tensor& phi_m, - ct::Tensor& sphi, - ct::Tensor& hphi) +bool DiagoCG::update_psi( + const ct::Tensor &pphi, + const ct::Tensor &cg, + const ct::Tensor &scg, + const double ðreshold, + Real &cg_norm, + Real &theta, + Real &eigen, + ct::Tensor &phi_m, + ct::Tensor &sphi, + ct::Tensor &hphi) { cg_norm = sqrt(ModuleBase::dot_real_op()(this->n_basis_, cg.data(), scg.data())); - if (cg_norm < 1.0e-10){ + if (cg_norm < 1.0e-10) + { return true; } - const Real a0 - = ModuleBase::dot_real_op()(this->n_basis_, phi_m.data(), pphi.data()) * 2.0 / cg_norm; - const Real b0 - = ModuleBase::dot_real_op()(this->n_basis_, cg.data(), pphi.data()) / (cg_norm * cg_norm); + const Real a0 = + ModuleBase::dot_real_op()(this->n_basis_, phi_m.data(), pphi.data()) * 2.0 / cg_norm; + const Real b0 = + ModuleBase::dot_real_op()(this->n_basis_, cg.data(), pphi.data()) / (cg_norm * cg_norm); const Real e0 = eigen; theta = atan(a0 / (e0 - b0)) / 2.0; @@ -446,12 +517,13 @@ bool DiagoCG::update_psi(const ct::Tensor& pphi, // } // haozhihan replace this 2022-10-6 - ModuleBase::vector_add_vector_op()(this->n_basis_, - phi_m.data(), - phi_m.data(), - cost, - cg.data(), - sint_norm); + ModuleBase::vector_add_vector_op()( + this->n_basis_, + phi_m.data(), + phi_m.data(), + cost, + cg.data(), + sint_norm); if (std::abs(eigen - e0) < ethreshold) { @@ -467,24 +539,30 @@ bool DiagoCG::update_psi(const ct::Tensor& pphi, // } // haozhihan replace this 2022-10-6 - ModuleBase::vector_add_vector_op()(this->n_basis_, - sphi.data(), - sphi.data(), - cost, - scg.data(), - sint_norm); - ModuleBase::vector_add_vector_op()(this->n_basis_, - hphi.data(), - hphi.data(), - cost, - pphi.data(), - sint_norm); + ModuleBase::vector_add_vector_op()( + this->n_basis_, + sphi.data(), + sphi.data(), + cost, + scg.data(), + sint_norm); + ModuleBase::vector_add_vector_op()( + this->n_basis_, + hphi.data(), + hphi.data(), + cost, + pphi.data(), + sint_norm); return false; } } template -void DiagoCG::schmit_orth(const int& m, const ct::Tensor& psi, const ct::Tensor& sphi, ct::Tensor& phi_m) +void DiagoCG::schmit_orth( + const int &m, + const ct::Tensor &psi, + const ct::Tensor &sphi, + ct::Tensor &phi_m) { const int matrix_ld = std::max(this->n_basis_, 1); // ModuleBase::TITLE("DiagoCG","schmit_orth"); @@ -506,34 +584,36 @@ void DiagoCG::schmit_orth(const int& m, const ct::Tensor& psi, const //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // haozhihan replace 2022-10-6 int inc = 1; - ModuleBase::gemv_op()('C', - this->n_basis_, - m + 1, - this->one_, - psi.data(), - matrix_ld, - sphi.data(), - inc, - this->zero_, - lagrange_so.data(), - inc); + ModuleBase::gemv_op()( + 'C', + this->n_basis_, + m + 1, + this->one_, + psi.data(), + matrix_ld, + sphi.data(), + inc, + this->zero_, + lagrange_so.data(), + inc); // be careful , here reduce m+1 Parallel_Reduce::reduce_pool(lagrange_so.data(), m + 1); //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // haozhihan replace 2022-10-6 - ModuleBase::gemv_op()('N', - this->n_basis_, - m, - this->neg_one_, - psi.data(), - matrix_ld, - lagrange_so.data(), - inc, - this->one_, - phi_m.data(), - inc); + ModuleBase::gemv_op()( + 'N', + this->n_basis_, + m, + this->neg_one_, + psi.data(), + matrix_ld, + lagrange_so.data(), + inc, + this->one_, + phi_m.data(), + inc); //====================================================================== /*for (int j = 0; j < m; j++) @@ -545,8 +625,11 @@ void DiagoCG::schmit_orth(const int& m, const ct::Tensor& psi, const psi_norm -= ( conj(lagrange[j]) * lagrange[j] ).real(); }*/ //>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> - auto psi_norm = ct::extract(lagrange_so[m]) - - dot_real_op()(m, lagrange_so.data(), lagrange_so.data(), false); + auto psi_norm = ct::extract(lagrange_so[m]) - dot_real_op()( + m, + lagrange_so.data(), + lagrange_so.data(), + false); if (psi_norm <= 0.0) { @@ -574,13 +657,22 @@ void DiagoCG::schmit_orth(const int& m, const ct::Tensor& psi, const // { // pphi_m[ig] /= psi_norm; // } - ModuleBase::vector_mul_real_op()(this->n_basis_, phi_m.data(), phi_m.data(), Real(1.0 / psi_norm)); + ModuleBase::vector_mul_real_op()( + this->n_basis_, + phi_m.data(), + phi_m.data(), + Real(1.0 / psi_norm)); // ModuleBase::timer::end("DiagoCG","schmit_orth"); } template -void DiagoCG::diag_subspace(const T* psi_in, T* psi_out, const int dim, const int nband, const bool S_orth) +void DiagoCG::diag_subspace( + const T *psi_in, + T *psi_out, + const int dim, + const int nband, + const bool S_orth) { // subspace diagonalization of the current nband vectors, packed with leading dimension dim; // the eigenvalues it produces are not needed, CG recomputes them. @@ -589,11 +681,21 @@ void DiagoCG::diag_subspace(const T* psi_in, T* psi_out, const int di // S-orthonormal after a CG restart, the results; wavefunction-sensitive outputs such as // the Wannier90 projections rely on the generalized path. std::vector eigen(nband, 0.0); - DiagoIterAssist::diag_subspace(*op_, psi_in, psi_out, nband, nband, dim, dim, eigen.data(), diag_comm_, false); + DiagoIterAssist::diag_subspace( + *op_, + psi_in, + psi_out, + nband, + nband, + dim, + dim, + eigen.data(), + diag_comm_, + false); } template -bool DiagoCG::test_exit_cond(const int& ntry, const int& notconv) const +bool DiagoCG::test_exit_cond(const int &ntry, const int ¬conv) const { const bool scf = calculation_ != "nscf"; // If ntry <=5, try to do it better, if ntry > 5, exit. @@ -607,37 +709,40 @@ bool DiagoCG::test_exit_cond(const int& ntry, const int& notconv) con } template -double DiagoCG::diag(const HSOperator& op, - const int ld_psi, - const int nband, - const int dim, - T* psi_in, - Real* eigenvalue_in, - const std::vector& ethr_band, - const Real* prec) +double DiagoCG::diag( + const HSOperator &op, + const int ld_psi, + const int nband, + const int dim, + T *psi_in, + Real *eigenvalue_in, + const std::vector ðr_band, + const Real *prec) { REQUIRES_OK(ld_psi >= dim, "DiagoCG::diag: ld_psi must be >= dim"); - REQUIRES_OK(static_cast(ethr_band.size()) >= nband, - "DiagoCG::diag: ethr_band size must be >= nband"); - - auto psi = ct::TensorMap(psi_in, - ct::DataTypeToEnum::value, - ct::DeviceTypeToEnum::value, - ct::TensorShape({nband, ld_psi})); - auto eigen = ct::TensorMap(eigenvalue_in, - ct::DataTypeToEnum::value, - ct::DeviceTypeToEnum::value, - ct::TensorShape({nband})); + REQUIRES_OK(static_cast(ethr_band.size()) >= nband, "DiagoCG::diag: ethr_band size must be >= nband"); + + auto psi = ct::TensorMap( + psi_in, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nband, ld_psi})); + auto eigen = ct::TensorMap( + eigenvalue_in, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nband})); // An empty LR rank can pass nullptr. Do not mistake Tensor's default // one-element float buffer for a supplied preconditioner. ct::Tensor prec_tensor(ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {0}); if (prec != nullptr) { - prec_tensor = ct::TensorMap(const_cast(prec), - ct::DataTypeToEnum::value, - ct::DeviceTypeToEnum::value, - ct::TensorShape({dim})) + prec_tensor = ct::TensorMap( + const_cast(prec), + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({dim})) .template to_device(); } @@ -667,21 +772,34 @@ double DiagoCG::diag(const HSOperator& op, { ct::TensorMap psi_map = ct::TensorMap(psi.data(), psi_temp); const bool assume_S_orthogonal = true; - this->diag_subspace(psi_temp.data(), psi_map.data(), dim, nband, assume_S_orthogonal); + this->diag_subspace( + psi_temp.data(), + psi_map.data(), + dim, + nband, + assume_S_orthogonal); psi_temp.sync(psi_map); } else if (need_subspace_) { ct::TensorMap psi_map = ct::TensorMap(psi.data(), psi_temp); const bool assume_S_orthogonal = false; - this->diag_subspace(psi_temp.data(), psi_map.data(), dim, nband, assume_S_orthogonal); + this->diag_subspace( + psi_temp.data(), + psi_map.data(), + dim, + nband, + assume_S_orthogonal); psi_temp.sync(psi_map); } - ++ntry; avg_iter_ += 1.0; - this->diag_once(prec_tensor, psi_temp, eigen, ethr_band); + this->diag_once( + prec_tensor, + psi_temp, + eigen, + ethr_band); } while (this->test_exit_cond(ntry, this->notconv_)); if (this->notconv_ > std::max(5, this->n_band_ / 4)) @@ -695,8 +813,8 @@ double DiagoCG::diag(const HSOperator& op, psi.sync(psi_temp); #ifdef __DEBUG -// only output iter count for each band if DEBUG! -// this should not be output in production log + // only output iter count for each band if DEBUG! + // this should not be output in production log std::cout << "\n DiagoCG::diag' avg_iter_ = " << avg_iter_; std::cout << "\n DiagoCG::diag' iter_band = "; for (auto iter_in_band : iter_band) diff --git a/source/source_hsolver/diago_cg.h b/source/source_hsolver/diago_cg.h index c5109cfb0cf..a0dd8502131 100644 --- a/source/source_hsolver/diago_cg.h +++ b/source/source_hsolver/diago_cg.h @@ -1,21 +1,19 @@ #ifndef MODULE_HSOLVER_DIAGO_CG_H_ #define MODULE_HSOLVER_DIAGO_CG_H_ -#include - -#include -#include - #include "source_hsolver/diag_comm_info.h" #include "source_hsolver/hs_operator.h" #include #include +#include +#include +#include -namespace hsolver { +namespace hsolver +{ -template -class DiagoCG final +template class DiagoCG final { // private: accessibility within class is private by default // Note GetTypeReal::type will @@ -23,34 +21,36 @@ class DiagoCG final // otherwise return the real type of T(complex, std::complex) using Real = typename GetTypeReal::type; using ct_Device = typename ct::PsiToContainer::type; + public: // Constructor need: // 1. basis type and calculation type of ABACUS // 2. diag_comm: the pool over which the subspace matrices are reduced - DiagoCG(const std::string& basis_type, const std::string& calculation); + DiagoCG(const std::string &basis_type, const std::string &calculation); DiagoCG( - const std::string& basis_type, - const std::string& calculation, - const bool& need_subspace, - const diag_comm_info& diag_comm, - const Real& pw_diag_thr, - const int& pw_diag_nmax); + const std::string &basis_type, + const std::string &calculation, + const bool &need_subspace, + const diag_comm_info &diag_comm, + const Real &pw_diag_thr, + const int &pw_diag_nmax); ~DiagoCG(); // this is the diag() function for CG method // returns avg_iter - double diag(const HSOperator& op, - const int ld_psi, - const int nband, - const int dim, - T* psi_in, - Real* eigenvalue_in, - const std::vector& ethr_band, - const Real* prec = nullptr); + double diag( + const HSOperator &op, + const int ld_psi, + const int nband, + const int dim, + T *psi_in, + Real *eigenvalue_in, + const std::vector ðr_band, + const Real *prec = nullptr); private: - Device * ctx_ = {}; + Device *ctx_ = {}; /// static variables, used for passing control variables /// record for how many bands not have convergence eigenvalues int notconv_ = 0; @@ -76,61 +76,71 @@ class DiagoCG final bool need_subspace_ = false; /// The H and S operator being diagonalized, set for the duration of diag(). - const HSOperator* op_ = nullptr; + const HSOperator *op_ = nullptr; void calc_grad( - const ct::Tensor& prec, - ct::Tensor& grad, - ct::Tensor& hphi, - ct::Tensor& sphi, - ct::Tensor& pphi); + const ct::Tensor &prec, + ct::Tensor &grad, + ct::Tensor &hphi, + ct::Tensor &sphi, + ct::Tensor &pphi); void orth_grad( - const ct::Tensor& psi, - const int& m, - ct::Tensor& grad, - ct::Tensor& scg, - ct::Tensor& lagrange); + const ct::Tensor &psi, + const int &m, + ct::Tensor &grad, + ct::Tensor &scg, + ct::Tensor &lagrange); void calc_gamma_cg( - const int& iter, - const Real& cg_norm, - const Real& theta, - const ct::Tensor& prec, - const ct::Tensor& scg, - const ct::Tensor& grad, - const ct::Tensor& phi_m, - Real& gg_last, - ct::Tensor& g0, - ct::Tensor& cg); + const int &iter, + const Real &cg_norm, + const Real &theta, + const ct::Tensor &prec, + const ct::Tensor &scg, + const ct::Tensor &grad, + const ct::Tensor &phi_m, + Real &gg_last, + ct::Tensor &g0, + ct::Tensor &cg); bool update_psi( - const ct::Tensor& pphi, - const ct::Tensor& cg, - const ct::Tensor& scg, - const double& ethreshold, + const ct::Tensor &pphi, + const ct::Tensor &cg, + const ct::Tensor &scg, + const double ðreshold, Real &cg_norm, Real &theta, Real &eigen, - ct::Tensor& phi_m, - ct::Tensor& sphi, - ct::Tensor& hphi); + ct::Tensor &phi_m, + ct::Tensor &sphi, + ct::Tensor &hphi); - void schmit_orth(const int& m, const ct::Tensor& psi, const ct::Tensor& sphi, ct::Tensor& phi_m); + void schmit_orth( + const int &m, + const ct::Tensor &psi, + const ct::Tensor &sphi, + ct::Tensor &phi_m); // used in diag() for template replace Hamilt with Hamilt_PW - void diag_once(const ct::Tensor& prec, - ct::Tensor& psi, - ct::Tensor& eigen, - const std::vector& ethr_band); + void diag_once( + const ct::Tensor &prec, + ct::Tensor &psi, + ct::Tensor &eigen, + const std::vector ðr_band); - bool test_exit_cond(const int& ntry, const int& notconv) const; + bool test_exit_cond(const int &ntry, const int ¬conv) const; /// subspace rotation of the current nband vectors (packed, leading dimension dim) - void diag_subspace(const T* psi_in, T* psi_out, const int dim, const int nband, const bool S_orth); + void diag_subspace( + const T *psi_in, + T *psi_out, + const int dim, + const int nband, + const bool S_orth); using dot_real_op = ModuleBase::dot_real_op; - const T * one_ = nullptr, * zero_ = nullptr, * neg_one_ = nullptr; + const T *one_ = nullptr, *zero_ = nullptr, *neg_one_ = nullptr; }; } // namespace hsolver diff --git a/source/source_hsolver/diago_cusolver.cpp b/source/source_hsolver/diago_cusolver.cpp index b597e40ac72..bce24c56026 100644 --- a/source/source_hsolver/diago_cusolver.cpp +++ b/source/source_hsolver/diago_cusolver.cpp @@ -1,10 +1,10 @@ #include "diago_cusolver.h" -#include "source_base/module_external/blas_connector.h" #include "source_base/module_external/blacs_connector.h" +#include "source_base/module_external/blas_connector.h" #include "source_base/module_external/scalapack_connector.h" -#include "source_base/tool_title.h" #include "source_base/timer.h" +#include "source_base/tool_title.h" #include #include @@ -16,37 +16,49 @@ using complex = std::complex; namespace hsolver { // Initialize the DecomposedState variable for real and complex numbers -template -int DiagoCusolver::DecomposedState = 0; +template int DiagoCusolver::DecomposedState = 0; template DiagoCusolver::DiagoCusolver(const int nlocal_in, const int nbands_in) : nlocal(nlocal_in), nbands(nbands_in) { } -template -DiagoCusolver::~DiagoCusolver() -{ -} +template DiagoCusolver::~DiagoCusolver() {} // Diagonalization function template void DiagoCusolver::diag( - ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in) + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in) { ModuleBase::TITLE("DiagoCusolver", "diag"); ModuleBase::timer::start("DiagoCusolver", "cusolver"); // Allocate memory for eigenvalues std::vector eigen(this->nlocal, 0.0); std::vector eigenvectors(h_mat.row * h_mat.col); - this->dc.Dngvd(h_mat.row, h_mat.col, h_mat.p, s_mat.p, eigen.data(), eigenvectors.data()); + this->dc.Dngvd( + h_mat.row, + h_mat.col, + h_mat.p, + s_mat.p, + eigen.data(), + eigenvectors.data()); const int size = psi.get_nbands() * psi.get_nbasis(); - BlasConnector::copy(size, eigenvectors.data(), 1, psi.get_pointer(), 1); + BlasConnector::copy( + size, + eigenvectors.data(), + 1, + psi.get_pointer(), + 1); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); ModuleBase::timer::end("DiagoCusolver", "cusolver"); } diff --git a/source/source_hsolver/diago_cusolver.h b/source/source_hsolver/diago_cusolver.h index d59daf87442..9c3077c8273 100644 --- a/source/source_hsolver/diago_cusolver.h +++ b/source/source_hsolver/diago_cusolver.h @@ -1,7 +1,7 @@ #ifndef DIAGOCUSOLVER_H #define DIAGOCUSOLVER_H -#include "source_base/macros.h" // GetRealType +#include "source_base/macros.h" // GetRealType #include "source_base/matrix_block.h" #include "source_hsolver/kernels/cuda/diag_cusolver.cuh" #include "source_psi/psi.h" @@ -10,15 +10,13 @@ namespace hsolver { // DiagoCusolver class for diagonalization using CUSOLVER -template -class DiagoCusolver +template class DiagoCusolver { private: // Real is the real part of the complex type T using Real = typename GetTypeReal::type; public: - /// @param nlocal_in global dimension of the NAO Hamiltonian /// @param nbands_in number of lowest eigenpairs to compute DiagoCusolver(const int nlocal_in, const int nbands_in); @@ -26,10 +24,10 @@ class DiagoCusolver // Override the diag function for CUSOLVER diagonalization void diag( - ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in); + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in); // Static variable to keep track of the decomposition state static int DecomposedState; @@ -40,7 +38,7 @@ class DiagoCusolver private: #ifdef __MPI // Function to check if ELPA handle needs to be created or reused in MPI settings - bool ifElpaHandle(const bool& newIteration, const bool& ifNSCF) const; + bool ifElpaHandle(const bool &newIteration, const bool &ifNSCF) const; #endif const int nlocal; diff --git a/source/source_hsolver/diago_cusolvermp.cpp b/source/source_hsolver/diago_cusolvermp.cpp index ec6ce1e7917..5e5efc22262 100644 --- a/source/source_hsolver/diago_cusolvermp.cpp +++ b/source/source_hsolver/diago_cusolvermp.cpp @@ -12,10 +12,11 @@ using complex = std::complex; namespace hsolver { template -void DiagoCusolverMP::diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in) +void DiagoCusolverMP::diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in) { ModuleBase::TITLE("DiagoCusolverMP", "diag"); @@ -24,16 +25,30 @@ void DiagoCusolverMP::diag(ModuleBase::MatrixBlock& h_mat, MPI_Comm COMM_DIAG = MPI_COMM_WORLD; // use all processes { - Diag_CusolverMP_gvd es(COMM_DIAG, (const int)h_mat.row, (const int)h_mat.col, (const int*)h_mat.desc); + Diag_CusolverMP_gvd es(COMM_DIAG, (const int)h_mat.row, (const int)h_mat.col, (const int *)h_mat.desc); ModuleBase::timer::start("DiagoCusolverMP", "Diag_CusolverMP_gvd"); - es.generalized_eigenvector(h_mat.p, s_mat.p, eigen.data(), eigenvectors.data()); + es.generalized_eigenvector( + h_mat.p, + s_mat.p, + eigen.data(), + eigenvectors.data()); ModuleBase::timer::end("DiagoCusolverMP", "Diag_CusolverMP_gvd"); } const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); const int size = psi.get_nbands() * psi.get_nbasis(); - BlasConnector::copy(size, eigenvectors.data(), inc, psi.get_pointer(), inc); + BlasConnector::copy( + size, + eigenvectors.data(), + inc, + psi.get_pointer(), + inc); } template class DiagoCusolverMP; template class DiagoCusolverMP; diff --git a/source/source_hsolver/diago_cusolvermp.h b/source/source_hsolver/diago_cusolvermp.h index 3794cf1c084..2ea2e07cc91 100644 --- a/source/source_hsolver/diago_cusolvermp.h +++ b/source/source_hsolver/diago_cusolvermp.h @@ -10,8 +10,7 @@ namespace hsolver { // DiagoCusolverMP class, for diagonalization using CUSOLVERMP -template -class DiagoCusolverMP +template class DiagoCusolverMP { private: using Real = typename GetTypeReal::type; @@ -19,14 +18,13 @@ class DiagoCusolverMP public: /// @param nlocal_in global dimension of the NAO Hamiltonian /// @param nbands_in number of lowest eigenpairs to compute - DiagoCusolverMP(const int nlocal_in, const int nbands_in) : nlocal(nlocal_in), nbands(nbands_in) - { - } + DiagoCusolverMP(const int nlocal_in, const int nbands_in) : nlocal(nlocal_in), nbands(nbands_in) {} // the diag function for CUSOLVERMP diagonalization - void diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in); + void diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in); private: const int nlocal; diff --git a/source/source_hsolver/diago_dav_subspace.cpp b/source/source_hsolver/diago_dav_subspace.cpp index 9b4764fc9d5..546105e80f8 100644 --- a/source/source_hsolver/diago_dav_subspace.cpp +++ b/source/source_hsolver/diago_dav_subspace.cpp @@ -1,45 +1,46 @@ -#include #include "diago_dav_subspace.h" #include "diago_iter_assist.h" - +#include "source_base/kernels/dsp/dsp_connector.h" +#include "source_base/kernels/math_kernel_op.h" #include "source_base/module_device/device.h" #include "source_base/parallel_cell.h" #include "source_base/timer.h" #include "source_base/tool_quit.h" -#include "source_base/kernels/math_kernel_op.h" -#include "source_base/kernels/dsp/dsp_connector.h" -// #include "source_base/module_container/ATen/kernels/lapack.h" -#include +#include +// #include "source_base/module_container/ATen/kernels/lapack.h" -#include "source_hsolver/kernels/hegvd_op.h" #include "source_hsolver/diag_hs_para.h" #include "source_hsolver/kernels/bpcg_kernel_op.h" // normalize_op, precondition_op, apply_eigenvalues_op +#include "source_hsolver/kernels/hegvd_op.h" +#include #include #ifdef __MPI -#include #include "source_base/parallel_comm.h" #include "source_base/parallel_device.h" + +#include #endif using namespace hsolver; template -Diago_DavSubspace::Diago_DavSubspace(const std::vector& precondition_in, - const int& nband_in, - const int& nbasis_in, - const int& david_ndim_in, - const double& diag_thr_in, - const int& diag_nmax_in, - const diag_comm_info& diag_comm_in, - const int diag_subspace_in, - const int diago_subspace_bs_in) +Diago_DavSubspace::Diago_DavSubspace( + const std::vector &precondition_in, + const int &nband_in, + const int &nbasis_in, + const int &david_ndim_in, + const double &diag_thr_in, + const int &diag_nmax_in, + const diag_comm_info &diag_comm_in, + const int diag_subspace_in, + const int diago_subspace_bs_in) : precondition(precondition_in), n_band(nband_in), dim(nbasis_in), nbase_x(nband_in * david_ndim_in), - diag_thr(diag_thr_in), iter_nmax(diag_nmax_in), diag_comm(diag_comm_in), - diag_subspace(diag_subspace_in), diago_subspace_bs(diago_subspace_bs_in) + diag_thr(diag_thr_in), iter_nmax(diag_nmax_in), diag_comm(diag_comm_in), diag_subspace(diag_subspace_in), + diago_subspace_bs(diago_subspace_bs_in) { this->device = base_device::get_device_type(this->ctx); @@ -96,8 +97,7 @@ Diago_DavSubspace::Diago_DavSubspace(const std::vector& precond #endif } -template -Diago_DavSubspace::~Diago_DavSubspace() +template Diago_DavSubspace::~Diago_DavSubspace() { delmem_complex_op()(this->psi_in_iter); @@ -118,11 +118,12 @@ Diago_DavSubspace::~Diago_DavSubspace() } template -int Diago_DavSubspace::diag_once(const HSOperator& op, - T* psi_in, - const int psi_in_dmax, - Real* eigenvalue_in_hsolver, - const std::vector& ethr_band) +int Diago_DavSubspace::diag_once( + const HSOperator &op, + T *psi_in, + const int psi_in_dmax, + Real *eigenvalue_in_hsolver, + const std::vector ðr_band) { ModuleBase::timer::start("Diago_DavSubspace", "diag_once"); const int matrix_ld = std::max(this->dim, 1); @@ -145,7 +146,13 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, ModuleBase::timer::start("Diago_DavSubspace", "first"); - syncmem_complex_2d_op()(this->psi_in_iter, this->dim, psi_in, psi_in_dmax, this->dim, this->n_band); + syncmem_complex_2d_op()( + this->psi_in_iter, + this->dim, + psi_in, + psi_in_dmax, + this->dim, + this->n_band); for (int m = 0; m < this->n_band; m++) { unconv[m] = m; @@ -154,15 +161,38 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, // compute h*psi_in_iter // NOTE: bands after the first n_band should yield zero // hpsi[:, 0:nbase_x] = H * psi_in_iter[:, 0:nbase_x] - op.hpsi(this->psi_in_iter, this->hpsi, this->dim, this->notconv); + op.hpsi( + this->psi_in_iter, + this->hpsi, + this->dim, + this->notconv); // compute s*psi_in_iter // spsi[:, 0:nbase_x] = S * psi_in_iter[:, 0:nbase_x] - op.spsi(this->psi_in_iter, this->spsi, this->dim, this->notconv); - - this->cal_elem(this->dim, nbase, this->notconv, this->psi_in_iter, this->spsi, this->hpsi, this->hcc, this->scc); - - this->diag_zhegvx(nbase, this->notconv, this->hcc, this->scc, this->nbase_x, &eigenvalue_iter, this->vcc); + op.spsi( + this->psi_in_iter, + this->spsi, + this->dim, + this->notconv); + + this->cal_elem( + this->dim, + nbase, + this->notconv, + this->psi_in_iter, + this->spsi, + this->hpsi, + this->hcc, + this->scc); + + this->diag_zhegvx( + nbase, + this->notconv, + this->hcc, + this->scc, + this->nbase_x, + &eigenvalue_iter, + this->vcc); for (size_t m = 0; m < this->n_band; m++) { @@ -177,27 +207,36 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, { dav_iter++; - this->cal_grad(op, - this->dim, - nbase, - this->notconv, - this->psi_in_iter, - this->hpsi, - this->spsi, - this->vcc, - unconv.data(), - &eigenvalue_iter); - - this->cal_elem(this->dim, - nbase, - this->notconv, - this->psi_in_iter, - this->spsi, - this->hpsi, - this->hcc, - this->scc); - - this->diag_zhegvx(nbase, this->n_band, this->hcc, this->scc, this->nbase_x, &eigenvalue_iter, this->vcc); + this->cal_grad( + op, + this->dim, + nbase, + this->notconv, + this->psi_in_iter, + this->hpsi, + this->spsi, + this->vcc, + unconv.data(), + &eigenvalue_iter); + + this->cal_elem( + this->dim, + nbase, + this->notconv, + this->psi_in_iter, + this->spsi, + this->hpsi, + this->hcc, + this->scc); + + this->diag_zhegvx( + nbase, + this->n_band, + this->hcc, + this->scc, + this->nbase_x, + &eigenvalue_iter, + this->vcc); // check convergence and update eigenvalues ModuleBase::timer::start("Diago_DavSubspace", "check_update"); @@ -232,18 +271,18 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, ModuleBase::gemm_op() #endif ('N', - 'N', - this->dim, - this->n_band, - nbase, - this->one, - this->psi_in_iter, - matrix_ld, - this->vcc, - this->nbase_x, - this->zero, - psi_in, - psi_ld); + 'N', + this->dim, + this->n_band, + nbase, + this->one, + this->psi_in_iter, + matrix_ld, + this->vcc, + this->nbase_x, + this->zero, + psi_in, + psi_ld); if (!this->notconv || (dav_iter == this->iter_nmax)) { @@ -259,18 +298,25 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, // estimate of the eigenvectors and set the basis dimension to N; // update this->psi_in_iter according to psi_in - syncmem_complex_2d_op()(this->psi_in_iter, this->dim, psi_in, psi_in_dmax, this->dim, this->n_band); - - this->refresh(this->dim, - this->n_band, - nbase, - eigenvalue_in_hsolver, - this->psi_in_iter, - this->hpsi, - this->spsi, - this->hcc, - this->scc, - this->vcc); + syncmem_complex_2d_op()( + this->psi_in_iter, + this->dim, + psi_in, + psi_in_dmax, + this->dim, + this->n_band); + + this->refresh( + this->dim, + this->n_band, + nbase, + eigenvalue_in_hsolver, + this->psi_in_iter, + this->hpsi, + this->spsi, + this->hcc, + this->scc, + this->vcc); ModuleBase::timer::end("Diago_DavSubspace", "last"); } @@ -284,16 +330,17 @@ int Diago_DavSubspace::diag_once(const HSOperator& op, } template -void Diago_DavSubspace::cal_grad(const HSOperator& op, - const int& dim, - const int& nbase, - const int& notconv, - T* psi_iter, - T* hpsi, - T* spsi, - T* vcc, - const int* unconv, - std::vector* eigenvalue_iter) +void Diago_DavSubspace::cal_grad( + const HSOperator &op, + const int &dim, + const int &nbase, + const int ¬conv, + T *psi_iter, + T *hpsi, + T *spsi, + T *vcc, + const int *unconv, + std::vector *eigenvalue_iter) { ModuleBase::timer::start("Diago_DavSubspace", "cal_grad"); const int matrix_ld = std::max(this->dim, 1); @@ -307,51 +354,52 @@ void Diago_DavSubspace::cal_grad(const HSOperator& op, } } - if (notconv > 1){ + if (notconv > 1) + { #ifdef __DSP - ModuleBase::gemm_op_mt() + ModuleBase::gemm_op_mt() #else - ModuleBase::gemm_op() + ModuleBase::gemm_op() #endif - ('N', - 'N', - this->dim, - notconv, - nbase, - this->one, - hpsi, - matrix_ld, - vcc, - this->nbase_x, - this->zero, - psi_iter + (nbase) * this->dim, - matrix_ld); - } else + ('N', + 'N', + this->dim, + notconv, + nbase, + this->one, + hpsi, + matrix_ld, + vcc, + this->nbase_x, + this->zero, + psi_iter + (nbase) * this->dim, + matrix_ld); + } + else { #ifdef __DSP - ModuleBase::gemv_op_mt() + ModuleBase::gemv_op_mt() #else - ModuleBase::gemv_op() + ModuleBase::gemv_op() #endif - ('N', - this->dim, // m: row of A - nbase, // n: col of A - this->one, // alpha - hpsi, // A dim * nbase - matrix_ld, // LDA: if(N) max(1,m) - vcc, // X nbase - 1, // incx - this->zero, // beta - psi_iter + (nbase) * this->dim, // Y dim - 1 // incy - ); + ('N', + this->dim, // m: row of A + nbase, // n: col of A + this->one, // alpha + hpsi, // A dim * nbase + matrix_ld, // LDA: if(N) max(1,m) + vcc, // X nbase + 1, // incx + this->zero, // beta + psi_iter + (nbase) * this->dim, // Y dim + 1 // incy + ); } - // Eigenvalues operation section - Real* e_temp_hd = eigenvalue_iter->data(); + Real *e_temp_hd = eigenvalue_iter->data(); if (this->device == base_device::GpuDevice) { syncmem_var_h2d_op()(this->d_eigenvalue, eigenvalue_iter->data(), this->nbase_x); @@ -359,87 +407,99 @@ void Diago_DavSubspace::cal_grad(const HSOperator& op, } // vcc = - vcc * eigenvalue - ModuleBase::matrix_mul_vector_op()(nbase, notconv, vcc, this->nbase_x, e_temp_hd, -1.0, vcc, this->nbase_x); - - if (notconv > 1){ + ModuleBase::matrix_mul_vector_op()( + nbase, + notconv, + vcc, + this->nbase_x, + e_temp_hd, + -1.0, + vcc, + this->nbase_x); + + if (notconv > 1) + { #ifdef __DSP - ModuleBase::gemm_op_mt() + ModuleBase::gemm_op_mt() #else - ModuleBase::gemm_op() + ModuleBase::gemm_op() #endif - ('N', - 'N', - this->dim, - notconv, - nbase, - this->one, - spsi, - matrix_ld, - vcc, - this->nbase_x, - this->one, - psi_iter + nbase * this->dim, - matrix_ld); - } else + ('N', + 'N', + this->dim, + notconv, + nbase, + this->one, + spsi, + matrix_ld, + vcc, + this->nbase_x, + this->one, + psi_iter + nbase * this->dim, + matrix_ld); + } + else { #ifdef __DSP - ModuleBase::gemv_op_mt() + ModuleBase::gemv_op_mt() #else - ModuleBase::gemv_op() + ModuleBase::gemv_op() #endif - ('N', - this->dim, // m: row of A - nbase, // n: col of A - this->one, // alpha - spsi, // A dim * nbase - matrix_ld, // LDA: if(N) max(1,m) - vcc, // X nbase - 1, // incx - this->one, // beta - psi_iter + nbase * this->dim, // Y dim - 1 // incy - ); + ('N', + this->dim, // m: row of A + nbase, // n: col of A + this->one, // alpha + spsi, // A dim * nbase + matrix_ld, // LDA: if(N) max(1,m) + vcc, // X nbase + 1, // incx + this->one, // beta + psi_iter + nbase * this->dim, // Y dim + 1 // incy + ); } // Precondition section #if defined(__CUDA) || defined(__ROCM) if (this->device == base_device::GpuDevice) { - precondition_op()(this->dim, - psi_iter, - nbase, - notconv, - d_precondition, - this->d_eigenvalue); + precondition_op()( + this->dim, + psi_iter, + nbase, + notconv, + d_precondition, + this->d_eigenvalue); } else #endif { - precondition_op()(this->dim, - psi_iter, - nbase, - notconv, - this->precondition.data(), - (*eigenvalue_iter).data()); + precondition_op()( + this->dim, + psi_iter, + nbase, + notconv, + this->precondition.data(), (*eigenvalue_iter).data()); } // Normalize section #if defined(__CUDA) || defined(__ROCM) if (this->device == base_device::GpuDevice) { - Real* psi_norm = nullptr; + Real *psi_norm = nullptr; resmem_real_op()(psi_norm, notconv); setmem_real_op()(psi_norm, 0.0, notconv); - normalize_op()(this->dim, - psi_iter, - nbase, - notconv, - psi_norm); + normalize_op()( + this->dim, + psi_iter, + nbase, + notconv, + psi_norm); // Check for zero norms (GPU path: copy norms from device to host) - Real* psi_norm_host = nullptr; + Real *psi_norm_host = nullptr; resmem_real_h_op()(psi_norm_host, notconv); syncmem_var_d2h_op()(psi_norm_host, psi_norm, notconv); for (int i = 0; i < notconv; i++) @@ -461,15 +521,16 @@ void Diago_DavSubspace::cal_grad(const HSOperator& op, else #endif { - Real* psi_norm = nullptr; + Real *psi_norm = nullptr; resmem_real_h_op()(psi_norm, notconv); setmem_real_h_op()(psi_norm, 0.0, notconv); - normalize_op()(this->dim, - psi_iter, - nbase, - notconv, - psi_norm); + normalize_op()( + this->dim, + psi_iter, + nbase, + notconv, + psi_norm); // Check for zero norms (CPU path) for (int i = 0; i < notconv; i++) @@ -489,22 +550,31 @@ void Diago_DavSubspace::cal_grad(const HSOperator& op, // update hpsi[:, nbase:nbase+notconv] // hpsi[:, nbase:nbase+notconv] = H * psi_iter[:, nbase:nbase+notconv] - op.hpsi(psi_iter + nbase * dim, hpsi + nbase * this->dim, this->dim, notconv); - op.spsi(psi_iter + nbase * dim, spsi + nbase * this->dim, this->dim, notconv); + op.hpsi( + psi_iter + nbase * dim, + hpsi + nbase * this->dim, + this->dim, + notconv); + op.spsi( + psi_iter + nbase * dim, + spsi + nbase * this->dim, + this->dim, + notconv); ModuleBase::timer::end("Diago_DavSubspace", "cal_grad"); return; } template -void Diago_DavSubspace::cal_elem(const int& dim, - int& nbase, - const int& notconv, - const T* psi_iter, - const T* spsi, - const T* hpsi, - T* hcc, - T* scc) +void Diago_DavSubspace::cal_elem( + const int &dim, + int &nbase, + const int ¬conv, + const T *psi_iter, + const T *spsi, + const T *hpsi, + T *hcc, + T *scc) { ModuleBase::timer::start("Diago_DavSubspace", "cal_elem"); const int matrix_ld = std::max(this->dim, 1); @@ -512,98 +582,104 @@ void Diago_DavSubspace::cal_elem(const int& dim, if (this->dim == 0) { const int contribution_size = this->nbase_x * notconv; - T* const h_contribution = hcc + nbase * this->nbase_x; - T* const s_contribution = scc + nbase * this->nbase_x; + T *const h_contribution = hcc + nbase * this->nbase_x; + T *const s_contribution = scc + nbase * this->nbase_x; // GEMV does not overwrite its output when the local row count is zero. setmem_complex_op()(h_contribution, 0, contribution_size); setmem_complex_op()(s_contribution, 0, contribution_size); } - if (notconv > 1){ + if (notconv > 1) + { #ifdef __DSP - ModuleBase::gemm_op_mt() + ModuleBase::gemm_op_mt() #else - ModuleBase::gemm_op() + ModuleBase::gemm_op() #endif - ('C', - 'N', - nbase + notconv, - notconv, - this->dim, - this->one, - psi_iter, - matrix_ld, - &hpsi[nbase * this->dim], - matrix_ld, - this->zero, - &hcc[nbase * this->nbase_x], - this->nbase_x); + ('C', + 'N', + nbase + notconv, + notconv, + this->dim, + this->one, + psi_iter, + matrix_ld, + &hpsi[nbase * this->dim], + matrix_ld, + this->zero, + &hcc[nbase * this->nbase_x], + this->nbase_x); #ifdef __DSP - ModuleBase::gemm_op_mt() + ModuleBase::gemm_op_mt() #else - ModuleBase::gemm_op() + ModuleBase::gemm_op() #endif - ('C', - 'N', - nbase + notconv, - notconv, - this->dim, - this->one, - psi_iter, - matrix_ld, - spsi + nbase * this->dim, - matrix_ld, - this->zero, - &scc[nbase * this->nbase_x], - this->nbase_x); - - } else { + ('C', + 'N', + nbase + notconv, + notconv, + this->dim, + this->one, + psi_iter, + matrix_ld, + spsi + nbase * this->dim, + matrix_ld, + this->zero, + &scc[nbase * this->nbase_x], + this->nbase_x); + } + else + { #ifdef __DSP - ModuleBase::gemv_op_mt() + ModuleBase::gemv_op_mt() #else - ModuleBase::gemv_op() + ModuleBase::gemv_op() #endif - ('C', - this->dim, // m: row of A - nbase + notconv, // n: col of A - this->one, // alpha - psi_iter, // A dim * nbase - matrix_ld, // LDA: if(N) max(1,m) - &hpsi[nbase * this->dim], // X nbase - 1, // incx - this->zero, // beta - &hcc[nbase * this->nbase_x], // Y dim - 1 // incy - ); + ('C', + this->dim, // m: row of A + nbase + notconv, // n: col of A + this->one, // alpha + psi_iter, // A dim * nbase + matrix_ld, // LDA: if(N) max(1,m) + &hpsi[nbase * this->dim], // X nbase + 1, // incx + this->zero, // beta + &hcc[nbase * this->nbase_x], // Y dim + 1 // incy + ); #ifdef __DSP - ModuleBase::gemv_op_mt() + ModuleBase::gemv_op_mt() #else - ModuleBase::gemv_op() + ModuleBase::gemv_op() #endif - ('C', - this->dim, // m: row of A - nbase + notconv, // n: col of A - this->one, // alpha - psi_iter, // A dim * nbase - matrix_ld, // LDA: if(N) max(1,m) - spsi + nbase * this->dim, // X nbase - 1, // incx - this->zero, // beta - &scc[nbase * this->nbase_x], // Y dim - 1 // incy - ); - - } - + ('C', + this->dim, // m: row of A + nbase + notconv, // n: col of A + this->one, // alpha + psi_iter, // A dim * nbase + matrix_ld, // LDA: if(N) max(1,m) + spsi + nbase * this->dim, // X nbase + 1, // incx + this->zero, // beta + &scc[nbase * this->nbase_x], // Y dim + 1 // incy + ); + } #ifdef __MPI if (this->diag_comm.nproc > 1) { #ifdef __DSP // Only on dsp hardware need an extra space to reduce data - mtfunc::dsp_dav_subspace_reduce(hcc, scc, nbase, this->nbase_x, this->notconv, this->diag_comm.comm); + mtfunc::dsp_dav_subspace_reduce( + hcc, + scc, + nbase, + this->nbase_x, + this->notconv, + this->diag_comm.comm); #else assert(this->diag_comm.comm == POOL_WORLD); Parallel_Reduce::reduce_pool(hcc + nbase * this->nbase_x, notconv * this->nbase_x); @@ -620,13 +696,14 @@ void Diago_DavSubspace::cal_elem(const int& dim, } template -void Diago_DavSubspace::diag_zhegvx(const int& nbase, - const int& nband, - T* hcc, - T* scc, - const int& nbase_x, - std::vector* eigenvalue_iter, - T* vcc) +void Diago_DavSubspace::diag_zhegvx( + const int &nbase, + const int &nband, + T *hcc, + T *scc, + const int &nbase_x, + std::vector *eigenvalue_iter, + T *vcc) { ModuleBase::timer::start("Diago_DavSubspace", "diag_zhegvx"); assert(nbase_x >= std::max(1, nbase)); @@ -637,7 +714,13 @@ void Diago_DavSubspace::diag_zhegvx(const int& nbase, if (this->diag_comm.rank == 0) { syncmem_complex_op()(this->d_scc, scc, nbase * this->nbase_x); - ct::kernels::lapack_hegvd()(nbase, this->nbase_x, this->hcc, this->d_scc, this->d_eigenvalue, this->vcc); + ct::kernels::lapack_hegvd()( + nbase, + this->nbase_x, + this->hcc, + this->d_scc, + this->d_eigenvalue, + this->vcc); syncmem_var_d2h_op()((*eigenvalue_iter).data(), this->d_eigenvalue, this->nbase_x); } #endif @@ -659,14 +742,14 @@ void Diago_DavSubspace::diag_zhegvx(const int& nbase, s_diag[i][j] = scc[i * this->nbase_x + j]; } } - hegvx_op()(this->ctx, - nbase, - this->nbase_x, - this->hcc, - this->scc, - nband, - (*eigenvalue_iter).data(), - this->vcc); + hegvx_op()( + this->ctx, + nbase, + this->nbase_x, + this->hcc, + this->scc, + nband, (*eigenvalue_iter).data(), + this->vcc); // reset: for (size_t i = 0; i < nbase; i++) { @@ -706,15 +789,15 @@ void Diago_DavSubspace::diag_zhegvx(const int& nbase, } } } - diago_hs_para(h_diag.data(), - s_diag.data(), - nbase, - nband, - (*eigenvalue_iter).data(), - vcc_tmp.data(), - this->diag_comm.comm, - this->diag_subspace, - this->diago_subspace_bs); + diago_hs_para( + h_diag.data(), + s_diag.data(), + nbase, + nband, (*eigenvalue_iter).data(), + vcc_tmp.data(), + this->diag_comm.comm, + this->diag_subspace, + this->diago_subspace_bs); if (this->diag_comm.rank == 0) { for (size_t i = 0; i < nband; i++) @@ -739,9 +822,17 @@ void Diago_DavSubspace::diag_zhegvx(const int& nbase, for (int i = 0; i < nband; i++) { // fix a bug here when vcc is real: bcast according to the true type - Parallel_Common::bcast_data(&vcc[i * this->nbase_x], nbase, this->diag_comm.comm, 0); + Parallel_Common::bcast_data( + &vcc[i * this->nbase_x], + nbase, + this->diag_comm.comm, + 0); } - Parallel_Common::bcast_data((*eigenvalue_iter).data(), nband, this->diag_comm.comm, 0); + Parallel_Common::bcast_data( + (*eigenvalue_iter).data(), + nband, + this->diag_comm.comm, + 0); } #endif @@ -750,17 +841,18 @@ void Diago_DavSubspace::diag_zhegvx(const int& nbase, } template -void Diago_DavSubspace::refresh(const int& dim, - const int& nband, - int& nbase, - const Real* eigenvalue_in_hsolver, - // const psi::Psi& psi, - T* psi_iter, - T* hpsi, - T* spsi, - T* hcc, - T* scc, - T* vcc) +void Diago_DavSubspace::refresh( + const int &dim, + const int &nband, + int &nbase, + const Real *eigenvalue_in_hsolver, + // const psi::Psi& psi, + T *psi_iter, + T *hpsi, + T *spsi, + T *hcc, + T *scc, + T *vcc) { ModuleBase::timer::start("Diago_DavSubspace", "refresh"); const int matrix_ld = std::max(this->dim, 1); @@ -770,19 +862,19 @@ void Diago_DavSubspace::refresh(const int& dim, #else ModuleBase::gemm_op() #endif - ('N', - 'N', - this->dim, - nband, - nbase, - this->one, - this->hpsi, - matrix_ld, - this->vcc, - this->nbase_x, - this->zero, - psi_iter + nband * this->dim, - matrix_ld); + ('N', + 'N', + this->dim, + nband, + nbase, + this->one, + this->hpsi, + matrix_ld, + this->vcc, + this->nbase_x, + this->zero, + psi_iter + nband * this->dim, + matrix_ld); // update hpsi syncmem_complex_op()(hpsi, psi_iter + nband * this->dim, this->dim * nband); @@ -793,18 +885,18 @@ void Diago_DavSubspace::refresh(const int& dim, ModuleBase::gemm_op() #endif ('N', - 'N', - this->dim, - nband, - nbase, - this->one, - this->spsi, - matrix_ld, - this->vcc, - this->nbase_x, - this->zero, - psi_iter + nband * this->dim, - matrix_ld); + 'N', + this->dim, + nband, + nbase, + this->one, + this->spsi, + matrix_ld, + this->vcc, + this->nbase_x, + this->zero, + psi_iter + nband * this->dim, + matrix_ld); // update spsi syncmem_complex_op()(spsi, psi_iter + nband * this->dim, this->dim * nband); @@ -812,16 +904,38 @@ void Diago_DavSubspace::refresh(const int& dim, nbase = nband; // set hcc/scc/vcc to 0 - setmem_complex_2d_op()(hcc, this->nbase_x, 0, nbase, nbase); - setmem_complex_2d_op()(scc, this->nbase_x, 0, nbase, nbase); - setmem_complex_2d_op()(vcc, this->nbase_x, 0, nbase, nbase); + setmem_complex_2d_op()( + hcc, + this->nbase_x, + 0, + nbase, + nbase); + setmem_complex_2d_op()( + scc, + this->nbase_x, + 0, + nbase, + nbase); + setmem_complex_2d_op()( + vcc, + this->nbase_x, + 0, + nbase, + nbase); if (this->device == base_device::GpuDevice) { #if defined(__CUDA) || defined(__ROCM) syncmem_var_h2d_op()(this->d_eigenvalue, eigenvalue_in_hsolver, nbase); #endif - refresh_hcc_scc_vcc_op()(nbase, hcc, scc, vcc, this->nbase_x, this->d_eigenvalue, this->one_); + refresh_hcc_scc_vcc_op()( + nbase, + hcc, + scc, + vcc, + this->nbase_x, + this->d_eigenvalue, + this->one_); } else { @@ -838,12 +952,13 @@ void Diago_DavSubspace::refresh(const int& dim, } template -int Diago_DavSubspace::diag(const HSOperator& op, - T* psi_in, - const int psi_in_dmax, - Real* eigenvalue_in_hsolver, - const std::vector& ethr_band, - const bool& scf_type) +int Diago_DavSubspace::diag( + const HSOperator &op, + T *psi_in, + const int psi_in_dmax, + Real *eigenvalue_in_hsolver, + const std::vector ðr_band, + const bool &scf_type) { /// record the times of trying iterative diagonalization this->notconv = 0; @@ -854,7 +969,12 @@ int Diago_DavSubspace::diag(const HSOperator& op, do { - sum_iter += this->diag_once(op, psi_in, psi_in_dmax, eigenvalue_in_hsolver, ethr_band); + sum_iter += this->diag_once( + op, + psi_in, + psi_in_dmax, + eigenvalue_in_hsolver, + ethr_band); ++ntry; @@ -870,7 +990,7 @@ int Diago_DavSubspace::diag(const HSOperator& op, } template -bool Diago_DavSubspace::test_exit_cond(const int& ntry, const int& notconv, const bool& scf) +bool Diago_DavSubspace::test_exit_cond(const int &ntry, const int ¬conv, const bool &scf) { // scf = true; // scf // scf = false; // nscf diff --git a/source/source_hsolver/diago_dav_subspace.h b/source/source_hsolver/diago_dav_subspace.h index 4e1e97fbbbc..002a0b00c8f 100644 --- a/source/source_hsolver/diago_dav_subspace.h +++ b/source/source_hsolver/diago_dav_subspace.h @@ -1,12 +1,10 @@ #ifndef DIAGO_NEW_DAV_H #define DIAGO_NEW_DAV_H -#include "source_base/macros.h" // GetRealType -#include "source_base/module_device/device.h" // base_device -#include "source_base/module_device/memory_op.h"// base_device::memory" - +#include "source_base/macros.h" // GetRealType #include "source_base/module_container/ATen/kernels/lapack.h" - +#include "source_base/module_device/device.h" // base_device +#include "source_base/module_device/memory_op.h" // base_device::memory" #include "source_hsolver/diag_comm_info.h" #include "source_hsolver/hs_operator.h" @@ -15,8 +13,7 @@ namespace hsolver { -template , typename Device = base_device::DEVICE_CPU> -class Diago_DavSubspace +template , typename Device = base_device::DEVICE_CPU> class Diago_DavSubspace { private: // Note GetTypeReal::type will @@ -25,24 +22,26 @@ class Diago_DavSubspace using Real = typename GetTypeReal::type; public: - Diago_DavSubspace(const std::vector& precondition_in, - const int& nband_in, - const int& nbasis_in, - const int& david_ndim_in, - const double& diag_thr_in, - const int& diag_nmax_in, - const diag_comm_info& diag_comm_in, - const int diago_dav_method_in, - const int block_size_in); + Diago_DavSubspace( + const std::vector &precondition_in, + const int &nband_in, + const int &nbasis_in, + const int &david_ndim_in, + const double &diag_thr_in, + const int &diag_nmax_in, + const diag_comm_info &diag_comm_in, + const int diago_dav_method_in, + const int block_size_in); ~Diago_DavSubspace(); - int diag(const HSOperator& op, - T* psi_in, - const int psi_in_dmax, - Real* eigenvalue_in, - const std::vector& ethr_band, - const bool& scf_type); + int diag( + const HSOperator &op, + T *psi_in, + const int psi_in_dmax, + Real *eigenvalue_in, + const std::vector ðr_band, + const bool &scf_type); private: /// for MPI communication @@ -64,67 +63,70 @@ class Diago_DavSubspace const int nbase_x = 0; /// precondition for diag - const std::vector& precondition; - Real* d_precondition = nullptr; + const std::vector &precondition; + Real *d_precondition = nullptr; /// record for how many bands not have convergence eigenvalues int notconv = 0; - T* psi_in_iter = nullptr; + T *psi_in_iter = nullptr; /// the product of H and psi in the reduced basis set - T* hpsi = nullptr; + T *hpsi = nullptr; /// the product of S and psi in the reduced basis set - T* spsi = nullptr; + T *spsi = nullptr; /// Hamiltonian on the reduced basis - T* hcc = nullptr; + T *hcc = nullptr; /// Overlap on the reduced basis - T* scc = nullptr; + T *scc = nullptr; /// Eigenvectors on the reduced basis - T* vcc = nullptr; + T *vcc = nullptr; - T* d_scc = nullptr; - Real* d_eigenvalue = nullptr; + T *d_scc = nullptr; + Real *d_eigenvalue = nullptr; /// device type of psi - Device* ctx = {}; - base_device::DEVICE_CPU* cpu_ctx = {}; + Device *ctx = {}; + base_device::DEVICE_CPU *cpu_ctx = {}; base_device::AbacusDevice_t device = {}; - void cal_grad(const HSOperator& op, - const int& dim, - const int& nbase, - const int& notconv, - T* psi_iter, - T* hpsi, - T* spsi, - T* vcc, - const int* unconv, - std::vector* eigenvalue_iter); - - void cal_elem(const int& dim, - int& nbase, - const int& notconv, - const T* psi_iter, - const T* spsi, - const T* hpsi, - T* hcc, - T* scc); - - void refresh(const int& dim, - const int& nband, - int& nbase, - const Real* eigenvalue, - T* psi_iter, - T* hpsi, - T* spsi, - T* hcc, - T* scc, - T* vcc); + void cal_grad( + const HSOperator &op, + const int &dim, + const int &nbase, + const int ¬conv, + T *psi_iter, + T *hpsi, + T *spsi, + T *vcc, + const int *unconv, + std::vector *eigenvalue_iter); + + void cal_elem( + const int &dim, + int &nbase, + const int ¬conv, + const T *psi_iter, + const T *spsi, + const T *hpsi, + T *hcc, + T *scc); + + void refresh( + const int &dim, + const int &nband, + int &nbase, + const Real *eigenvalue, + T *psi_iter, + T *hpsi, + T *spsi, + T *hcc, + T *scc, + T *vcc); // void diagH_LAPACK(const int nstart, // const int nbands, @@ -134,23 +136,25 @@ class Diago_DavSubspace // Real* e, // T* vcc); - void diag_zhegvx(const int& nbase, - const int& nband, - T* hcc, - T* scc, - const int& nbase_x, - std::vector* eigenvalue_iter, - T* vcc); - - int diag_once(const HSOperator& op, - T* psi_in, - const int psi_in_dmax, - Real* eigenvalue_in, - const std::vector& ethr_band); - - bool test_exit_cond(const int& ntry, const int& notconv, const bool& scf); - - int diag_subspace; // 0: LAPACK, 1: Gen-ELPA, 2: ScaLAPACK + void diag_zhegvx( + const int &nbase, + const int &nband, + T *hcc, + T *scc, + const int &nbase_x, + std::vector *eigenvalue_iter, + T *vcc); + + int diag_once( + const HSOperator &op, + T *psi_in, + const int psi_in_dmax, + Real *eigenvalue_in, + const std::vector ðr_band); + + bool test_exit_cond(const int &ntry, const int ¬conv, const bool &scf); + + int diag_subspace; // 0: LAPACK, 1: Gen-ELPA, 2: ScaLAPACK int diago_subspace_bs = 0; // the block size in 2d block cyclic distribution if use elpa or scalapack #ifdef __DSP diff --git a/source/source_hsolver/diago_david.cpp b/source/source_hsolver/diago_david.cpp index 31e6fcfb0c2..ad4b880dba5 100644 --- a/source/source_hsolver/diago_david.cpp +++ b/source/source_hsolver/diago_david.cpp @@ -1,26 +1,26 @@ -#include #include "diago_david.h" -#include "source_base/parallel_cell.h" -#include "source_base/timer.h" // ModuleBase::timer::tick -#include "source_base/tool_title.h" // ModuleBase::TITLE -#include "source_base/module_device/device.h" - -#include "source_hsolver/kernels/hegvd_op.h" #include "source_base/kernels/math_kernel_op.h" +#include "source_base/module_device/device.h" +#include "source_base/parallel_cell.h" #include "source_base/parallel_comm.h" +#include "source_base/timer.h" // ModuleBase::timer::tick +#include "source_base/tool_title.h" // ModuleBase::TITLE +#include "source_hsolver/kernels/hegvd_op.h" +#include using namespace hsolver; - template -DiagoDavid::DiagoDavid(const Real* precondition_in, - const int nband_in, - const int dim_in, - const int david_ndim_in, - const diag_comm_info& diag_comm_in) - : nband(nband_in), dim(dim_in), nbase_x(david_ndim_in * nband_in), david_ndim(david_ndim_in), diag_comm(diag_comm_in) +DiagoDavid::DiagoDavid( + const Real *precondition_in, + const int nband_in, + const int dim_in, + const int david_ndim_in, + const diag_comm_info &diag_comm_in) + : nband(nband_in), dim(dim_in), nbase_x(david_ndim_in * nband_in), david_ndim(david_ndim_in), + diag_comm(diag_comm_in) { this->device = base_device::get_device_type(this->ctx); this->precondition = precondition_in; @@ -35,7 +35,6 @@ DiagoDavid::DiagoDavid(const Real* precondition_in, // 3: check the eigenvalues and errors of the last result // default: no check - // set auxiliary memory // !!! // allocate operaions MUST be paired with deallocate operations @@ -50,14 +49,16 @@ DiagoDavid::DiagoDavid(const Real* precondition_in, assert(this->nbase_x < global_dimension); /// initialize variables - /// k_first = 0 means that nks is more like a dimension of "basis" to be contracted in "HPsi".In LR-TDDFT the formula writes : + /// k_first = 0 means that nks is more like a dimension of "basis" to be contracted in "HPsi".In LR-TDDFT the + /// formula writes : /// $$\sum_{ jb\mathbf{k}'}A^I_{ia\mathbf{k}, jb\mathbf{k}' }X ^ I_{ jb\mathbf{k}'}$$ /// In the code : /// - "H" means A /// - "Psi" means X /// - "band" means the superscript I : the number of excited states to be solved /// - k : k-points, the same meaning as the ground state - /// - "basis" : number of occupied ks-orbitals(subscripts i,j) * number of unoccupied ks-orbitals(subscripts a,b), corresponding to "bands" of the ground state + /// - "basis" : number of occupied ks-orbitals(subscripts i,j) * number of unoccupied ks-orbitals(subscripts a,b), + /// corresponding to "bands" of the ground state // the lowest N eigenvalues base_device::memory::resize_memory_op()(this->eigenvalue, nbase_x, "DAV::eig"); @@ -100,8 +101,7 @@ DiagoDavid::DiagoDavid(const Real* precondition_in, #endif } -template -DiagoDavid::~DiagoDavid() +template DiagoDavid::~DiagoDavid() { delmem_complex_op()(this->basis); delmem_complex_op()(this->hpsi); @@ -120,14 +120,15 @@ DiagoDavid::~DiagoDavid() } template -int DiagoDavid::diag_once(const HSOperator& op, - const int dim, - const int nband, - const int ld_psi, - T *psi_in, - Real* eigenvalue_in, - const std::vector& ethr_band, - const int david_maxiter) +int DiagoDavid::diag_once( + const HSOperator &op, + const int dim, + const int nband, + const int ld_psi, + T *psi_in, + Real *eigenvalue_in, + const std::vector ðr_band, + const int david_maxiter) { if (test_david == 1) { @@ -146,7 +147,8 @@ int DiagoDavid::diag_once(const HSOperator& op, this->notconv = nband; // the number of unconverged eigenvalues - for (int m = 0; m < nband; m++) { + for (int m = 0; m < nband; m++) + { unconv[m] = m; } @@ -163,24 +165,33 @@ int DiagoDavid::diag_once(const HSOperator& op, { { // phm_in->sPsi(psi_in + m*ld_psi, &this->spsi[m * dim], dim, dim, 1); - op.spsi(psi_in + m*ld_psi,&this->spsi[m*dim],dim, 1); + op.spsi( + psi_in + m * ld_psi, + &this->spsi[m * dim], + dim, + 1); } } // begin SchmidtOrth for (int m = 0; m < nband; m++) { - syncmem_complex_op()(basis + dim*m, psi_in + m*ld_psi, dim); - - this->SchmidtOrth(dim, - nband, - m, - this->spsi, - &this->lagrange_matrix[m * nband], - pre_matrix_mm_m[m], - pre_matrix_mv_m[m]); + syncmem_complex_op()(basis + dim * m, psi_in + m * ld_psi, dim); + + this->SchmidtOrth( + dim, + nband, + m, + this->spsi, + &this->lagrange_matrix[m * nband], + pre_matrix_mm_m[m], + pre_matrix_mv_m[m]); { // phm_in->sPsi(basis + dim*m, &this->spsi[m * dim], dim, dim, 1); - op.spsi(basis + dim*m, &this->spsi[m * dim], dim, 1); + op.spsi( + basis + dim * m, + &this->spsi[m * dim], + dim, + 1); } } @@ -189,11 +200,28 @@ int DiagoDavid::diag_once(const HSOperator& op, // phm_in->ops->hPsi(dav_hpsi_in); // hpsi[:, 0:nband] = H basis[:, 0:nband] // slice index in this piece of code is in C manner. i.e. 0:id stands for [0,id) - op.hpsi(basis, hpsi, dim, nband); - - this->cal_elem(dim, nbase, nbase_x, this->notconv, this->hpsi, this->spsi, this->hcc); - - this->diag_zhegvx(nbase, nband, this->hcc, nbase_x, this->eigenvalue, this->vcc); + op.hpsi( + basis, + hpsi, + dim, + nband); + + this->cal_elem( + dim, + nbase, + nbase_x, + this->notconv, + this->hpsi, + this->spsi, + this->hcc); + + this->diag_zhegvx( + nbase, + nband, + this->hcc, + nbase_x, + this->eigenvalue, + this->vcc); for (int m = 0; m < nband; m++) { @@ -207,20 +235,34 @@ int DiagoDavid::diag_once(const HSOperator& op, { dav_iter++; - this->cal_grad(op, - dim, - nbase, - nbase_x, - this->notconv, - this->hpsi, - this->spsi, - this->vcc, - unconv.data(), - this->eigenvalue); - - this->cal_elem(dim, nbase, nbase_x, this->notconv, this->hpsi, this->spsi, this->hcc); - - this->diag_zhegvx(nbase, nband, this->hcc, nbase_x, this->eigenvalue, this->vcc); + this->cal_grad( + op, + dim, + nbase, + nbase_x, + this->notconv, + this->hpsi, + this->spsi, + this->vcc, + unconv.data(), + this->eigenvalue); + + this->cal_elem( + dim, + nbase, + nbase_x, + this->notconv, + this->hpsi, + this->spsi, + this->hcc); + + this->diag_zhegvx( + nbase, + nband, + this->hcc, + nbase_x, + this->eigenvalue, + this->vcc); // check convergence and update eigenvalues ModuleBase::timer::start("DiagoDavid", "check_update"); @@ -238,8 +280,7 @@ int DiagoDavid::diag_once(const HSOperator& op, } ModuleBase::timer::end("DiagoDavid", "check_update"); - if (!this->notconv || (nbase + this->notconv > nbase_x) - || (dav_iter == david_maxiter)) + if (!this->notconv || (nbase + this->notconv > nbase_x) || (dav_iter == david_maxiter)) { ModuleBase::timer::start("DiagoDavid", "last"); @@ -247,19 +288,20 @@ int DiagoDavid::diag_once(const HSOperator& op, setmem_complex_op()(psi_in, 0, nband * ld_psi); //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - ModuleBase::gemm_op()('N', - 'N', - dim, // m: row of A,C - nband, // n: col of B,C - nbase, // k: col of A, row of B - this->one, - basis, // A dim * nbase - matrix_ld, - this->vcc, // B nbase * nband - nbase_x, - this->zero, - psi_in, // C dim * nband - psi_ld); + ModuleBase::gemm_op()( + 'N', + 'N', + dim, // m: row of A,C + nband, // n: col of B,C + nbase, // k: col of A, row of B + this->one, + basis, // A dim * nbase + matrix_ld, + this->vcc, // B nbase * nband + nbase_x, + this->zero, + psi_in, // C dim * nband + psi_ld); if (!this->notconv || (dav_iter == david_maxiter)) { @@ -272,17 +314,18 @@ int DiagoDavid::diag_once(const HSOperator& op, // if the dimension of the reduced basis set is becoming too large, // then replace the first N (=nband) basis vectors with the current // estimate of the eigenvectors and set the basis dimension to N; - this->refresh(dim, - nband, - nbase, - nbase_x, - eigenvalue_in, - psi_in, - ld_psi, - this->hpsi, - this->spsi, - this->hcc, - this->vcc); + this->refresh( + dim, + nband, + nbase, + nbase_x, + eigenvalue_in, + psi_in, + ld_psi, + this->hpsi, + this->spsi, + this->hcc, + this->vcc); ModuleBase::timer::end("DiagoDavid", "last"); } @@ -295,24 +338,26 @@ int DiagoDavid::diag_once(const HSOperator& op, return dav_iter; } - template -void DiagoDavid::cal_grad(const HSOperator& op, - const int& dim, - const int& nbase, // current dimension of the reduced basis - const int nbase_x, // maximum dimension of the reduced basis set - const int& notconv, - T* hpsi, - T* spsi, - const T* vcc, - const int* unconv, - const Real* eigenvalue) +void DiagoDavid::cal_grad( + const HSOperator &op, + const int &dim, + const int &nbase, // current dimension of the reduced basis + const int nbase_x, // maximum dimension of the reduced basis set + const int ¬conv, + T *hpsi, + T *spsi, + const T *vcc, + const int *unconv, + const Real *eigenvalue) { - if (test_david == 1) { + if (test_david == 1) + { ModuleBase::TITLE("DiagoDavid", "cal_grad"); } // if all eigenpairs have converged, return - if (notconv == 0) { + if (notconv == 0) + { return; } ModuleBase::timer::start("DiagoDavid", "cal_grad"); @@ -328,7 +373,7 @@ void DiagoDavid::cal_grad(const HSOperator& op, // vc_ev_vector(notconv, nbase); // eigenvectors of unconverged index extracted from vcc - T* vc_ev_vector = nullptr; + T *vc_ev_vector = nullptr; resmem_complex_op()(vc_ev_vector, notconv * nbase); setmem_complex_op()(vc_ev_vector, 0, notconv * nbase); @@ -348,44 +393,45 @@ void DiagoDavid::cal_grad(const HSOperator& op, // vc_ev_vector[m * nbase + i] = vcc[i * nbase_x + unconv[m]]; for (int m = 0; m < notconv; m++) { - syncmem_complex_op()(vc_ev_vector + m * nbase, - vcc + unconv[m] * nbase_x, - nbase); + syncmem_complex_op()(vc_ev_vector + m * nbase, vcc + unconv[m] * nbase_x, nbase); } // basis[nbase] = hpsi * vc_ev_vector = hpsi*vcc // basis' = vc_ev_vector' * hpsi' // (dim, notconv) (dim, nbase) (nbase, notconv) - if (notconv == 1){ - //Reuse gemv for vector case to avoid potential bug using gemm call with n=1 - ModuleBase::gemv_op()('N', - dim, // m: row of A - nbase, // n: col of A - this->one, // alpha - hpsi, // A dim * nbase - matrix_ld, // LDA: if(N) max(1,m) - vc_ev_vector, // X nbase - 1, // incx - this->zero, // beta - basis + dim * nbase, // Y dim - 1 // incy + if (notconv == 1) + { + // Reuse gemv for vector case to avoid potential bug using gemm call with n=1 + ModuleBase::gemv_op()( + 'N', + dim, // m: row of A + nbase, // n: col of A + this->one, // alpha + hpsi, // A dim * nbase + matrix_ld, // LDA: if(N) max(1,m) + vc_ev_vector, // X nbase + 1, // incx + this->zero, // beta + basis + dim * nbase, // Y dim + 1 // incy ); - - }else + } + else { - ModuleBase::gemm_op()('N', - 'N', - dim, // m: row of A,C - notconv, // n: col of B,C - nbase, // k: col of A, row of B - this->one, // alpha - hpsi, // A dim * nbase - matrix_ld, // LDA: if(N) max(1,m) if(T) max(1,k) - vc_ev_vector, // B nbase * notconv - nbase, // LDB: if(N) max(1,k) if(T) max(1,n) - this->zero, // belta - basis + dim * nbase, // C dim * notconv - matrix_ld // LDC: if(N) max(1, m) + ModuleBase::gemm_op()( + 'N', + 'N', + dim, // m: row of A,C + notconv, // n: col of B,C + nbase, // k: col of A, row of B + this->one, // alpha + hpsi, // A dim * nbase + matrix_ld, // LDA: if(N) max(1,m) if(T) max(1,k) + vc_ev_vector, // B nbase * notconv + nbase, // LDB: if(N) max(1,k) if(T) max(1,n) + this->zero, // belta + basis + dim * nbase, // C dim * notconv + matrix_ld // LDC: if(N) max(1, m) ); } @@ -409,24 +455,26 @@ void DiagoDavid::cal_grad(const HSOperator& op, if (this->device == base_device::GpuDevice) { #if defined(__CUDA) || defined(__ROCM) - Real* e_temp_gpu = nullptr; + Real *e_temp_gpu = nullptr; resmem_var_op()(e_temp_gpu, nbase); syncmem_var_h2d_op()(e_temp_gpu, e_temp_cpu.data(), nbase); - ModuleBase::vector_mul_vector_op()(nbase, - vc_ev_vector + m * nbase, - vc_ev_vector + m * nbase, - e_temp_gpu, - false); + ModuleBase::vector_mul_vector_op()( + nbase, + vc_ev_vector + m * nbase, + vc_ev_vector + m * nbase, + e_temp_gpu, + false); delmem_var_op()(e_temp_gpu); #endif } else { - ModuleBase::vector_mul_vector_op()(nbase, - vc_ev_vector + m * nbase, - vc_ev_vector + m * nbase, - e_temp_cpu.data(), - false); + ModuleBase::vector_mul_vector_op()( + nbase, + vc_ev_vector + m * nbase, + vc_ev_vector + m * nbase, + e_temp_cpu.data(), + false); } } //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< @@ -436,35 +484,39 @@ void DiagoDavid::cal_grad(const HSOperator& op, // = (H - lambda * S) * psi * vcc // = (H - lambda * S) * psi_new //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - if (notconv == 1){ - //Use gemv for vector case to avoid potential bug using gemm call with n=1 - ModuleBase::gemv_op()('N', - dim, // m: row of A - nbase, // n: col of A - this->one, // alpha - spsi, // A dim * nbase - matrix_ld, // LDA: if(N) max(1,m) - vc_ev_vector, // X nbase - 1, // incx - this->one, // beta - basis + dim * nbase, // Y dim - 1 //incy + if (notconv == 1) + { + // Use gemv for vector case to avoid potential bug using gemm call with n=1 + ModuleBase::gemv_op()( + 'N', + dim, // m: row of A + nbase, // n: col of A + this->one, // alpha + spsi, // A dim * nbase + matrix_ld, // LDA: if(N) max(1,m) + vc_ev_vector, // X nbase + 1, // incx + this->one, // beta + basis + dim * nbase, // Y dim + 1 // incy ); - } else + } + else { - ModuleBase::gemm_op()('N', - 'N', - dim, // m: row of A,C - notconv, // n: col of B,C - nbase, // k: col of A, row of B - this->one, // alpha - spsi, // A - matrix_ld, // LDA: if(N) max(1,m) if(T) max(1,k) - vc_ev_vector, // B - nbase, // LDB: if(N) max(1,k) if(T) max(1,n) - this->one, // beta - basis + dim * nbase, // C dim * notconv - matrix_ld // LDC: if(N) max(1, m) + ModuleBase::gemm_op()( + 'N', + 'N', + dim, // m: row of A,C + notconv, // n: col of B,C + nbase, // k: col of A, row of B + this->one, // alpha + spsi, // A + matrix_ld, // LDA: if(N) max(1,m) if(T) max(1,k) + vc_ev_vector, // B + nbase, // LDB: if(N) max(1,k) if(T) max(1,n) + this->one, // beta + basis + dim * nbase, // C dim * notconv + matrix_ld // LDC: if(N) max(1, m) ); } //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< @@ -480,18 +532,20 @@ void DiagoDavid::cal_grad(const HSOperator& op, if (this->device == base_device::GpuDevice) { #if defined(__CUDA) || defined(__ROCM) - ModuleBase::vector_div_vector_op()(dim, - basis + dim * (nbase + m), - basis + dim * (nbase + m), - this->d_precondition); + ModuleBase::vector_div_vector_op()( + dim, + basis + dim * (nbase + m), + basis + dim * (nbase + m), + this->d_precondition); #endif } else { - ModuleBase::vector_div_vector_op()(dim, - basis + dim * (nbase + m), - basis + dim * (nbase + m), - this->precondition); + ModuleBase::vector_div_vector_op()( + dim, + basis + dim * (nbase + m), + basis + dim * (nbase + m), + this->precondition); } //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // for (int ig = 0; ig < dim; ig++) @@ -503,7 +557,7 @@ void DiagoDavid::cal_grad(const HSOperator& op, // there is a nbase to nbase + notconv band orthogonalise // plan for SchmidtOrth - T* lagrange = nullptr; + T *lagrange = nullptr; resmem_complex_op()(lagrange, notconv * (nbase + notconv)); setmem_complex_op()(lagrange, 0, notconv * (nbase + notconv)); @@ -514,56 +568,68 @@ void DiagoDavid::cal_grad(const HSOperator& op, { { // phm_in->sPsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1); - op.spsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1); + op.spsi( + basis + dim * (nbase + m), + &spsi[(nbase + m) * dim], + dim, + 1); } } // first nbase bands psi* dot notconv bands spsi to prepare lagrange_matrix // calculate the square matrix for future lagranges - if (notconv == 1){ - //Use gemv for vector case to avoid potential bug using gemm call with n=1 - ModuleBase::gemv_op()('C', - dim, // m: row of A - nbase, // n: col of A - this->one, // alpha - basis, // A dim * nbase - matrix_ld, // LDA: if(N) max(1,m) - &spsi[nbase * dim], // X dim - 1, // incx - this->zero, // beta - lagrange, // Y nbase - 1 - ); - } else + if (notconv == 1) { - ModuleBase::gemm_op()('C', - 'N', - nbase, // m: row of A,C - notconv, // n: col of B,C - dim, // k: col of A, row of B - this->one, // alpha - basis, // A - matrix_ld, // LDA: if(N) max(1,m) if(T) max(1,k) - &spsi[nbase * dim], // B - matrix_ld, // LDB: if(N) max(1,k) if(T) max(1,n) - this->zero, // belta - lagrange, // C - nbase + notconv // LDC: if(N) max(1, m) + // Use gemv for vector case to avoid potential bug using gemm call with n=1 + ModuleBase::gemv_op()( + 'C', + dim, // m: row of A + nbase, // n: col of A + this->one, // alpha + basis, // A dim * nbase + matrix_ld, // LDA: if(N) max(1,m) + &spsi[nbase * dim], // X dim + 1, // incx + this->zero, // beta + lagrange, // Y nbase + 1); + } + else + { + ModuleBase::gemm_op()( + 'C', + 'N', + nbase, // m: row of A,C + notconv, // n: col of B,C + dim, // k: col of A, row of B + this->one, // alpha + basis, // A + matrix_ld, // LDA: if(N) max(1,m) if(T) max(1,k) + &spsi[nbase * dim], // B + matrix_ld, // LDB: if(N) max(1,k) if(T) max(1,n) + this->zero, // belta + lagrange, // C + nbase + notconv // LDC: if(N) max(1, m) ); } for (int m = 0; m < notconv; m++) { - this->SchmidtOrth(dim, - nbase + notconv, - nbase + m, - spsi, - &lagrange[m * (nbase + notconv)], - pre_matrix_mm_m[m], - pre_matrix_mv_m[m]); + this->SchmidtOrth( + dim, + nbase + notconv, + nbase + m, + spsi, + &lagrange[m * (nbase + notconv)], + pre_matrix_mm_m[m], + pre_matrix_mv_m[m]); { // phm_in->sPsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, dim, 1); - op.spsi(basis + dim*(nbase + m), &spsi[(nbase + m) * dim], dim, 1); + op.spsi( + basis + dim * (nbase + m), + &spsi[(nbase + m) * dim], + dim, + 1); } } // calculate H|psi> for not convergence bands @@ -572,7 +638,11 @@ void DiagoDavid::cal_grad(const HSOperator& op, // &hpsi[nbase * dim]); // &hp(nbase, 0) // phm_in->ops->hPsi(dav_hpsi_in); // hpsi[:, nbase:nbase+notcnv] = H basis[:, nbase:nbase+notcnv] - op.hpsi(basis + nbase * dim, hpsi + nbase * dim, dim, notconv); + op.hpsi( + basis + nbase * dim, + hpsi + nbase * dim, + dim, + notconv); delmem_complex_op()(lagrange); delmem_complex_op()(vc_ev_vector); @@ -581,51 +651,61 @@ void DiagoDavid::cal_grad(const HSOperator& op, return; } - template -void DiagoDavid::cal_elem(const int& dim, - int& nbase, // current dimension of the reduced basis - const int nbase_x, // maximum dimension of the reduced basis set - const int& notconv, // number of newly added basis vectors - const T* hpsi, - const T* spsi, - T* hcc) +void DiagoDavid::cal_elem( + const int &dim, + int &nbase, // current dimension of the reduced basis + const int nbase_x, // maximum dimension of the reduced basis set + const int ¬conv, // number of newly added basis vectors + const T *hpsi, + const T *spsi, + T *hcc) { - if (test_david == 1) { + if (test_david == 1) + { ModuleBase::TITLE("DiagoDavid", "cal_elem"); } - if (notconv == 0) { + if (notconv == 0) + { return; } ModuleBase::timer::start("DiagoDavid", "cal_elem"); const int matrix_ld = std::max(dim, 1); // hcc[nbase](notconv, nbase + notconv)= basis[nbase]' * hpsi - ModuleBase::gemm_op()('C', - 'N', - notconv, - nbase + notconv, - dim, - this->one, - basis + dim * nbase, // basis(:,nbase:) dim * notconv - matrix_ld, - hpsi, // dim * (nbase + notconv) - matrix_ld, - this->zero, - hcc + nbase, // notconv * (nbase + notconv) - nbase_x); - + ModuleBase::gemm_op()( + 'C', + 'N', + notconv, + nbase + notconv, + dim, + this->one, + basis + dim * nbase, // basis(:,nbase:) dim * notconv + matrix_ld, + hpsi, // dim * (nbase + notconv) + matrix_ld, + this->zero, + hcc + nbase, // notconv * (nbase + notconv) + nbase_x); #ifdef __MPI if (diag_comm.nproc > 1) { - ModuleBase::matrixTranspose_op()(nbase_x, nbase_x, hcc, hcc); + ModuleBase::matrixTranspose_op()( + nbase_x, + nbase_x, + hcc, + hcc); assert(diag_comm.comm == POOL_WORLD); Parallel_Reduce::reduce_pool(hcc + nbase * nbase_x, notconv * nbase_x); - ModuleBase::matrixTranspose_op()(nbase_x, nbase_x, hcc, hcc); + ModuleBase::matrixTranspose_op()( + nbase_x, + nbase_x, + hcc, + hcc); } #endif @@ -646,12 +726,13 @@ void DiagoDavid::cal_elem(const int& dim, // we needn't creat another two complexmatrixs in diag_zhegvx(). //============================================================================== template -void DiagoDavid::diag_zhegvx(const int& nbase, - const int& nband, - const T* hcc, - const int& nbase_x, - Real* eigenvalue, // in CPU - T* vcc) +void DiagoDavid::diag_zhegvx( + const int &nbase, + const int &nband, + const T *hcc, + const int &nbase_x, + Real *eigenvalue, // in CPU + T *vcc) { ModuleBase::timer::start("DiagoDavid", "diag_zhegvx"); if (diag_comm.rank == 0) @@ -661,12 +742,18 @@ void DiagoDavid::diag_zhegvx(const int& nbase, if (this->device == base_device::GpuDevice) { #if defined(__CUDA) || defined(__ROCM) - Real* eigenvalue_gpu = nullptr; + Real *eigenvalue_gpu = nullptr; resmem_var_op()(eigenvalue_gpu, nbase_x); syncmem_var_h2d_op()(eigenvalue_gpu, this->eigenvalue, nbase_x); // heevx_op()(this->ctx, nbase, nbase_x, hcc, nband, eigenvalue_gpu, vcc); - ct::kernels::lapack_heevx()(nbase, nbase_x, hcc, nband, eigenvalue_gpu, vcc); + ct::kernels::lapack_heevx()( + nbase, + nbase_x, + hcc, + nband, + eigenvalue_gpu, + vcc); syncmem_var_d2h_op()(this->eigenvalue, eigenvalue_gpu, nbase_x); delmem_var_op()(eigenvalue_gpu); @@ -674,8 +761,14 @@ void DiagoDavid::diag_zhegvx(const int& nbase, } else { - //heevx_op()(this->ctx, nbase, nbase_x, hcc, nband, this->eigenvalue, vcc); - ct::kernels::lapack_heevx()(nbase, nbase_x, hcc, nband, this->eigenvalue, vcc); + // heevx_op()(this->ctx, nbase, nbase_x, hcc, nband, this->eigenvalue, vcc); + ct::kernels::lapack_heevx()( + nbase, + nbase_x, + hcc, + nband, + this->eigenvalue, + vcc); } } @@ -685,9 +778,19 @@ void DiagoDavid::diag_zhegvx(const int& nbase, // vcc: nbase * nband for (int i = 0; i < nband; i++) { - MPI_Bcast(&vcc[i * nbase_x], nbase, MPI_DOUBLE_COMPLEX, 0, diag_comm.comm); + MPI_Bcast( + &vcc[i * nbase_x], + nbase, + MPI_DOUBLE_COMPLEX, + 0, + diag_comm.comm); } - MPI_Bcast(this->eigenvalue, nband, MPI_DOUBLE, 0, diag_comm.comm); + MPI_Bcast( + this->eigenvalue, + nband, + MPI_DOUBLE, + 0, + diag_comm.comm); } #endif @@ -695,63 +798,66 @@ void DiagoDavid::diag_zhegvx(const int& nbase, return; } - template -void DiagoDavid::refresh(const int& dim, - const int& nband, - int& nbase, - const int nbase_x, // maximum dimension of the reduced basis set - const Real* eigenvalue_in, - const T *psi_in, - const int ld_psi, - T* hpsi, - T* spsi, - T* hcc, - T* vcc) +void DiagoDavid::refresh( + const int &dim, + const int &nband, + int &nbase, + const int nbase_x, // maximum dimension of the reduced basis set + const Real *eigenvalue_in, + const T *psi_in, + const int ld_psi, + T *hpsi, + T *spsi, + T *hcc, + T *vcc) { - if (test_david == 1) { + if (test_david == 1) + { ModuleBase::TITLE("DiagoDavid", "refresh"); } ModuleBase::timer::start("DiagoDavid", "refresh"); const int matrix_ld = std::max(dim, 1); // update hp,sp - setmem_complex_op()(basis , 0, nbase_x * dim); + setmem_complex_op()(basis, 0, nbase_x * dim); // basis(dim, nband) = hpsi(dim, nbase) * vcc(nbase, nband) - ModuleBase::gemm_op()('N', - 'N', - dim, // m: row of A,C - nband, // n: col of B,C - nbase, // k: col of A, row of B - this->one, - hpsi, // A dim * nbase - matrix_ld, - vcc, // B nbase * nband - nbase_x, - zero, - basis, // C dim * nband - matrix_ld); + ModuleBase::gemm_op()( + 'N', + 'N', + dim, // m: row of A,C + nband, // n: col of B,C + nbase, // k: col of A, row of B + this->one, + hpsi, // A dim * nbase + matrix_ld, + vcc, // B nbase * nband + nbase_x, + zero, + basis, // C dim * nband + matrix_ld); //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // basis[nband] = spsi * vcc - ModuleBase::gemm_op()('N', - 'N', - dim, // m: row of A,C - nband, // n: col of B,C - nbase, // k: col of A, row of B - this->one, - spsi, // A dim * nbase - matrix_ld, - vcc, // B nbase * nband - nbase_x, - this->zero, - basis + dim * nband, // C dim * nband - matrix_ld); + ModuleBase::gemm_op()( + 'N', + 'N', + dim, // m: row of A,C + nband, // n: col of B,C + nbase, // k: col of A, row of B + this->one, + spsi, // A dim * nbase + matrix_ld, + vcc, // B nbase * nband + nbase_x, + this->zero, + basis + dim * nband, // C dim * nband + matrix_ld); // hpsi = basis, spsi = basis[nband] syncmem_complex_op()(hpsi, basis, dim * nband); - syncmem_complex_op()(spsi, basis + dim*nband, dim * nband); + syncmem_complex_op()(spsi, basis + dim * nband, dim * nband); /*for (int m = 0; m < nband; m++) { for (int ig = 0; ig < dim; ig++) { @@ -761,11 +867,11 @@ void DiagoDavid::refresh(const int& dim, }*/ // update basis - setmem_complex_op()(basis , 0, nbase_x * dim); + setmem_complex_op()(basis, 0, nbase_x * dim); for (int m = 0; m < nband; m++) { - syncmem_complex_op()(basis + dim*m,psi_in + m*ld_psi, dim); + syncmem_complex_op()(basis + dim * m, psi_in + m * ld_psi, dim); /*for (int ig = 0; ig < npw; ig++) basis(m, ig) = psi(m, ig);*/ } @@ -779,14 +885,10 @@ void DiagoDavid::refresh(const int& dim, if (this->device == base_device::GpuDevice) { #if defined(__CUDA) || defined(__ROCM) - T* hcc_cpu = nullptr; - T* vcc_cpu = nullptr; - base_device::memory::resize_memory_op()(hcc_cpu, - nbase_x * nbase_x, - "DAV::hcc"); - base_device::memory::resize_memory_op()(vcc_cpu, - nbase_x * nbase_x, - "DAV::vcc"); + T *hcc_cpu = nullptr; + T *vcc_cpu = nullptr; + base_device::memory::resize_memory_op()(hcc_cpu, nbase_x * nbase_x, "DAV::hcc"); + base_device::memory::resize_memory_op()(vcc_cpu, nbase_x * nbase_x, "DAV::vcc"); syncmem_d2h_op()(hcc_cpu, hcc, nbase_x * nbase_x); syncmem_d2h_op()(vcc_cpu, vcc, nbase_x * nbase_x); @@ -818,15 +920,15 @@ void DiagoDavid::refresh(const int& dim, return; } - template -void DiagoDavid::SchmidtOrth(const int& dim, - const int nband, - const int m, - const T* spsi, - T* lagrange_m, - const int mm_size, - const int mv_size) +void DiagoDavid::SchmidtOrth( + const int &dim, + const int nband, + const int m, + const T *spsi, + T *lagrange_m, + const int mm_size, + const int mv_size) { // if(test_david == 1) ModuleBase::TITLE("DiagoDavid","SchmidtOrth"); ModuleBase::timer::start("DiagoDavid", "SchmidtOrth"); @@ -850,7 +952,7 @@ void DiagoDavid::SchmidtOrth(const int& dim, } // psi_m = basis[m] - T* psi_m = basis + dim*m; + T *psi_m = basis + dim * m; // std::complex *lagrange = new std::complex[m + 1]; // ModuleBase::GlobalFunc::ZEROS(lagrange, m + 1); @@ -860,35 +962,37 @@ void DiagoDavid::SchmidtOrth(const int& dim, { // lagrange_m[m - mv_size + 1 - mm_size] // = basis[m - mv_size + 1 - mm_size]' * spsi[m] - ModuleBase::gemm_op()('C', - 'N', - mm_size, // m: row of A,C - mm_size, // n: col of B,C - dim, // k: col of A, row of B - this->one, // alpha - basis + dim * (m - mv_size + 1 - mm_size), // A - matrix_ld, // LDA: if(N) max(1,m) if(T) max(1,k) - &spsi[m * dim], // B - matrix_ld, // LDB: if(N) max(1,k) if(T) max(1,n) - this->zero, // belta - &lagrange_m[m - mv_size + 1 - mm_size], // C - nband // LDC: if(N) max(1, m) + ModuleBase::gemm_op()( + 'C', + 'N', + mm_size, // m: row of A,C + mm_size, // n: col of B,C + dim, // k: col of A, row of B + this->one, // alpha + basis + dim * (m - mv_size + 1 - mm_size), // A + matrix_ld, // LDA: if(N) max(1,m) if(T) max(1,k) + &spsi[m * dim], // B + matrix_ld, // LDB: if(N) max(1,k) if(T) max(1,n) + this->zero, // belta + &lagrange_m[m - mv_size + 1 - mm_size], // C + nband // LDC: if(N) max(1, m) ); } // calculate other lagranges for this band // lagrange_m[m - mv_size + 1] // = basis[m - mv_size + 1]' * spsi[m] - ModuleBase::gemv_op()('C', - dim, - mv_size, - this->one, - basis + dim * (m - mv_size + 1), - matrix_ld, - &spsi[m * dim], - 1, - this->zero, - &lagrange_m[m - mv_size + 1], - 1); + ModuleBase::gemv_op()( + 'C', + dim, + mv_size, + this->one, + basis + dim * (m - mv_size + 1), + matrix_ld, + &spsi[m * dim], + 1, + this->zero, + &lagrange_m[m - mv_size + 1], + 1); Parallel_Reduce::reduce_pool(lagrange_m, m + 1); @@ -900,20 +1004,25 @@ void DiagoDavid::SchmidtOrth(const int& dim, // / psi_m = psi_m - \sum_{i < m} \langle psi(i)|S|psi(m) \rangle psi(i) // psi_m = psi_m - basis * lagrange_m - ModuleBase::gemv_op()('N', - dim, - m, - this->neg_one, - basis, - matrix_ld, - lagrange_m, - 1, - this->one, - psi_m, - 1); + ModuleBase::gemv_op()( + 'N', + dim, + m, + this->neg_one, + basis, + matrix_ld, + lagrange_m, + 1, + this->one, + psi_m, + 1); // psi_norm = psi_norm - lagrange_m · lagrange_m - psi_norm -= ModuleBase::dot_real_op()(m, lagrange_m, lagrange_m, false); + psi_norm -= ModuleBase::dot_real_op()( + m, + lagrange_m, + lagrange_m, + false); // for (int j = 0; j < m; j++) // { @@ -943,7 +1052,11 @@ void DiagoDavid::SchmidtOrth(const int& dim, else { // psi_m = psi_m / psi_norm - ModuleBase::vector_mul_real_op()(dim, psi_m, psi_m, Real(1.0 / psi_norm)); + ModuleBase::vector_mul_real_op()( + dim, + psi_m, + psi_m, + Real(1.0 / psi_norm)); // for (int i = 0; i < npw; i++) // { // psi_m[i] /= psi_norm; @@ -955,11 +1068,12 @@ void DiagoDavid::SchmidtOrth(const int& dim, return; } - template -void DiagoDavid::planSchmidtOrth(const int nband, std::vector& pre_matrix_mm_m, std::vector& pre_matrix_mv_m) +void DiagoDavid::planSchmidtOrth( + const int nband, std::vector &pre_matrix_mm_m, std::vector &pre_matrix_mv_m) { - if (nband <= 0) { + if (nband <= 0) + { return; } std::fill(pre_matrix_mm_m.begin(), pre_matrix_mm_m.end(), 0); @@ -976,11 +1090,14 @@ void DiagoDavid::planSchmidtOrth(const int nband, std::vector& p { divide_points[0] = index; pre_matrix_mm_m[index] = matrix_size; - if (res_nband == matrix_size) { + if (res_nband == matrix_size) + { pre_matrix_mv_m[index] = 1; - } else { + } + else + { pre_matrix_mv_m[index] = 2; -} + } divide_times = 1; } else @@ -1019,16 +1136,16 @@ void DiagoDavid::planSchmidtOrth(const int nband, std::vector& p } } - template -int DiagoDavid::diag(const HSOperator& op, - const int ld_psi, - T *psi_in, - Real* eigenvalue_in, - const std::vector& ethr_band, - const int david_maxiter, - const int ntry_max, - const int notconv_max) +int DiagoDavid::diag( + const HSOperator &op, + const int ld_psi, + T *psi_in, + Real *eigenvalue_in, + const std::vector ðr_band, + const int david_maxiter, + const int ntry_max, + const int notconv_max) { /// record the times of trying iterative diagonalization int ntry = 0; @@ -1037,9 +1154,21 @@ int DiagoDavid::diag(const HSOperator& op, int sum_dav_iter = 0; do { - sum_dav_iter += this->diag_once(op, dim, nband, ld_psi, psi_in, eigenvalue_in, ethr_band, david_maxiter); + sum_dav_iter += this->diag_once( + op, + dim, + nband, + ld_psi, + psi_in, + eigenvalue_in, + ethr_band, + david_maxiter); ++ntry; - } while (!check_block_conv(ntry, this->notconv, ntry_max, notconv_max)); + } while (!check_block_conv( + ntry, + this->notconv, + ntry_max, + notconv_max)); if (notconv > std::max(5, nband / 4)) { @@ -1049,20 +1178,20 @@ int DiagoDavid::diag(const HSOperator& op, return sum_dav_iter; } - template -inline bool DiagoDavid::check_block_conv(const int& ntry, - const int& notconv, - const int& ntry_max, - const int& notconv_max) +inline bool DiagoDavid::check_block_conv( + const int &ntry, + const int ¬conv, + const int &ntry_max, + const int ¬conv_max) { // Allow at most 5 tries at diag. If more than 5 then exit loop. - if(ntry > ntry_max) + if (ntry > ntry_max) { return true; } // If notconv <= notconv_max allowed, set convergence to true and exit loop. - if(notconv <= notconv_max) + if (notconv <= notconv_max) { return true; } @@ -1070,7 +1199,8 @@ inline bool DiagoDavid::check_block_conv(const int& ntry, return false; } -namespace hsolver { +namespace hsolver +{ template class DiagoDavid, base_device::DEVICE_CPU>; template class DiagoDavid, base_device::DEVICE_CPU>; #if ((defined __CUDA) || (defined __ROCM)) diff --git a/source/source_hsolver/diago_david.h b/source/source_hsolver/diago_david.h index 5dd29805378..8557f14903e 100644 --- a/source/source_hsolver/diago_david.h +++ b/source/source_hsolver/diago_david.h @@ -1,12 +1,10 @@ #ifndef DIAGODAVID_H #define DIAGODAVID_H -#include "source_base/macros.h" // GetRealType -#include "source_base/module_device/device.h" // base_device -#include "source_base/module_device/memory_op.h"// base_device::memory - +#include "source_base/macros.h" // GetRealType #include "source_base/module_container/ATen/kernels/lapack.h" // container::kernels - +#include "source_base/module_device/device.h" // base_device +#include "source_base/module_device/memory_op.h" // base_device::memory #include "source_hsolver/diag_comm_info.h" #include "source_hsolver/hs_operator.h" #include "source_hsolver/kernels/hegvd_op.h" @@ -25,8 +23,7 @@ namespace hsolver * @tparam T The data type of the matrices and arrays (e.g., float, double, std::complex, std::complex). * @tparam Device The device type (e.g., base_device::DEVICE_CPU or DEVICE_GPU). */ -template , typename Device = base_device::DEVICE_CPU> -class DiagoDavid +template , typename Device = base_device::DEVICE_CPU> class DiagoDavid { private: // Note GetTypeReal::type will @@ -35,7 +32,6 @@ class DiagoDavid using Real = typename GetTypeReal::type; public: - /** * @brief Constructor for the DiagoDavid class. * @@ -52,11 +48,12 @@ class DiagoDavid * * @note Auxiliary memory is allocated in the constructor and deallocated in the destructor. */ - DiagoDavid(const Real* precondition_in, - const int nband_in, - const int dim_in, - const int david_ndim_in, - const diag_comm_info& diag_comm_in); + DiagoDavid( + const Real *precondition_in, + const int nband_in, + const int dim_in, + const int david_ndim_in, + const diag_comm_info &diag_comm_in); /** * @brief Destructor for the DiagoDavid class. @@ -67,7 +64,6 @@ class DiagoDavid */ ~DiagoDavid(); - /** * @brief Performs iterative diagonalization using the David algorithm. * @@ -87,14 +83,14 @@ class DiagoDavid * notconv_max is determined by the accuracy required for the calculation, default 0 */ int diag( - const HSOperator& op, // applies H and S to block vectors - const int ld_psi, // Leading dimension of the psi input - T *psi_in, // Pointer to eigenvectors - Real* eigenvalue_in, // Pointer to store the resulting eigenvalues - const std::vector& ethr_band, // Convergence threshold for the Davidson iteration - const int david_maxiter, // Maximum allowed iterations for the Davidson method - const int ntry_max = 5, // Maximum number of diagonalization attempts (5 by default) - const int notconv_max = 0); // Maximum number of allowed non-converged eigenvectors + const HSOperator &op, // applies H and S to block vectors + const int ld_psi, // Leading dimension of the psi input + T *psi_in, // Pointer to eigenvectors + Real *eigenvalue_in, // Pointer to store the resulting eigenvalues + const std::vector ðr_band, // Convergence threshold for the Davidson iteration + const int david_maxiter, // Maximum allowed iterations for the Davidson method + const int ntry_max = 5, // Maximum number of diagonalization attempts (5 by default) + const int notconv_max = 0); // Maximum number of allowed non-converged eigenvectors private: int test_david = 0; @@ -113,37 +109,38 @@ class DiagoDavid int notconv = 0; /// precondition for diag, diagonal approximation of matrix A(i.e. Hamilt) - const Real* precondition = nullptr; - Real* d_precondition = nullptr; + const Real *precondition = nullptr; + Real *d_precondition = nullptr; /// eigenvalue results - Real* eigenvalue = nullptr; + Real *eigenvalue = nullptr; - T *basis = nullptr; /// pointer to basis set(dim, nbase_x), leading dimension = dim + T *basis = nullptr; /// pointer to basis set(dim, nbase_x), leading dimension = dim - T* hpsi = nullptr; /// the product of H and psi in the reduced basis set + T *hpsi = nullptr; /// the product of H and psi in the reduced basis set - T* spsi = nullptr; /// the Product of S and psi in the reduced basis set + T *spsi = nullptr; /// the Product of S and psi in the reduced basis set - T* hcc = nullptr; /// Hamiltonian on the reduced basis + T *hcc = nullptr; /// Hamiltonian on the reduced basis - T* vcc = nullptr; /// eigenvectors of hc + T *vcc = nullptr; /// eigenvectors of hc - T* lagrange_matrix = nullptr; + T *lagrange_matrix = nullptr; /// device type of psi - Device* ctx = {}; - base_device::DEVICE_CPU* cpu_ctx = {}; + Device *ctx = {}; + base_device::DEVICE_CPU *cpu_ctx = {}; base_device::AbacusDevice_t device = {}; - int diag_once(const HSOperator& op, - const int dim, - const int nband, - const int ld_psi, - T *psi_in, - Real* eigenvalue_in, - const std::vector& ethr_band, - const int david_maxiter); + int diag_once( + const HSOperator &op, + const int dim, + const int nband, + const int ld_psi, + T *psi_in, + Real *eigenvalue_in, + const std::vector ðr_band, + const int david_maxiter); /** * Calculates the preconditioned gradient of the eigenvectors in Davidson method. @@ -159,16 +156,17 @@ class DiagoDavid * @param unconv The array of indices for the unconverged eigenpairs. * @param eigenvalue The array of eigenvalues. */ - void cal_grad(const HSOperator& op, - const int& dim, - const int& nbase, - const int nbase_x, - const int& notconv, - T* hpsi, - T* spsi, - const T* vcc, - const int* unconv, - const Real* eigenvalue); + void cal_grad( + const HSOperator &op, + const int &dim, + const int &nbase, + const int nbase_x, + const int ¬conv, + T *hpsi, + T *spsi, + const T *vcc, + const int *unconv, + const Real *eigenvalue); /** * Calculates the elements of the diagonalization matrix for the DiagoDavid class. @@ -181,13 +179,14 @@ class DiagoDavid * @param spsi The output array for the overlap matrix S times blockvector psi. * @param hcc Pointer to the array where the calculated Hamiltonian matrix elements will be stored. */ - void cal_elem(const int& dim, - int& nbase, - const int nbase_x, - const int& notconv, - const T* hpsi, - const T* spsi, - T* hcc); + void cal_elem( + const int &dim, + int &nbase, + const int nbase_x, + const int ¬conv, + const T *hpsi, + const T *spsi, + T *hcc); /** * Refreshes the diagonalization solver by updating the basis and the reduced Hamiltonian. @@ -205,17 +204,18 @@ class DiagoDavid * @param vcc Pointer to the output array for the updated eigenvector matrix. * */ - void refresh(const int& dim, - const int& nband, - int& nbase, - const int nbase_x, - const Real* eigenvalue, - const T *psi_in, - const int ld_psi, - T* hpsi, - T* spsi, - T* hcc, - T* vcc); + void refresh( + const int &dim, + const int &nband, + int &nbase, + const int nbase_x, + const Real *eigenvalue, + const T *psi_in, + const int ld_psi, + T *hpsi, + T *spsi, + T *hcc, + T *vcc); /** * SchmidtOrth function performs orthogonalization of the starting eigenfunction to those already calculated. @@ -230,13 +230,14 @@ class DiagoDavid * @param mm_size The size of the square matrix for future lagranges. * @param mv_size The size of the lagrange_m array. */ - void SchmidtOrth(const int& dim, - const int nband, - const int m, - const T* spsi, - T* lagrange_m, - const int mm_size, - const int mv_size); + void SchmidtOrth( + const int &dim, + const int nband, + const int m, + const T *spsi, + T *lagrange_m, + const int mm_size, + const int mv_size); /** * @brief Plans the Schmidt orthogonalization for a given number of bands. @@ -247,14 +248,15 @@ class DiagoDavid * @param pre_matrix_mm_m The vector to store the matrix sizes. * @param pre_matrix_mv_m The vector to store the number of matrix-vector multiplications. */ - void planSchmidtOrth(const int nband, std::vector& pre_matrix_mm_m, std::vector& pre_matrix_mv_m); + void planSchmidtOrth(const int nband, std::vector &pre_matrix_mm_m, std::vector &pre_matrix_mv_m); - void diag_zhegvx(const int& nbase, - const int& nband, - const T* hcc, - const int& nbase_x, - Real* eigenvalue, - T* vcc); + void diag_zhegvx( + const int &nbase, + const int &nband, + const T *hcc, + const int &nbase_x, + Real *eigenvalue, + T *vcc); /** * @brief Check the convergence of block eigenvectors in the Davidson iteration. @@ -276,7 +278,11 @@ class DiagoDavid * @note Exits the diagonalization loop if either the convergence criteria * are met or the maximum number of tries is exceeded. */ - bool check_block_conv(const int &ntry, const int ¬conv, const int &ntry_max, const int ¬conv_max); + bool check_block_conv( + const int &ntry, + const int ¬conv, + const int &ntry_max, + const int ¬conv_max); using resmem_complex_op = base_device::memory::resize_memory_op; using delmem_complex_op = base_device::memory::delete_memory_op; diff --git a/source/source_hsolver/diago_elpa.cpp b/source/source_hsolver/diago_elpa.cpp index eceb85da86c..a4e163929ba 100644 --- a/source/source_hsolver/diago_elpa.cpp +++ b/source/source_hsolver/diago_elpa.cpp @@ -4,29 +4,41 @@ #include "source_base/module_external/blacs_connector.h" #include "source_base/module_external/blas_connector.h" #include "source_base/timer.h" -#include "source_base/tool_title.h" #include "source_base/tool_quit.h" +#include "source_base/tool_title.h" -namespace hsolver { +namespace hsolver +{ #ifdef __MPI -template <> -MPI_Comm DiagoElpa::setmpicomm() { - if (this->elpa_num_thread == -1) { +template <> MPI_Comm DiagoElpa::setmpicomm() +{ + if (this->elpa_num_thread == -1) + { return MPI_COMM_WORLD; - } else { + } + else + { int _num = 0; MPI_Comm_size(MPI_COMM_WORLD, &_num); - if (elpa_num_thread > _num || elpa_num_thread <= 0) { + if (elpa_num_thread > _num || elpa_num_thread <= 0) + { return MPI_COMM_WORLD; - } else { + } + else + { lastmpinum++; - int* _ranks = new int[elpa_num_thread]; - for (int i = 0; i < elpa_num_thread; i++) { + int *_ranks = new int[elpa_num_thread]; + for (int i = 0; i < elpa_num_thread; i++) + { _ranks[i] = (lastmpinum + i) % _num; } MPI_Group _tempgroup, _oldgroup; MPI_Comm_group(MPI_COMM_WORLD, &_oldgroup); - MPI_Group_incl(_oldgroup, elpa_num_thread, _ranks, &_tempgroup); + MPI_Group_incl( + _oldgroup, + elpa_num_thread, + _ranks, + &_tempgroup); MPI_Comm _new_comm; MPI_Comm_create(MPI_COMM_WORLD, _tempgroup, &_new_comm); delete[] _ranks; @@ -34,24 +46,35 @@ MPI_Comm DiagoElpa::setmpicomm() { } } } -template <> -MPI_Comm DiagoElpa>::setmpicomm() { - if (this->elpa_num_thread == -1) { +template <> MPI_Comm DiagoElpa>::setmpicomm() +{ + if (this->elpa_num_thread == -1) + { return MPI_COMM_WORLD; - } else { + } + else + { int _num = 0; MPI_Comm_size(MPI_COMM_WORLD, &_num); - if (elpa_num_thread > _num || elpa_num_thread <= 0) { + if (elpa_num_thread > _num || elpa_num_thread <= 0) + { return MPI_COMM_WORLD; - } else { + } + else + { lastmpinum++; - int* _ranks = new int[elpa_num_thread]; - for (int i = 0; i < elpa_num_thread; i++) { + int *_ranks = new int[elpa_num_thread]; + for (int i = 0; i < elpa_num_thread; i++) + { _ranks[i] = (lastmpinum + i) % _num; } MPI_Group _tempgroup, _oldgroup; MPI_Comm_group(MPI_COMM_WORLD, &_oldgroup); - MPI_Group_incl(_oldgroup, elpa_num_thread, _ranks, &_tempgroup); + MPI_Group_incl( + _oldgroup, + elpa_num_thread, + _ranks, + &_tempgroup); MPI_Comm _new_comm; MPI_Comm_create(MPI_COMM_WORLD, _tempgroup, &_new_comm); delete[] _ranks; @@ -62,146 +85,158 @@ MPI_Comm DiagoElpa>::setmpicomm() { #endif template <> void DiagoElpa>::diag( - ModuleBase::MatrixBlock>& h_mat, - ModuleBase::MatrixBlock>& s_mat, - psi::Psi>& psi, - Real* eigenvalue_in) { + ModuleBase::MatrixBlock> &h_mat, + ModuleBase::MatrixBlock> &s_mat, + psi::Psi> &psi, + Real *eigenvalue_in) +{ ModuleBase::TITLE("DiagoElpa", "diag"); #ifdef __MPI std::vector eigen(this->nlocal, 0.0); bool isReal = false; MPI_Comm COMM_DIAG = setmpicomm(); // set mpi_comm needed - ELPA_Solver es((const bool)isReal, - COMM_DIAG, - (const int)this->nbands, - (const int)h_mat.row, - (const int)h_mat.col, - (const int*)h_mat.desc); - this->DecomposedState - = 0; // for k pointer, the decomposed s_mat can not be reused + ELPA_Solver es( + (const bool)isReal, + COMM_DIAG, (const int)this->nbands, (const int)h_mat.row, (const int)h_mat.col, (const int *)h_mat.desc); + this->DecomposedState = 0; // for k pointer, the decomposed s_mat can not be reused ModuleBase::timer::start("DiagoElpa", "elpa_solve"); - es.generalized_eigenvector(h_mat.p, - s_mat.p, - this->DecomposedState, - eigen.data(), - psi.get_pointer()); + es.generalized_eigenvector( + h_mat.p, + s_mat.p, + this->DecomposedState, + eigen.data(), + psi.get_pointer()); ModuleBase::timer::end("DiagoElpa", "elpa_solve"); es.exit(); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); #else - ModuleBase::WARNING_QUIT("DiagoElpa", - "DiagoElpa only can be used with macro __MPI"); + ModuleBase::WARNING_QUIT("DiagoElpa", "DiagoElpa only can be used with macro __MPI"); #endif } template <> -void DiagoElpa::diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in) { +void DiagoElpa::diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in) +{ ModuleBase::TITLE("DiagoElpa", "diag"); #ifdef __MPI std::vector eigen(this->nlocal, 0.0); bool isReal = true; MPI_Comm COMM_DIAG = setmpicomm(); // set mpi_comm needed - ELPA_Solver es((const bool)isReal, - COMM_DIAG, - (const int)this->nbands, - (const int)h_mat.row, - (const int)h_mat.col, - (const int*)h_mat.desc); + ELPA_Solver es( + (const bool)isReal, + COMM_DIAG, (const int)this->nbands, (const int)h_mat.row, (const int)h_mat.col, (const int *)h_mat.desc); ModuleBase::timer::start("DiagoElpa", "elpa_solve"); - es.generalized_eigenvector(h_mat.p, - s_mat.p, - this->DecomposedState, - eigen.data(), - psi.get_pointer()); + es.generalized_eigenvector( + h_mat.p, + s_mat.p, + this->DecomposedState, + eigen.data(), + psi.get_pointer()); ModuleBase::timer::end("DiagoElpa", "elpa_solve"); es.exit(); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); #else - ModuleBase::WARNING_QUIT("DiagoElpa", - "DiagoElpa only can be used with macro __MPI"); + ModuleBase::WARNING_QUIT("DiagoElpa", "DiagoElpa only can be used with macro __MPI"); #endif } - #ifdef __MPI template <> -void DiagoElpa>::diag_pool(ModuleBase::MatrixBlock>& h_mat, - ModuleBase::MatrixBlock>& s_mat, - psi::Psi>& psi, - Real* eigenvalue_in, - MPI_Comm& comm) +void DiagoElpa>::diag_pool( + ModuleBase::MatrixBlock> &h_mat, + ModuleBase::MatrixBlock> &s_mat, + psi::Psi> &psi, + Real *eigenvalue_in, + MPI_Comm &comm) { std::vector eigen(this->nlocal, 0.0); bool isReal = false; - ELPA_Solver es((const bool)isReal, - comm, - (const int)this->nbands, - (const int)h_mat.row, - (const int)h_mat.col, - (const int*)h_mat.desc); - this->DecomposedState - = 0; // for k pointer, the decomposed s_mat can not be reused + ELPA_Solver es( + (const bool)isReal, + comm, (const int)this->nbands, (const int)h_mat.row, (const int)h_mat.col, (const int *)h_mat.desc); + this->DecomposedState = 0; // for k pointer, the decomposed s_mat can not be reused ModuleBase::timer::start("DiagoElpa", "elpa_solve"); - es.generalized_eigenvector(h_mat.p, - s_mat.p, - this->DecomposedState, - eigen.data(), - psi.get_pointer()); + es.generalized_eigenvector( + h_mat.p, + s_mat.p, + this->DecomposedState, + eigen.data(), + psi.get_pointer()); ModuleBase::timer::end("DiagoElpa", "elpa_solve"); es.exit(); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); } template <> -void DiagoElpa::diag_pool(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in, - MPI_Comm& comm) +void DiagoElpa::diag_pool( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in, + MPI_Comm &comm) { std::vector eigen(this->nlocal, 0.0); bool isReal = true; - ELPA_Solver es((const bool)isReal, - comm, - (const int)this->nbands, - (const int)h_mat.row, - (const int)h_mat.col, - (const int*)h_mat.desc); + ELPA_Solver es( + (const bool)isReal, + comm, (const int)this->nbands, (const int)h_mat.row, (const int)h_mat.col, (const int *)h_mat.desc); ModuleBase::timer::start("DiagoElpa", "elpa_solve"); - es.generalized_eigenvector(h_mat.p, - s_mat.p, - this->DecomposedState, - eigen.data(), - psi.get_pointer()); + es.generalized_eigenvector( + h_mat.p, + s_mat.p, + this->DecomposedState, + eigen.data(), + psi.get_pointer()); ModuleBase::timer::end("DiagoElpa", "elpa_solve"); es.exit(); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); } #endif - - #ifdef __MPI -template -bool DiagoElpa::ifElpaHandle(const bool& newIteration, const bool& ifNSCF) const { +template bool DiagoElpa::ifElpaHandle(const bool &newIteration, const bool &ifNSCF) const +{ int doHandle = false; - if (newIteration) { + if (newIteration) + { doHandle = true; } - if (ifNSCF) { + if (ifNSCF) + { doHandle = true; } return doHandle; diff --git a/source/source_hsolver/diago_elpa.h b/source/source_hsolver/diago_elpa.h index e7e528f380f..aeced362504 100644 --- a/source/source_hsolver/diago_elpa.h +++ b/source/source_hsolver/diago_elpa.h @@ -1,7 +1,7 @@ #ifndef DIAGOELPA_H #define DIAGOELPA_H -#include "source_base/macros.h" // GetRealType +#include "source_base/macros.h" // GetRealType #include "source_base/matrix_block.h" #include "source_basis/module_ao/parallel_orbitals.h" #include "source_psi/psi.h" @@ -9,8 +9,7 @@ namespace hsolver { -template -class DiagoElpa +template class DiagoElpa { private: using Real = typename GetTypeReal::type; @@ -20,22 +19,28 @@ class DiagoElpa /// @param nbands_in number of lowest eigenpairs to compute DiagoElpa(const int nlocal_in, const int nbands_in) : nlocal(nlocal_in), nbands(nbands_in) {}; - void diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in); + void diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in); #ifdef __MPI // diagnolization used in parallel-k case - void diag_pool(ModuleBase::MatrixBlock& h_mat, ModuleBase::MatrixBlock& s_mat, psi::Psi& psi, Real* eigenvalue_in, MPI_Comm& comm); - MPI_Comm setmpicomm(); // set mpi comm; - static int elpa_num_thread; // need to set mpi_comm or not,-1 not,else the number of mpi needed + void diag_pool( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in, + MPI_Comm &comm); + MPI_Comm setmpicomm(); // set mpi comm; + static int elpa_num_thread; // need to set mpi_comm or not,-1 not,else the number of mpi needed #endif static int DecomposedState; private: #ifdef __MPI - bool ifElpaHandle(const bool& newIteration, const bool& ifNSCF) const; + bool ifElpaHandle(const bool &newIteration, const bool &ifNSCF) const; static int lastmpinum; // last using mpi; #endif @@ -43,12 +48,9 @@ class DiagoElpa const int nbands; }; -template -int DiagoElpa::DecomposedState = 0; -template -int DiagoElpa::lastmpinum = -1; -template -int DiagoElpa::elpa_num_thread = -1; +template int DiagoElpa::DecomposedState = 0; +template int DiagoElpa::lastmpinum = -1; +template int DiagoElpa::elpa_num_thread = -1; } // namespace hsolver #endif diff --git a/source/source_hsolver/diago_elpa_native.cpp b/source/source_hsolver/diago_elpa_native.cpp index 9d30afa590b..86057654346 100644 --- a/source/source_hsolver/diago_elpa_native.cpp +++ b/source/source_hsolver/diago_elpa_native.cpp @@ -1,19 +1,18 @@ #include "diago_elpa_native.h" +#include "omp.h" #include "source_base/global_function.h" -#include "source_base/module_external/blas_connector.h" #include "source_base/module_external/blacs_connector.h" +#include "source_base/module_external/blas_connector.h" #include "source_base/module_external/scalapack_connector.h" #include "source_base/timer.h" #include "source_base/tool_quit.h" #include "source_hsolver/module_genelpa/elpa_new.h" -#include "omp.h" namespace hsolver { #ifdef __MPI -template -MPI_Comm DiagoElpaNative::setmpicomm() +template MPI_Comm DiagoElpaNative::setmpicomm() { if (this->elpa_num_thread == -1) { @@ -30,14 +29,18 @@ MPI_Comm DiagoElpaNative::setmpicomm() else { lastmpinum++; - int* _ranks = new int[elpa_num_thread]; + int *_ranks = new int[elpa_num_thread]; for (int i = 0; i < elpa_num_thread; i++) { _ranks[i] = (lastmpinum + i) % _num; } MPI_Group _tempgroup, _oldgroup; MPI_Comm_group(MPI_COMM_WORLD, &_oldgroup); - MPI_Group_incl(_oldgroup, elpa_num_thread, _ranks, &_tempgroup); + MPI_Group_incl( + _oldgroup, + elpa_num_thread, + _ranks, + &_tempgroup); MPI_Comm _new_comm; MPI_Comm_create(MPI_COMM_WORLD, _tempgroup, &_new_comm); delete[] _ranks; @@ -49,11 +52,12 @@ MPI_Comm DiagoElpaNative::setmpicomm() #ifdef __MPI template -void DiagoElpaNative::diag_pool(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in, - MPI_Comm& comm) +void DiagoElpaNative::diag_pool( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in, + MPI_Comm &comm) { ModuleBase::timer::start("DiagoElpaNative", "elpa_solve"); @@ -68,7 +72,12 @@ void DiagoElpaNative::diag_pool(ModuleBase::MatrixBlock& h_mat, // cout<<"parameters are passed\n"; int nprows, npcols, myprow, mypcol; - Cblacs_gridinfo(cblacs_ctxt, &nprows, &npcols, &myprow, &mypcol); + Cblacs_gridinfo( + cblacs_ctxt, + &nprows, + &npcols, + &myprow, + &mypcol); std::vector eigen(this->nlocal, 0.0); std::vector eigenvectors(narows * nacols); @@ -78,60 +87,70 @@ void DiagoElpaNative::diag_pool(ModuleBase::MatrixBlock& h_mat, // private Hermitian copies before the solve. std::vector h_work; std::vector s_work; - T* h_elpa = h_mat.p; - T* s_elpa = s_mat.p; + T *h_elpa = h_mat.p; + T *s_elpa = s_mat.p; int decomposed_state = this->DecomposedState; if (!std::is_same::value) { h_work.resize(narows * nacols); s_work.resize(narows * nacols); const int one = 1; - ScalapackConnector::tranc(nFull, - nFull, - T(1.0), - h_mat.p, - one, - one, - h_mat.desc, - T(0.0), - h_work.data(), - one, - one, - h_mat.desc); - ScalapackConnector::tranc(nFull, - nFull, - T(1.0), - s_mat.p, - one, - one, - s_mat.desc, - T(0.0), - s_work.data(), - one, - one, - s_mat.desc); + ScalapackConnector::tranc( + nFull, + nFull, + T(1.0), + h_mat.p, + one, + one, + h_mat.desc, + T(0.0), + h_work.data(), + one, + one, + h_mat.desc); + ScalapackConnector::tranc( + nFull, + nFull, + T(1.0), + s_mat.p, + one, + one, + s_mat.desc, + T(0.0), + s_work.data(), + one, + one, + s_mat.desc); const auto local_to_global = [](const int local_index, - const int block_size, - const int process_coordinate, - const int source_coordinate, - const int process_count) { + const int block_size, + const int process_coordinate, + const int source_coordinate, + const int process_count) + { if (source_coordinate < 0) { return local_index; } - const int process_offset - = (process_coordinate - source_coordinate + process_count) % process_count; - return ((local_index / block_size) * process_count + process_offset) * block_size - + local_index % block_size; + const int process_offset = (process_coordinate - source_coordinate + process_count) % process_count; + return ((local_index / block_size) * process_count + process_offset) * block_size + + local_index % block_size; }; for (int local_col = 0; local_col < nacols; ++local_col) { - const int global_col - = local_to_global(local_col, h_mat.desc[5], mypcol, h_mat.desc[7], npcols); + const int global_col = local_to_global( + local_col, + h_mat.desc[5], + mypcol, + h_mat.desc[7], + npcols); for (int local_row = 0; local_row < narows; ++local_row) { - const int global_row - = local_to_global(local_row, h_mat.desc[4], myprow, h_mat.desc[6], nprows); + const int global_row = local_to_global( + local_row, + h_mat.desc[4], + myprow, + h_mat.desc[6], + nprows); if (global_row <= global_col) { const int local_index = local_row + local_col * narows; @@ -158,18 +177,53 @@ void DiagoElpaNative::diag_pool(ModuleBase::MatrixBlock& h_mat, #else int num_threads = 1; #endif - elpa_set(handle, "omp_threads", num_threads, &success); - elpa_set(handle, "na", (int)nFull, &success); - elpa_set(handle, "nev", (int)nev, &success); - elpa_set(handle, "local_nrows", (int)narows, &success); - elpa_set(handle, "local_ncols", (int)nacols, &success); - elpa_set(handle, "nblk", (int)nblk, &success); - elpa_set(handle, "mpi_comm_parent", (int)(MPI_Comm_c2f(comm)), &success); - elpa_set(handle, "process_row", (int)myprow, &success); - elpa_set(handle, "process_col", (int)mypcol, &success); - elpa_set(handle, "blacs_context", (int)cblacs_ctxt, &success); + elpa_set( + handle, + "omp_threads", + num_threads, + &success); + elpa_set( + handle, + "na", (int)nFull, + &success); + elpa_set( + handle, + "nev", (int)nev, + &success); + elpa_set( + handle, + "local_nrows", (int)narows, + &success); + elpa_set( + handle, + "local_ncols", (int)nacols, + &success); + elpa_set( + handle, + "nblk", (int)nblk, + &success); + elpa_set( + handle, + "mpi_comm_parent", (int)(MPI_Comm_c2f(comm)), + &success); + elpa_set( + handle, + "process_row", (int)myprow, + &success); + elpa_set( + handle, + "process_col", (int)mypcol, + &success); + elpa_set( + handle, + "blacs_context", (int)cblacs_ctxt, + &success); elpa_setup(handle); - elpa_set(handle, "solver", ELPA_SOLVER_1STAGE, &success); + elpa_set( + handle, + "solver", + ELPA_SOLVER_1STAGE, + &success); /* ELPA_WITH_NVIDIA_GPU_VERSION is a symbol defined in elpa/elpa_configured_options.h For example: @@ -181,19 +235,28 @@ void DiagoElpaNative::diag_pool(ModuleBase::MatrixBlock& h_mat, #if ELPA_WITH_NVIDIA_GPU_VERSION if (this->use_gpu) { - elpa_set(handle, "nvidia-gpu", 1, &success); - elpa_set(handle, "real_kernel", ELPA_2STAGE_REAL_NVIDIA_GPU, &success); + elpa_set( + handle, + "nvidia-gpu", + 1, + &success); + elpa_set( + handle, + "real_kernel", + ELPA_2STAGE_REAL_NVIDIA_GPU, + &success); elpa_setup_gpu(handle); } #endif - elpa_generalized_eigenvectors(handle, - h_elpa, - s_elpa, - eigen.data(), - eigenvectors.data(), - decomposed_state, - &success); + elpa_generalized_eigenvectors( + handle, + h_elpa, + s_elpa, + eigen.data(), + eigenvectors.data(), + decomposed_state, + &success); elpa_deallocate(handle, &success); elpa_uninit(&success); @@ -210,22 +273,38 @@ void DiagoElpaNative::diag_pool(ModuleBase::MatrixBlock& h_mat, } const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); const int size = psi.get_nbands() * psi.get_nbasis(); - BlasConnector::copy(size, eigenvectors.data(), inc, psi.get_pointer(), inc); + BlasConnector::copy( + size, + eigenvectors.data(), + inc, + psi.get_pointer(), + inc); } #endif template -void DiagoElpaNative::diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in) +void DiagoElpaNative::diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in) { ModuleBase::TITLE("DiagoElpaNative", "diag"); #ifdef __MPI MPI_Comm COMM_DIAG = setmpicomm(); // set mpi_comm needed - diag_pool(h_mat, s_mat, psi, eigenvalue_in, COMM_DIAG); + diag_pool( + h_mat, + s_mat, + psi, + eigenvalue_in, + COMM_DIAG); #else ModuleBase::WARNING_QUIT("DiagoElpaNative", "DiagoElpaNative only can be used with macro __MPI"); #endif diff --git a/source/source_hsolver/diago_elpa_native.h b/source/source_hsolver/diago_elpa_native.h index b938240299d..2b4ec1dec3d 100644 --- a/source/source_hsolver/diago_elpa_native.h +++ b/source/source_hsolver/diago_elpa_native.h @@ -1,7 +1,7 @@ #ifndef DIAGOELPANATIVE_H #define DIAGOELPANATIVE_H -#include "source_base/macros.h" // GetRealType +#include "source_base/macros.h" // GetRealType #include "source_base/matrix_block.h" #include "source_basis/module_ao/parallel_orbitals.h" #include "source_psi/psi.h" @@ -9,8 +9,7 @@ namespace hsolver { -template -class DiagoElpaNative +template class DiagoElpaNative { private: using Real = typename GetTypeReal::type; @@ -22,16 +21,22 @@ class DiagoElpaNative DiagoElpaNative(const int nlocal_in, const int nbands_in, const bool use_gpu_in) : nlocal(nlocal_in), nbands(nbands_in), use_gpu(use_gpu_in) {}; - void diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in); + void diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in); #ifdef __MPI // diagnolization used in parallel-k case - void diag_pool(ModuleBase::MatrixBlock& h_mat, ModuleBase::MatrixBlock& s_mat, psi::Psi& psi, Real* eigenvalue_in, MPI_Comm& comm); - MPI_Comm setmpicomm(); // set mpi comm; - static int elpa_num_thread; // need to set mpi_comm or not,-1 not,else the number of mpi needed - static int lastmpinum; // last using mpi; + void diag_pool( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in, + MPI_Comm &comm); + MPI_Comm setmpicomm(); // set mpi comm; + static int elpa_num_thread; // need to set mpi_comm or not,-1 not,else the number of mpi needed + static int lastmpinum; // last using mpi; #endif @@ -43,14 +48,11 @@ class DiagoElpaNative const bool use_gpu; }; -template -int DiagoElpaNative::DecomposedState = 0; +template int DiagoElpaNative::DecomposedState = 0; #ifdef __MPI -template -int DiagoElpaNative::lastmpinum = -1; -template -int DiagoElpaNative::elpa_num_thread = -1; +template int DiagoElpaNative::lastmpinum = -1; +template int DiagoElpaNative::elpa_num_thread = -1; #endif } // namespace hsolver diff --git a/source/source_hsolver/diago_iter_assist.cpp b/source/source_hsolver/diago_iter_assist.cpp index 92ebac863cb..61ee9f9aff9 100644 --- a/source/source_hsolver/diago_iter_assist.cpp +++ b/source/source_hsolver/diago_iter_assist.cpp @@ -1,4 +1,3 @@ -#include #include "diago_iter_assist.h" #include "source_base/constants.h" @@ -10,6 +9,7 @@ #include "source_hsolver/diag_comm_info.h" #include "source_hsolver/kernels/hegvd_op.h" +#include #include namespace hsolver @@ -21,16 +21,17 @@ namespace hsolver // Produces on output n_band eigenvectors (n_band <= nstart) in evc. //---------------------------------------------------------------------- template -void DiagoIterAssist::diag_subspace(const HSOperator& op, - const T* psi, - T* evc, - const int nstart, - const int n_band, - const int dmin, - const int dmax, - Real* en, - const diag_comm_info& diag_comm, - const bool S_orth) +void DiagoIterAssist::diag_subspace( + const HSOperator &op, + const T *psi, + T *evc, + const int nstart, + const int n_band, + const int dmin, + const int dmax, + Real *en, + const diag_comm_info &diag_comm, + const bool S_orth) { ModuleBase::TITLE("DiagoIterAssist", "diag_subspace"); ModuleBase::timer::start("DiagoIterAssist", "diag_subspace"); @@ -47,58 +48,70 @@ void DiagoIterAssist::diag_subspace(const HSOperator& op, setmem_complex_op()(hcc, 0, nstart * nstart); // scc is overlap matrix, only needed when psi is not orthogonal - if(!S_orth){ + if (!S_orth) + { resmem_complex_op()(scc, nstart * nstart, "DiagSub::scc"); setmem_complex_op()(scc, 0, nstart * nstart); } - + // vcc is eigenvector matrix of the reduced generalized eigenvalue problem resmem_complex_op()(vcc, nstart * nstart, "DiagSub::vcc"); setmem_complex_op()(vcc, 0, nstart * nstart); // temp holds H|psi>, then S|psi>, then the rotated vectors; it is separate // from evc so that evc may alias psi - T* temp = nullptr; + T *temp = nullptr; const int local_storage = std::max(nstart * dmax, 1); resmem_complex_op()(temp, local_storage, "DiagSub::temp"); setmem_complex_op()(temp, 0, nstart * dmax); { // code block to calculate hcc and scc T *hpsi = temp; - op.hpsi(psi, hpsi, dmax, nstart); - - ModuleBase::gemm_op()('C', - 'N', - nstart, - nstart, - dmin, - &one, - psi, - matrix_ld, - hpsi, - matrix_ld, - &zero, - hcc, - nstart); - - if(!S_orth){ + op.hpsi( + psi, + hpsi, + dmax, + nstart); + + ModuleBase::gemm_op()( + 'C', + 'N', + nstart, + nstart, + dmin, + &one, + psi, + matrix_ld, + hpsi, + matrix_ld, + &zero, + hcc, + nstart); + + if (!S_orth) + { // Only calculate S_sub if not orthogonal T *spsi = temp; - op.spsi(psi, spsi, dmax, nstart); - - ModuleBase::gemm_op()('C', - 'N', - nstart, - nstart, - dmin, - &one, - psi, - matrix_ld, - spsi, - matrix_ld, - &zero, - scc, - nstart); + op.spsi( + psi, + spsi, + dmax, + nstart); + + ModuleBase::gemm_op()( + 'C', + 'N', + nstart, + nstart, + dmin, + &one, + psi, + matrix_ld, + spsi, + matrix_ld, + &zero, + scc, + nstart); } } @@ -106,42 +119,67 @@ void DiagoIterAssist::diag_subspace(const HSOperator& op, { #ifdef __MPI Parallel_Common::reduce_dev(hcc, nstart * nstart, diag_comm.comm); - if(!S_orth){ + if (!S_orth) + { Parallel_Common::reduce_dev(scc, nstart * nstart, diag_comm.comm); } #endif } // after generation of H and (optionally) S matrix, diag them - if (S_orth) { + if (S_orth) + { // Solve standard eigenproblem: H_sub * y = lambda * y - DiagoIterAssist::diag_heevx(nstart, n_band, hcc, nstart, en, vcc); - } else { + DiagoIterAssist::diag_heevx( + nstart, + n_band, + hcc, + nstart, + en, + vcc); + } + else + { // Solve generalized eigenproblem: H_sub * y = lambda * S_sub * y - DiagoIterAssist::diag_hegvd(nstart, n_band, hcc, scc, nstart, en, vcc); + DiagoIterAssist::diag_hegvd( + nstart, + n_band, + hcc, + scc, + nstart, + en, + vcc); } { // code block to calculate evc - ModuleBase::gemm_op()('N', - 'N', - dmin, - n_band, - nstart, - &one, - psi, // dmin * nstart - matrix_ld, - vcc, // nstart * n_band - nstart, - &zero, - temp, - rotated_ld); + ModuleBase::gemm_op()( + 'N', + 'N', + dmin, + n_band, + nstart, + &one, + psi, // dmin * nstart + matrix_ld, + vcc, // nstart * n_band + nstart, + &zero, + temp, + rotated_ld); } - ModuleBase::matrixCopy()(n_band, dmin, temp, dmin, evc, dmax); + ModuleBase::matrixCopy()( + n_band, + dmin, + temp, + dmin, + evc, + dmax); delmem_complex_op()(temp); delmem_complex_op()(hcc); - if(!S_orth){ + if (!S_orth) + { delmem_complex_op()(scc); } delmem_complex_op()(vcc); @@ -150,14 +188,15 @@ void DiagoIterAssist::diag_subspace(const HSOperator& op, } template -void DiagoIterAssist::diag_subspace(const HSOperator& op, - const psi::Psi& psi, // [in] wavefunction - psi::Psi& evc, // [out] wavefunction, eigenvectors - Real* en, // [out] eigenvalues - const diag_comm_info& diag_comm, - int n_band, // [in] number of bands to be calculated, also number of rows - // of evc, if set to 0, n_band = nstart, default 0 - const bool S_orth // [in] if true, psi is assumed to be already S-orthogonalized +void DiagoIterAssist::diag_subspace( + const HSOperator &op, + const psi::Psi &psi, // [in] wavefunction + psi::Psi &evc, // [out] wavefunction, eigenvectors + Real *en, // [out] eigenvalues + const diag_comm_info &diag_comm, + int n_band, // [in] number of bands to be calculated, also number of rows + // of evc, if set to 0, n_band = nstart, default 0 + const bool S_orth // [in] if true, psi is assumed to be already S-orthogonalized ) { // two case: @@ -172,28 +211,30 @@ void DiagoIterAssist::diag_subspace(const HSOperator& op, // dmin is the active number of plane waves or atomic orbitals // dmax is the leading dimension of psi - diag_subspace(op, - psi.get_pointer(), - evc.get_pointer(), - nstart, - n_band, - psi.get_current_ngk(), - psi.get_nbasis(), - en, - diag_comm, - S_orth); + diag_subspace( + op, + psi.get_pointer(), + evc.get_pointer(), + nstart, + n_band, + psi.get_current_ngk(), + psi.get_nbasis(), + en, + diag_comm, + S_orth); } template -void DiagoIterAssist::diag_subspace_init(const HSOperator& op, - const T* psi, - int psi_nr, - int psi_nc, - psi::Psi& evc, - Real* en, - const std::string& basis_type, - const std::string& calculation, - const diag_comm_info& diag_comm) +void DiagoIterAssist::diag_subspace_init( + const HSOperator &op, + const T *psi, + int psi_nr, + int psi_nc, + psi::Psi &evc, + Real *en, + const std::string &basis_type, + const std::string &calculation, + const diag_comm_info &diag_comm) { ModuleBase::TITLE("DiagoIterAssist", "diag_subspace_init"); ModuleBase::timer::start("DiagoIterAssist", "diag_subspace_init"); @@ -218,55 +259,109 @@ void DiagoIterAssist::diag_subspace_init(const HSOperator& if (base_device::get_device_type(ctx) == base_device::GpuDevice) { // band by band on the GPU: the scratch buffer holds one vector only - T* temp = nullptr; + T *temp = nullptr; resmem_complex_op()(temp, psi_nc, "DiagSub::temp"); setmem_complex_op()(temp, 0, psi_nc); - T* hpsi = temp; + T *hpsi = temp; for (int i = 0; i < nstart; i++) { // H|Psi> to get hpsi for target band - op.hpsi(psi + i * psi_nc, hpsi, psi_nc, 1); + op.hpsi( + psi + i * psi_nc, + hpsi, + psi_nc, + 1); // calculate the related elements in hcc - ModuleBase::gemv_op()('C', psi_nc, nstart, &one, psi, psi_nc, hpsi, 1, &zero, hcc + i * nstart, 1); + ModuleBase::gemv_op()( + 'C', + psi_nc, + nstart, + &one, + psi, + psi_nc, + hpsi, + 1, + &zero, + hcc + i * nstart, + 1); } - T* spsi = temp; + T *spsi = temp; for (int i = 0; i < nstart; i++) { - op.spsi(psi + i * psi_nc, spsi, psi_nc, 1); - - ModuleBase::gemv_op()('C', - psi_nc, - nstart, - &one, - psi, - psi_nc, // nbasis - spsi, - 1, - &zero, - scc + i * nstart, - 1); + op.spsi( + psi + i * psi_nc, + spsi, + psi_nc, + 1); + + ModuleBase::gemv_op()( + 'C', + psi_nc, + nstart, + &one, + psi, + psi_nc, // nbasis + spsi, + 1, + &zero, + scc + i * nstart, + 1); } delmem_complex_op()(temp); } else if (base_device::get_device_type(ctx) == base_device::CpuDevice) { // hpsi and spsi share the temp space - T* temp = nullptr; + T *temp = nullptr; resmem_complex_op()(temp, nstart * psi_nc, "DiagSub::temp"); setmem_complex_op()(temp, 0, nstart * psi_nc); - T* hpsi = temp; - op.hpsi(psi, hpsi, psi_nc, nstart); - - ModuleBase::gemm_op()('C', 'N', nstart, nstart, dmin, &one, psi, psi_nc, hpsi, psi_nc, &zero, hcc, nstart); - - T* spsi = temp; - op.spsi(psi, spsi, psi_nc, nstart); - - ModuleBase::gemm_op()('C', 'N', nstart, nstart, dmin, &one, psi, psi_nc, spsi, psi_nc, &zero, scc, nstart); + T *hpsi = temp; + op.hpsi( + psi, + hpsi, + psi_nc, + nstart); + + ModuleBase::gemm_op()( + 'C', + 'N', + nstart, + nstart, + dmin, + &one, + psi, + psi_nc, + hpsi, + psi_nc, + &zero, + hcc, + nstart); + + T *spsi = temp; + op.spsi( + psi, + spsi, + psi_nc, + nstart); + + ModuleBase::gemm_op()( + 'C', + 'N', + nstart, + nstart, + dmin, + &one, + psi, + psi_nc, + spsi, + psi_nc, + &zero, + scc, + nstart); delmem_complex_op()(temp); } @@ -282,7 +377,14 @@ void DiagoIterAssist::diag_subspace_init(const HSOperator& } // after generation of H and S matrix, diag them - DiagoIterAssist::diag_hegvd(nstart, n_band, hcc, scc, nstart, en, vcc); + DiagoIterAssist::diag_hegvd( + nstart, + n_band, + hcc, + scc, + nstart, + en, + vcc); op.export_subspace_vec(vcc, nstart, n_band); @@ -293,44 +395,45 @@ void DiagoIterAssist::diag_subspace_init(const HSOperator& { // The caller requested eigenvalues only, so no wavefunction rotation is needed. } - else if ((basis_type == "lcao" || basis_type == "lcao_in_pw" || basis_type == "pw") - && (calculation == "scf" || calculation == "md" - || calculation == "relax")) // pengfei 2014-10-13 + else if ((basis_type == "lcao" || basis_type == "lcao_in_pw" || basis_type == "pw") && + (calculation == "scf" || calculation == "md" || calculation == "relax")) // pengfei 2014-10-13 { // because psi and evc are different here, // I think if psi and evc are the same, // there may be problems, mohan 2011-01-01 - ModuleBase::gemm_op()('N', - 'N', - dmax, - n_band, - nstart, - &one, - psi, // dmax * nstart - dmax, - vcc, // nstart * n_band - nstart, - &zero, - evc.get_pointer(), - dmax); + ModuleBase::gemm_op()( + 'N', + 'N', + dmax, + n_band, + nstart, + &one, + psi, // dmax * nstart + dmax, + vcc, // nstart * n_band + nstart, + &zero, + evc.get_pointer(), + dmax); } else { assert(psi != evc.get_pointer()); - ModuleBase::gemm_op()('N', - 'N', - dmin, - n_band, - nstart, - &one, - psi, // dmin * nstart - dmax, - vcc, // nstart * n_band - nstart, - &zero, - evc.get_pointer(), - dmax); + ModuleBase::gemm_op()( + 'N', + 'N', + dmin, + n_band, + nstart, + &one, + psi, // dmin * nstart + dmax, + vcc, // nstart * n_band + nstart, + &zero, + evc.get_pointer(), + dmax); } delmem_complex_op()(hcc); @@ -340,12 +443,13 @@ void DiagoIterAssist::diag_subspace_init(const HSOperator& } template -void DiagoIterAssist::diag_heevx(const int matrix_size, - const int num_eigenpairs, - const T *h, - const int ldh, - Real *e, // always in CPU - T *v) +void DiagoIterAssist::diag_heevx( + const int matrix_size, + const int num_eigenpairs, + const T *h, + const int ldh, + Real *e, // always in CPU + T *v) { ModuleBase::TITLE("DiagoIterAssist", "diag_heevx"); ModuleBase::timer::start("DiagoIterAssist", "diag_heevx"); @@ -355,8 +459,16 @@ void DiagoIterAssist::diag_heevx(const int matrix_size, resmem_var_op()(eigenvalues, matrix_size); setmem_var_op()(eigenvalues, 0, matrix_size); - // (const Device *d, const int matrix_size, const int lda, const T *A, const int num_eigenpairs, Real *eigenvalues, T *eigenvectors); - heevx_op()(ctx, matrix_size, ldh, h, num_eigenpairs, eigenvalues, v); + // (const Device *d, const int matrix_size, const int lda, const T *A, const int num_eigenpairs, Real *eigenvalues, + // T *eigenvectors); + heevx_op()( + ctx, + matrix_size, + ldh, + h, + num_eigenpairs, + eigenvalues, + v); if (base_device::get_device_type(ctx) == base_device::GpuDevice) { @@ -377,13 +489,14 @@ void DiagoIterAssist::diag_heevx(const int matrix_size, } template -void DiagoIterAssist::diag_hegvd(const int nstart, - const int nbands, - const T *hcc, - T *scc, - const int ldh, // nstart - Real *e, // always in CPU - T *vcc) +void DiagoIterAssist::diag_hegvd( + const int nstart, + const int nbands, + const T *hcc, + T *scc, + const int ldh, // nstart + Real *e, // always in CPU + T *vcc) { ModuleBase::TITLE("DiagoIterAssist", "diag_hegvd"); ModuleBase::timer::start("DiagoIterAssist", "diag_hegvd"); @@ -392,7 +505,14 @@ void DiagoIterAssist::diag_hegvd(const int nstart, resmem_var_op()(eigenvalues, nstart); setmem_var_op()(eigenvalues, 0, nstart); - hegvd_op()(ctx, nstart, ldh, hcc, scc, eigenvalues, vcc); + hegvd_op()( + ctx, + nstart, + ldh, + hcc, + scc, + eigenvalues, + vcc); if (base_device::get_device_type(ctx) == base_device::GpuDevice) { @@ -428,58 +548,69 @@ void DiagoIterAssist::diag_hegvd(const int nstart, } template -void DiagoIterAssist::cal_hs_subspace(const HSOperator& op, - const psi::Psi& psi, // [in] wavefunction - T* hcc, - T* scc, - const diag_comm_info& diag_comm) +void DiagoIterAssist::cal_hs_subspace( + const HSOperator &op, + const psi::Psi &psi, // [in] wavefunction + T *hcc, + T *scc, + const diag_comm_info &diag_comm) { const int nstart = psi.get_nbands(); - + setmem_complex_op()(hcc, 0, nstart * nstart); setmem_complex_op()(scc, 0, nstart * nstart); const int dmin = psi.get_current_ngk(); const int dmax = psi.get_nbasis(); - T* temp = nullptr; + T *temp = nullptr; resmem_complex_op()(temp, nstart * dmax, "DiagSub::temp"); setmem_complex_op()(temp, 0, nstart * dmax); { // code block to calculate hcc and scc - T* hpsi = temp; - op.hpsi(psi.get_pointer(), hpsi, dmax, nstart); - - ModuleBase::gemm_op()('C', - 'N', - nstart, - nstart, - dmin, - &one, - psi.get_pointer(), - dmax, - hpsi, - dmax, - &zero, - hcc, - nstart); - - T* spsi = temp; - op.spsi(psi.get_pointer(), spsi, dmax, nstart); - - ModuleBase::gemm_op()('C', - 'N', - nstart, - nstart, - dmin, - &one, - psi.get_pointer(), - dmax, - spsi, - dmax, - &zero, - scc, - nstart); + T *hpsi = temp; + op.hpsi( + psi.get_pointer(), + hpsi, + dmax, + nstart); + + ModuleBase::gemm_op()( + 'C', + 'N', + nstart, + nstart, + dmin, + &one, + psi.get_pointer(), + dmax, + hpsi, + dmax, + &zero, + hcc, + nstart); + + T *spsi = temp; + op.spsi( + psi.get_pointer(), + spsi, + dmax, + nstart); + + ModuleBase::gemm_op()( + 'C', + 'N', + nstart, + nstart, + dmin, + &one, + psi.get_pointer(), + dmax, + spsi, + dmax, + &zero, + scc, + nstart); } if (diag_comm.nproc > 1) @@ -494,13 +625,14 @@ void DiagoIterAssist::cal_hs_subspace(const HSOperator& op } template -void DiagoIterAssist::diag_responce( const T* hcc, - T* scc, - const int nbands, - const T* mat_in, // [out] target matrix to be multiplied - T* mat_out, - int mat_col, // [in] number of columns of target matrix - Real* en // [out] eigenvalues +void DiagoIterAssist::diag_responce( + const T *hcc, + T *scc, + const int nbands, + const T *mat_in, // [out] target matrix to be multiplied + T *mat_out, + int mat_col, // [in] number of columns of target matrix + Real *en // [out] eigenvalues ) { ModuleBase::TITLE("DiagoIterAssist", "diag_responce"); @@ -513,22 +645,30 @@ void DiagoIterAssist::diag_responce( const T* hcc, setmem_complex_op()(vcc, 0, nstart * nstart); // after generation of H and S matrix, diag them - DiagoIterAssist::diag_hegvd(nstart, nstart, hcc, scc, nstart, en, vcc); + DiagoIterAssist::diag_hegvd( + nstart, + nstart, + hcc, + scc, + nstart, + en, + vcc); { // code block to calculate tar_mat - ModuleBase::gemm_op()('N', - 'N', - mat_col, - nstart, - nstart, - &one, - mat_in, // mat_col * nstart - mat_col, - vcc, // nstart * nstart - nstart, - &zero, - mat_out, - mat_col); + ModuleBase::gemm_op()( + 'N', + 'N', + mat_col, + nstart, + nstart, + &one, + mat_in, // mat_col * nstart + mat_col, + vcc, // nstart * nstart + nstart, + &zero, + mat_out, + mat_col); } delmem_complex_op()(vcc); @@ -537,12 +677,12 @@ void DiagoIterAssist::diag_responce( const T* hcc, } template -void DiagoIterAssist::diag_subspace_psi(const T* hcc, - T* scc, - const int dim_subspace, - psi::Psi& evc, - Real* en -) +void DiagoIterAssist::diag_subspace_psi( + const T *hcc, + T *scc, + const int dim_subspace, + psi::Psi &evc, + Real *en) { ModuleBase::TITLE("DiagoIterAssist", "diag_subspace_psi"); ModuleBase::timer::start("DiagoIterAssist", "diag_subspace_psi"); @@ -555,28 +695,42 @@ void DiagoIterAssist::diag_subspace_psi(const T* hcc, setmem_complex_op()(vcc, 0, nstart * nstart); // after generation of H and S matrix, diag them - DiagoIterAssist::diag_hegvd(nstart, nstart, hcc, scc, nstart, en, vcc); + DiagoIterAssist::diag_hegvd( + nstart, + nstart, + hcc, + scc, + nstart, + en, + vcc); { // code block to calculate tar_mat const int dmin = evc.get_current_ngk(); const int dmax = evc.get_nbasis(); - T* temp = nullptr; + T *temp = nullptr; resmem_complex_op()(temp, nstart * dmax, "DiagSub::temp"); setmem_complex_op()(temp, 0, nstart * dmax); - ModuleBase::gemm_op()('N', - 'N', - dmin, - n_band, - nstart, - &one, - evc.get_pointer(), // dmin * nstart - dmax, - vcc, // nstart * n_band - nstart, - &zero, - temp, - dmin); - ModuleBase::matrixCopy()(n_band, dmin, temp, dmin, evc.get_pointer(), dmax); + ModuleBase::gemm_op()( + 'N', + 'N', + dmin, + n_band, + nstart, + &one, + evc.get_pointer(), // dmin * nstart + dmax, + vcc, // nstart * n_band + nstart, + &zero, + temp, + dmin); + ModuleBase::matrixCopy()( + n_band, + dmin, + temp, + dmin, + evc.get_pointer(), + dmax); delmem_complex_op()(temp); } diff --git a/source/source_hsolver/diago_iter_assist.h b/source/source_hsolver/diago_iter_assist.h index ef89de59bee..e5cbe10850d 100644 --- a/source/source_hsolver/diago_iter_assist.h +++ b/source/source_hsolver/diago_iter_assist.h @@ -13,8 +13,7 @@ namespace hsolver struct diag_comm_info; -template -class DiagoIterAssist +template class DiagoIterAssist { private: using Real = typename GetTypeReal::type; @@ -47,26 +46,28 @@ class DiagoIterAssist * @param is_S_orthogonal if true, psi is already S-orthonormal and the * standard eigenproblem is solved instead of the generalized one */ - static void diag_subspace(const HSOperator& op, - const T* psi, - T* evc, - const int nstart, - const int n_band, - const int dmin, - const int dmax, - Real* en, - const diag_comm_info& diag_comm, - const bool is_S_orthogonal = false); + static void diag_subspace( + const HSOperator &op, + const T *psi, + T *evc, + const int nstart, + const int n_band, + const int dmin, + const int dmax, + Real *en, + const diag_comm_info &diag_comm, + const bool is_S_orthogonal = false); /// psi::Psi flavour of diag_subspace(): nstart = psi.get_nbands(), /// n_band = 0 means all of them, dimensions taken from psi. - static void diag_subspace(const HSOperator& op, - const psi::Psi& psi, - psi::Psi& evc, - Real* en, - const diag_comm_info& diag_comm, - int n_band = 0, - const bool is_S_orthogonal = false); + static void diag_subspace( + const HSOperator &op, + const psi::Psi &psi, + psi::Psi &evc, + Real *en, + const diag_comm_info &diag_comm, + int n_band = 0, + const bool is_S_orthogonal = false); /// @brief subspace diagonalization used to build the starting wavefunction /// @param op interface to H and S; op.add_to_subspace_h() and @@ -79,40 +80,44 @@ class DiagoIterAssist /// @param basis_type "lcao", "lcao_in_pw" or "pw"; together with calculation it selects /// how the rotation matrix is applied to psi /// @param calculation "scf", "nscf", "md", "relax", ... - static void diag_subspace_init(const HSOperator& op, - const T* psi, - int psi_nr, - int psi_nc, - psi::Psi& evc, - Real* en, - const std::string& basis_type, - const std::string& calculation, - const diag_comm_info& diag_comm); - - static void diag_heevx(const int nstart, - const int nbands, - const T *hcc, - const int ldh, - Real *e, - T *vcc); - static void diag_hegvd(const int nstart, - const int nbands, - const T *hcc, - T *sc, - const int ldh, // nstart - Real *e, - T *vcc); + static void diag_subspace_init( + const HSOperator &op, + const T *psi, + int psi_nr, + int psi_nc, + psi::Psi &evc, + Real *en, + const std::string &basis_type, + const std::string &calculation, + const diag_comm_info &diag_comm); + + static void diag_heevx( + const int nstart, + const int nbands, + const T *hcc, + const int ldh, + Real *e, + T *vcc); + static void diag_hegvd( + const int nstart, + const int nbands, + const T *hcc, + T *sc, + const int ldh, // nstart + Real *e, + T *vcc); /// @brief calculate Hamiltonian and overlap matrix in subspace spanned by nstart states psi /// @param op : applies H and S /// @param psi : wavefunction /// @param hcc : Hamiltonian matrix /// @param scc : overlap matrix - static void cal_hs_subspace(const HSOperator& op, - const psi::Psi& psi, // [in] wavefunction - T* hcc, - T* scc, - const diag_comm_info& diag_comm); + static void cal_hs_subspace( + const HSOperator &op, + const psi::Psi &psi, // [in] wavefunction + T *hcc, + T *scc, + const diag_comm_info &diag_comm); /// @brief calculate the response matrix from rotation matrix solved by diagonalization of H and S matrix /// @param hcc : Hamiltonian matrix @@ -122,32 +127,34 @@ class DiagoIterAssist /// @param mat_out : output matrix to be rotated /// @param mat_col : number of columns of target matrix /// @param en : eigenvalues - static void diag_responce(const T* hcc, - T* scc, - const int nbands, - const T* mat_in, - T* mat_out, - int mat_col, - Real* en); - + static void diag_responce( + const T *hcc, + T *scc, + const int nbands, + const T *mat_in, + T *mat_out, + int mat_col, + Real *en); + /// @brief calculate the response wavefunction psi from rotation matrix solved by diagonalization of H and S matrix - static void diag_subspace_psi(const T* hcc, - T* scc, - const int dim_subspace, - psi::Psi& evc, - Real* en); + static void diag_subspace_psi( + const T *hcc, + T *scc, + const int dim_subspace, + psi::Psi &evc, + Real *en); private: - constexpr static const Device* ctx = {}; + constexpr static const Device *ctx = {}; using setmem_var_op = base_device::memory::set_memory_op; using resmem_var_op = base_device::memory::resize_memory_op; using delmem_var_op = base_device::memory::delete_memory_op; using syncmem_var_op = base_device::memory::synchronize_memory_op; - using syncmem_var_h2d_op - = base_device::memory::synchronize_memory_op; - using syncmem_var_d2h_op - = base_device::memory::synchronize_memory_op; + using syncmem_var_h2d_op = + base_device::memory::synchronize_memory_op; + using syncmem_var_d2h_op = + base_device::memory::synchronize_memory_op; using setmem_complex_op = base_device::memory::set_memory_op; using resmem_complex_op = base_device::memory::resize_memory_op; @@ -163,23 +170,18 @@ class DiagoIterAssist template typename DiagoIterAssist::Real DiagoIterAssist::avg_iter = 0.0; -template -int DiagoIterAssist::PW_DIAG_NMAX = 30; +template int DiagoIterAssist::PW_DIAG_NMAX = 30; template typename DiagoIterAssist::Real DiagoIterAssist::PW_DIAG_THR = 1.0e-2; -template -bool DiagoIterAssist::need_subspace = false; +template bool DiagoIterAssist::need_subspace = false; -template -int DiagoIterAssist::SCF_ITER = 0; +template int DiagoIterAssist::SCF_ITER = 0; -template -T DiagoIterAssist::one = static_cast(1.0); +template T DiagoIterAssist::one = static_cast(1.0); -template -T DiagoIterAssist::zero = static_cast(0.0); +template T DiagoIterAssist::zero = static_cast(0.0); } // namespace hsolver #endif diff --git a/source/source_hsolver/diago_lapack.cpp b/source/source_hsolver/diago_lapack.cpp index daf6e2f5889..a0ea65eda1e 100644 --- a/source/source_hsolver/diago_lapack.cpp +++ b/source/source_hsolver/diago_lapack.cpp @@ -13,26 +13,25 @@ namespace hsolver namespace { template -void check_lapack_layout(const ModuleBase::MatrixBlock& h_mat, - const ModuleBase::MatrixBlock& s_mat, - const std::size_t n) +void check_lapack_layout( + const ModuleBase::MatrixBlock &h_mat, const ModuleBase::MatrixBlock &s_mat, const std::size_t n) { if (h_mat.row != n || h_mat.col != n || s_mat.row != n || s_mat.col != n) { ModuleBase::WARNING_QUIT( - "DiagoLapack", - "The LAPACK eigensolver requires replicated " + std::to_string(n) + " x " - + std::to_string(n) + " Hamiltonian and overlap matrices, but received " - + std::to_string(h_mat.row) + " x " + std::to_string(h_mat.col) - + " local blocks. Please use ScaLAPACK or ELPA for distributed matrices."); + "DiagoLapack", "The LAPACK eigensolver requires replicated " + std::to_string(n) + " x " + + std::to_string(n) + " Hamiltonian and overlap matrices, but received " + + std::to_string(h_mat.row) + " x " + std::to_string(h_mat.col) + + " local blocks. Please use ScaLAPACK or ELPA for distributed matrices."); } } } // namespace template <> -void DiagoLapack::diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in) +void DiagoLapack::diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in) { ModuleBase::TITLE("DiagoLapack", "diag"); assert(h_mat.col == s_mat.col && h_mat.row == s_mat.row && h_mat.desc == s_mat.desc); @@ -40,67 +39,115 @@ void DiagoLapack::diag(ModuleBase::MatrixBlock& h_mat, check_lapack_layout(h_mat, s_mat, eigen.size()); // Diag - this->dsygvx_diag(h_mat.col, h_mat.row, h_mat.p, s_mat.p, eigen.data(), psi); + this->dsygvx_diag( + h_mat.col, + h_mat.row, + h_mat.p, + s_mat.p, + eigen.data(), + psi); // Copy result const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); } template <> -void DiagoLapack>::diag(ModuleBase::MatrixBlock>& h_mat, - ModuleBase::MatrixBlock>& s_mat, - psi::Psi>& psi, - Real* eigenvalue_in) +void DiagoLapack>::diag( + ModuleBase::MatrixBlock> &h_mat, + ModuleBase::MatrixBlock> &s_mat, + psi::Psi> &psi, + Real *eigenvalue_in) { ModuleBase::TITLE("DiagoLapack", "diag"); assert(h_mat.col == s_mat.col && h_mat.row == s_mat.row && h_mat.desc == s_mat.desc); std::vector eigen(this->nlocal, 0.0); check_lapack_layout(h_mat, s_mat, eigen.size()); - this->zhegvx_diag(h_mat.col, h_mat.row, h_mat.p, s_mat.p, eigen.data(), psi); + this->zhegvx_diag( + h_mat.col, + h_mat.row, + h_mat.p, + s_mat.p, + eigen.data(), + psi); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); } #ifdef __MPI - template<> - void DiagoLapack::diag_pool(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in, - MPI_Comm& comm) +template <> +void DiagoLapack::diag_pool( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in, + MPI_Comm &comm) { ModuleBase::TITLE("DiagoLapack", "diag_pool"); assert(h_mat.col == s_mat.col && h_mat.row == s_mat.row && h_mat.desc == s_mat.desc); std::vector eigen(this->nlocal, 0.0); check_lapack_layout(h_mat, s_mat, eigen.size()); - this->dsygvx_diag(h_mat.col, h_mat.row, h_mat.p, s_mat.p, eigen.data(), psi); + this->dsygvx_diag( + h_mat.col, + h_mat.row, + h_mat.p, + s_mat.p, + eigen.data(), + psi); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); } - template<> - void DiagoLapack>::diag_pool(ModuleBase::MatrixBlock>& h_mat, - ModuleBase::MatrixBlock>& s_mat, - psi::Psi>& psi, - Real* eigenvalue_in, - MPI_Comm& comm) +template <> +void DiagoLapack>::diag_pool( + ModuleBase::MatrixBlock> &h_mat, + ModuleBase::MatrixBlock> &s_mat, + psi::Psi> &psi, + Real *eigenvalue_in, + MPI_Comm &comm) { ModuleBase::TITLE("DiagoLapack", "diag_pool"); assert(h_mat.col == s_mat.col && h_mat.row == s_mat.row && h_mat.desc == s_mat.desc); std::vector eigen(this->nlocal, 0.0); check_lapack_layout(h_mat, s_mat, eigen.size()); - this->zhegvx_diag(h_mat.col, h_mat.row, h_mat.p, s_mat.p, eigen.data(), psi); + this->zhegvx_diag( + h_mat.col, + h_mat.row, + h_mat.p, + s_mat.p, + eigen.data(), + psi); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); } #endif template -std::pair> DiagoLapack::dsygvx_once(const int ncol, - const int nrow, - const double* const h_mat, - const double* const s_mat, - double* const ekb, - psi::Psi& wfc_2d) const +std::pair> DiagoLapack::dsygvx_once( + const int ncol, + const int nrow, + const double *const h_mat, + const double *const s_mat, + double *const ekb, + psi::Psi &wfc_2d) const { ModuleBase::matrix h_tmp(ncol, nrow, false); memcpy(h_tmp.c, h_mat, sizeof(double) * ncol * nrow); @@ -125,7 +172,8 @@ std::pair> DiagoLapack::dsygvx_once(const int ncol, // ABSTOL, M, W, Z, LDZ, WORK, LWORK, IWORK, IFAIL, INFO) int n = this->nlocal; int lda = n, ldb = n, ldz = n; - dsygvx_(&itype, + dsygvx_( + &itype, &jobz, &range, &uplo, @@ -148,10 +196,11 @@ std::pair> DiagoLapack::dsygvx_once(const int ncol, iwork.data(), ifail.data(), &info); - if (info) { - throw std::runtime_error("info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" - + std::string(__FILE__) + " line " - + std::to_string(__LINE__)); + if (info) + { + throw std::runtime_error( + "info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" + std::string(__FILE__) + " line " + + std::to_string(__LINE__)); } // Query returned optimal lwork in work[0] @@ -161,56 +210,71 @@ std::pair> DiagoLapack::dsygvx_once(const int ncol, liwork = std::max(1, 5 * n); iwork.resize(liwork, 0); - dsygvx_(&itype, - &jobz, - &range, - &uplo, - &n, - h_tmp.c, - &lda, - s_tmp.c, - &ldb, - &vl, - &vu, - &il, - &iu, - &abstol, - &M, - ekb, - wfc_2d.get_pointer(), - &ldz, - work.data(), - &lwork, - iwork.data(), - ifail.data(), - &info); + dsygvx_( + &itype, + &jobz, + &range, + &uplo, + &n, + h_tmp.c, + &lda, + s_tmp.c, + &ldb, + &vl, + &vu, + &il, + &iu, + &abstol, + &M, + ekb, + wfc_2d.get_pointer(), + &ldz, + work.data(), + &lwork, + iwork.data(), + ifail.data(), + &info); - if (info == 0) { + if (info == 0) + { return std::make_pair(info, std::vector{}); - } else if (info < 0) { + } + else if (info < 0) + { return std::make_pair(info, std::vector{}); - } else if (info % 2) { + } + else if (info % 2) + { return std::make_pair(info, ifail); - } else if (info / 2 % 2) { + } + else if (info / 2 % 2) + { return std::make_pair(info, iclustr); - } else if (info / 4 % 2) { + } + else if (info / 4 % 2) + { return std::make_pair(info, std::vector{M, NZ}); - } else if (info / 16 % 2) { + } + else if (info / 16 % 2) + { return std::make_pair(info, ifail); - } else { - throw std::runtime_error("info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" - + std::string(__FILE__) + " line " - + std::to_string(__LINE__)); + } + else + { + throw std::runtime_error( + "info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" + std::string(__FILE__) + " line " + + std::to_string(__LINE__)); } } template -std::pair> DiagoLapack::zhegvx_once(const int ncol, - const int nrow, - const std::complex* const h_mat, - const std::complex* const s_mat, - double* const ekb, - psi::Psi>& wfc_2d) const +std::pair> DiagoLapack::zhegvx_once( + const int ncol, + const int nrow, + const std::complex *const h_mat, + const std::complex *const s_mat, + double *const ekb, + psi::Psi> &wfc_2d) const { ModuleBase::ComplexMatrix h_tmp(ncol, nrow, false); memcpy(h_tmp.c, h_mat, sizeof(std::complex) * ncol * nrow); @@ -221,7 +285,7 @@ std::pair> DiagoLapack::zhegvx_once(const int ncol, const int itype = 1, il = 1, iu = this->nbands, one = 1; int M = 0, NZ = 0, lwork = -1, lrwork = -1, liwork = -1, info = 0; const double abstol = LAPACK_ABSTOL, orfac = LAPACK_ORFAC; - + const double vl = 0, vu = 0; std::vector> work(1, 0); std::vector rwork(3, 0); @@ -266,10 +330,11 @@ std::pair> DiagoLapack::zhegvx_once(const int ncol, iwork.data(), ifail.data(), &info); - if (info) { - throw std::runtime_error("info=" + ModuleBase::GlobalFunc::TO_STRING(info) + ". " - + std::string(__FILE__) + " line " - + std::to_string(__LINE__)); + if (info) + { + throw std::runtime_error( + "info=" + ModuleBase::GlobalFunc::TO_STRING(info) + ". " + std::string(__FILE__) + " line " + + std::to_string(__LINE__)); } // Query returned optimal lwork in work[0] @@ -287,92 +352,119 @@ std::pair> DiagoLapack::zhegvx_once(const int ncol, #else zhegvx_( #endif - &itype, - &jobz, - &range, - &uplo, - &n, - h_tmp.c, - &lda, - s_tmp.c, - &ldb, - &vl, - &vu, - &il, - &iu, - &abstol, - &M, - ekb, - wfc_2d.get_pointer(), - &ldz, - work.data(), - &lwork, - rwork.data(), - iwork.data(), - ifail.data(), - &info); + &itype, + &jobz, + &range, + &uplo, + &n, + h_tmp.c, + &lda, + s_tmp.c, + &ldb, + &vl, + &vu, + &il, + &iu, + &abstol, + &M, + ekb, + wfc_2d.get_pointer(), + &ldz, + work.data(), + &lwork, + rwork.data(), + iwork.data(), + ifail.data(), + &info); - if (info == 0) { + if (info == 0) + { return std::make_pair(info, std::vector{}); - } else if (info < 0) { + } + else if (info < 0) + { return std::make_pair(info, std::vector{}); - } else if (info % 2) { + } + else if (info % 2) + { return std::make_pair(info, ifail); - } else if (info / 2 % 2) { + } + else if (info / 2 % 2) + { return std::make_pair(info, iclustr); - } else if (info / 4 % 2) { + } + else if (info / 4 % 2) + { return std::make_pair(info, std::vector{M, NZ}); - } else if (info / 16 % 2) { + } + else if (info / 16 % 2) + { return std::make_pair(info, ifail); - } else { - throw std::runtime_error("info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" - + std::string(__FILE__) + " line " - + std::to_string(__LINE__)); + } + else + { + throw std::runtime_error( + "info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" + std::string(__FILE__) + " line " + + std::to_string(__LINE__)); } } template -void DiagoLapack::dsygvx_diag(const int ncol, - const int nrow, - const double* const h_mat, - const double* const s_mat, - double* const ekb, - psi::Psi& wfc_2d) +void DiagoLapack::dsygvx_diag( + const int ncol, + const int nrow, + const double *const h_mat, + const double *const s_mat, + double *const ekb, + psi::Psi &wfc_2d) { while (true) { - const std::pair> info_vec = dsygvx_once(ncol, nrow, h_mat, s_mat, ekb, wfc_2d); + const std::pair> info_vec = dsygvx_once( + ncol, + nrow, + h_mat, + s_mat, + ekb, + wfc_2d); post_processing(info_vec.first, info_vec.second); - if (info_vec.first == 0) { + if (info_vec.first == 0) + { break; } } } template -void DiagoLapack::zhegvx_diag(const int ncol, - const int nrow, - const std::complex* const h_mat, - const std::complex* const s_mat, - double* const ekb, - psi::Psi>& wfc_2d) +void DiagoLapack::zhegvx_diag( + const int ncol, + const int nrow, + const std::complex *const h_mat, + const std::complex *const s_mat, + double *const ekb, + psi::Psi> &wfc_2d) { while (true) { - const std::pair> info_vec = zhegvx_once(ncol, nrow, h_mat, s_mat, ekb, wfc_2d); + const std::pair> info_vec = zhegvx_once( + ncol, + nrow, + h_mat, + s_mat, + ekb, + wfc_2d); post_processing(info_vec.first, info_vec.second); - if (info_vec.first == 0) { + if (info_vec.first == 0) + { break; } } } -template -void DiagoLapack::post_processing(const int info, const std::vector& vec) +template void DiagoLapack::post_processing(const int info, const std::vector &vec) { const std::string str_info = "info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n"; - const std::string str_FILE - = std::string(__FILE__) + " line " + std::to_string(__LINE__) + ".\n"; + const std::string str_FILE = std::string(__FILE__) + " line " + std::to_string(__LINE__) + ".\n"; const std::string str_info_FILE = str_info + str_FILE; if (info == 0) @@ -382,17 +474,17 @@ void DiagoLapack::post_processing(const int info, const std::vector& vec else if (info < 0) { const int info_negative = -info; - const std::string str_index - = (info_negative > 100) - ? ModuleBase::GlobalFunc::TO_STRING(info_negative / 100) + "-th argument " - + ModuleBase::GlobalFunc::TO_STRING(info_negative % 100) + "-entry is illegal.\n" - : ModuleBase::GlobalFunc::TO_STRING(info_negative) + "-th argument is illegal.\n"; + const std::string str_index = + (info_negative > 100) ? ModuleBase::GlobalFunc::TO_STRING(info_negative / 100) + "-th argument " + + ModuleBase::GlobalFunc::TO_STRING(info_negative % 100) + "-entry is illegal.\n" + : ModuleBase::GlobalFunc::TO_STRING(info_negative) + "-th argument is illegal.\n"; throw std::runtime_error(str_info_FILE + str_index); } else if (info % 2) { std::string str_ifail = "ifail = "; - for (const int i: vec) { + for (const int i : vec) + { str_ifail += ModuleBase::GlobalFunc::TO_STRING(i) + " "; } throw std::runtime_error(str_info_FILE + str_ifail); @@ -401,14 +493,15 @@ void DiagoLapack::post_processing(const int info, const std::vector& vec { int degeneracy_need = 0; // `vec` is iclustr, holding one [begin, end] pair per process - for (std::size_t irank = 0; 2 * irank + 1 < vec.size(); ++irank) { + for (std::size_t irank = 0; 2 * irank + 1 < vec.size(); ++irank) + { degeneracy_need = std::max(degeneracy_need, vec[2 * irank + 1] - vec[2 * irank]); } if (degeneracy_need <= this->degeneracy_max) { throw std::runtime_error( - str_info_FILE + "degeneracy_need = " + ModuleBase::GlobalFunc::TO_STRING(degeneracy_need) + ".\n" - + "degeneracy_saved = " + ModuleBase::GlobalFunc::TO_STRING(this->degeneracy_max) + ".\n"); + str_info_FILE + "degeneracy_need = " + ModuleBase::GlobalFunc::TO_STRING(degeneracy_need) + ".\n" + + "degeneracy_saved = " + ModuleBase::GlobalFunc::TO_STRING(this->degeneracy_max) + ".\n"); } else { diff --git a/source/source_hsolver/diago_lapack.h b/source/source_hsolver/diago_lapack.h index 75f98dfc0c0..49ea3e5aeed 100644 --- a/source/source_hsolver/diago_lapack.h +++ b/source/source_hsolver/diago_lapack.h @@ -8,9 +8,9 @@ #ifndef DIAGOLAPACK_H #define DIAGOLAPACK_H -#include "source_base/macros.h" // GetRealType -#include "source_base/matrix_block.h" +#include "source_base/macros.h" // GetRealType #include "source_base/matrix.h" +#include "source_base/matrix_block.h" #include "source_basis/module_ao/parallel_orbitals.h" #include "source_psi/psi.h" @@ -20,8 +20,7 @@ namespace hsolver { -template -class DiagoLapack +template class DiagoLapack { private: using Real = typename GetTypeReal::type; @@ -31,44 +30,54 @@ class DiagoLapack /// @param nbands_in number of lowest eigenpairs to compute DiagoLapack(const int nlocal_in, const int nbands_in) : nlocal(nlocal_in), nbands(nbands_in) {}; - void diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in); - #ifdef __MPI + void diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in); +#ifdef __MPI // diagnolization used in parallel-k case - void diag_pool(ModuleBase::MatrixBlock& h_mat, ModuleBase::MatrixBlock& s_mat, psi::Psi& psi, Real* eigenvalue_in, MPI_Comm& comm); + void diag_pool( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in, + MPI_Comm &comm); #endif - void dsygvx_diag(const int ncol, - const int nrow, - const double* const h_mat, - const double* const s_mat, - double* const ekb, - psi::Psi& wfc_2d); - void zhegvx_diag(const int ncol, - const int nrow, - const std::complex* const h_mat, - const std::complex* const s_mat, - double* const ekb, - psi::Psi>& wfc_2d); + void dsygvx_diag( + const int ncol, + const int nrow, + const double *const h_mat, + const double *const s_mat, + double *const ekb, + psi::Psi &wfc_2d); + void zhegvx_diag( + const int ncol, + const int nrow, + const std::complex *const h_mat, + const std::complex *const s_mat, + double *const ekb, + psi::Psi> &wfc_2d); - std::pair> dsygvx_once(const int ncol, - const int nrow, - const double* const h_mat, - const double* const s_mat, - double* const ekb, - psi::Psi& wfc_2d) const; - std::pair> zhegvx_once(const int ncol, - const int nrow, - const std::complex* const h_mat, - const std::complex* const s_mat, - double* const ekb, - psi::Psi>& wfc_2d) const; + std::pair> dsygvx_once( + const int ncol, + const int nrow, + const double *const h_mat, + const double *const s_mat, + double *const ekb, + psi::Psi &wfc_2d) const; + std::pair> zhegvx_once( + const int ncol, + const int nrow, + const std::complex *const h_mat, + const std::complex *const s_mat, + double *const ekb, + psi::Psi> &wfc_2d) const; int degeneracy_max = 12; // For reorthogonalized memory. 12 followes siesta. - void post_processing(const int info, const std::vector& vec); + void post_processing(const int info, const std::vector &vec); private: const int nlocal; diff --git a/source/source_hsolver/diago_params.cpp b/source/source_hsolver/diago_params.cpp index 28e1040a974..40cf006db53 100644 --- a/source/source_hsolver/diago_params.cpp +++ b/source/source_hsolver/diago_params.cpp @@ -1,14 +1,16 @@ #include "diago_params.h" + #include "diago_iter_assist.h" namespace hsolver { template -void setup_diago_params_pw(const int istep, - const int iter, - const double ethr, - const Input_para& inp) +void setup_diago_params_pw( + const int istep, + const int iter, + const double ethr, + const Input_para &inp) { /// choose if psi should be diag in subspace /// be careful that istep start from 0 and iter start from 1 @@ -23,10 +25,11 @@ void setup_diago_params_pw(const int istep, } template -void setup_diago_params_sdft(const int istep, - const int iter, - const double ethr, - const Input_para& inp) +void setup_diago_params_sdft( + const int istep, + const int iter, + const double ethr, + const Input_para &inp) { /// choose if psi should be diag in subspace /// be careful that istep start from 0 and iter start from 1 @@ -44,63 +47,75 @@ void setup_diago_params_sdft(const int istep, } /// Template instantiation for CPU -template void setup_diago_params_pw, base_device::DEVICE_CPU>(const int istep, - const int iter, - const double ethr, - const Input_para& inp); -template void setup_diago_params_pw, base_device::DEVICE_CPU>(const int istep, - const int iter, - const double ethr, - const Input_para& inp); -template void setup_diago_params_pw(const int istep, - const int iter, - const double ethr, - const Input_para& inp); +template void setup_diago_params_pw, base_device::DEVICE_CPU>( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); +template void setup_diago_params_pw, base_device::DEVICE_CPU>( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); +template void setup_diago_params_pw( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); /// Template instantiation for GPU #if ((defined __CUDA) || (defined __ROCM)) -template void setup_diago_params_pw, base_device::DEVICE_GPU>(const int istep, - const int iter, - const double ethr, - const Input_para& inp); -template void setup_diago_params_pw, base_device::DEVICE_GPU>(const int istep, - const int iter, - const double ethr, - const Input_para& inp); -template void setup_diago_params_pw(const int istep, - const int iter, - const double ethr, - const Input_para& inp); +template void setup_diago_params_pw, base_device::DEVICE_GPU>( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); +template void setup_diago_params_pw, base_device::DEVICE_GPU>( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); +template void setup_diago_params_pw( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); #endif /// Template instantiation for SDFT CPU -template void setup_diago_params_sdft, base_device::DEVICE_CPU>(const int istep, - const int iter, - const double ethr, - const Input_para& inp); -template void setup_diago_params_sdft, base_device::DEVICE_CPU>(const int istep, - const int iter, - const double ethr, - const Input_para& inp); -template void setup_diago_params_sdft(const int istep, - const int iter, - const double ethr, - const Input_para& inp); +template void setup_diago_params_sdft, base_device::DEVICE_CPU>( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); +template void setup_diago_params_sdft, base_device::DEVICE_CPU>( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); +template void setup_diago_params_sdft( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); /// Template instantiation for SDFT GPU #if ((defined __CUDA) || (defined __ROCM)) -template void setup_diago_params_sdft, base_device::DEVICE_GPU>(const int istep, - const int iter, - const double ethr, - const Input_para& inp); -template void setup_diago_params_sdft, base_device::DEVICE_GPU>(const int istep, - const int iter, - const double ethr, - const Input_para& inp); -template void setup_diago_params_sdft(const int istep, - const int iter, - const double ethr, - const Input_para& inp); +template void setup_diago_params_sdft, base_device::DEVICE_GPU>( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); +template void setup_diago_params_sdft, base_device::DEVICE_GPU>( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); +template void setup_diago_params_sdft( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); #endif } // namespace hsolver diff --git a/source/source_hsolver/diago_params.h b/source/source_hsolver/diago_params.h index 5d46b010466..21ceb6fcb31 100644 --- a/source/source_hsolver/diago_params.h +++ b/source/source_hsolver/diago_params.h @@ -19,10 +19,11 @@ namespace hsolver * @param inp Input parameters */ template -void setup_diago_params_pw(const int istep, - const int iter, - const double ethr, - const Input_para& inp); +void setup_diago_params_pw( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); /** * @brief Setup diagonalization parameters for SDFT method @@ -37,10 +38,11 @@ void setup_diago_params_pw(const int istep, * @param inp Input parameters */ template -void setup_diago_params_sdft(const int istep, - const int iter, - const double ethr, - const Input_para& inp); +void setup_diago_params_sdft( + const int istep, + const int iter, + const double ethr, + const Input_para &inp); } // namespace hsolver diff --git a/source/source_hsolver/diago_pexsi.cpp b/source/source_hsolver/diago_pexsi.cpp index 6e310f8e2c9..9f68b146ba4 100644 --- a/source/source_hsolver/diago_pexsi.cpp +++ b/source/source_hsolver/diago_pexsi.cpp @@ -1,25 +1,25 @@ -#include #include #include +#include #ifdef __PEXSI #include "diago_pexsi.h" +#include "module_pexsi/pexsi_solver.h" #include "source_base/matrix_block.h" -#include "source_base/tool_title.h" #include "source_base/tool_quit.h" +#include "source_base/tool_title.h" #include "source_basis/module_ao/parallel_orbitals.h" -#include "module_pexsi/pexsi_solver.h" namespace hsolver { -template -std::vector DiagoPexsi::mu_buffer; +template std::vector DiagoPexsi::mu_buffer; template -DiagoPexsi::DiagoPexsi(const Parallel_Orbitals* ParaV_in, - const int nspin_in, - const int nlocal_in, - const double nelec_in, - const int world_nproc_in) +DiagoPexsi::DiagoPexsi( + const Parallel_Orbitals *ParaV_in, + const int nspin_in, + const int nlocal_in, + const double nelec_in, + const int world_nproc_in) { this->nspin_dm = (nspin_in == 4) ? 1 : nspin_in; this->nlocal = nlocal_in; @@ -42,38 +42,37 @@ DiagoPexsi::DiagoPexsi(const Parallel_Orbitals* ParaV_in, this->DM[i] = new T[ParaV->nrow * ParaV->ncol]; this->EDM[i] = new T[ParaV->nrow * ParaV->ncol]; } - } -template -DiagoPexsi::~DiagoPexsi() +template DiagoPexsi::~DiagoPexsi() { for (int i = 0; i < this->nspin_dm; i++) { delete[] this->DM[i]; delete[] this->EDM[i]; } - } template <> -void DiagoPexsi::diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - double* eigenvalue_in) +void DiagoPexsi::diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + double *eigenvalue_in) { ModuleBase::TITLE("DiagoPEXSI", "diag"); int ik = psi.get_current_k(); - this->ps->prepare(this->ParaV->blacs_ctxt, - this->ParaV->nb, - this->ParaV->nrow, - this->ParaV->ncol, - this->nlocal, - this->nelec, - h_mat.p, - s_mat.p, - DM[ik], - EDM[ik]); + this->ps->prepare( + this->ParaV->blacs_ctxt, + this->ParaV->nb, + this->ParaV->nrow, + this->ParaV->ncol, + this->nlocal, + this->nelec, + h_mat.p, + s_mat.p, + DM[ik], + EDM[ik]); this->ps->solve(mu_buffer[ik], this->world_nproc); this->totalFreeEnergy = this->ps->get_totalFreeEnergy(); this->totalEnergyH = this->ps->get_totalEnergyH(); @@ -82,17 +81,18 @@ void DiagoPexsi::diag(ModuleBase::MatrixBlock& h_mat, } template <> -void DiagoPexsi>::diag(ModuleBase::MatrixBlock>& h_mat, - ModuleBase::MatrixBlock>& s_mat, - psi::Psi>& psi, - double* eigenvalue_in) +void DiagoPexsi>::diag( + ModuleBase::MatrixBlock> &h_mat, + ModuleBase::MatrixBlock> &s_mat, + psi::Psi> &psi, + double *eigenvalue_in) { ModuleBase::TITLE("DiagoPEXSI", "diag"); ModuleBase::WARNING_QUIT("DiagoPEXSI", "PEXSI is not completed for multi-k case"); } template class DiagoPexsi; -template class DiagoPexsi >; +template class DiagoPexsi>; } // namespace hsolver #endif diff --git a/source/source_hsolver/diago_pexsi.h b/source/source_hsolver/diago_pexsi.h index 953f2a06571..16b07cb19c8 100644 --- a/source/source_hsolver/diago_pexsi.h +++ b/source/source_hsolver/diago_pexsi.h @@ -1,37 +1,39 @@ #ifndef DIGAOPEXSI_H #define DIGAOPEXSI_H -#include -#include -#include "source_base/macros.h" // GetRealType +#include "module_pexsi/pexsi_solver.h" +#include "source_base/macros.h" // GetRealType #include "source_base/matrix_block.h" #include "source_basis/module_ao/parallel_orbitals.h" #include "source_psi/psi.h" -#include "module_pexsi/pexsi_solver.h" + +#include +#include namespace hsolver { -template -class DiagoPexsi +template class DiagoPexsi { private: using Real = typename GetTypeReal::type; static std::vector mu_buffer; public: - DiagoPexsi(const Parallel_Orbitals* ParaV_in, - const int nspin_in, - const int nlocal_in, - const double nelec_in, - const int world_nproc_in); - void diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in); - const Parallel_Orbitals* ParaV = nullptr; - std::vector DM; - std::vector EDM; + DiagoPexsi( + const Parallel_Orbitals *ParaV_in, + const int nspin_in, + const int nlocal_in, + const double nelec_in, + const int world_nproc_in); + void diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in); + const Parallel_Orbitals *ParaV = nullptr; + std::vector DM; + std::vector EDM; double totalEnergyH; double totalEnergyS; double totalFreeEnergy; diff --git a/source/source_hsolver/diago_pxxxgvx.cpp b/source/source_hsolver/diago_pxxxgvx.cpp index e9ee523be0a..9adc701b0c7 100644 --- a/source/source_hsolver/diago_pxxxgvx.cpp +++ b/source/source_hsolver/diago_pxxxgvx.cpp @@ -33,7 +33,7 @@ namespace hsolver /** * @brief Wrapper function for Scalapack's generalized eigensolver routines. - * + * * @param itype Specifies the problem type to be solved. * @param jobz Specifies whether to compute eigenvectors. * @param range Specifies the range of eigenvalues to be found. @@ -70,316 +70,325 @@ namespace hsolver * @param iclustr The array to store the indices of the eigenvalue clusters. * @param gap The array to store the gaps between eigenvalue clusters. * @param info Output status of the computation. - * + * * @note for a uniform interface, rwork and lrwork are input arguments, but not used in pdsygvx_/pssygvx_ */ -void pxxxgvx_(const int* itype, - const char* jobz, - const char* range, - const char* uplo, - const int* n, - double* A, - const int* ia, - const int* ja, - const int* desca, - double* B, - const int* ib, - const int* jb, - const int* descb, - const double* vl, - const double* vu, - const int* il, - const int* iu, - const double* abstol, - int* m, - int* nz, - double* w, - const double* orfac, - double* Z, - const int* iz, - const int* jz, - const int* descz, - double* work, - int* lwork, - double* rwork, - int* lrwork, - int* iwork, - int* liwork, - int* ifail, - int* iclustr, - double* gap, - int* info) +void pxxxgvx_( + const int *itype, + const char *jobz, + const char *range, + const char *uplo, + const int *n, + double *A, + const int *ia, + const int *ja, + const int *desca, + double *B, + const int *ib, + const int *jb, + const int *descb, + const double *vl, + const double *vu, + const int *il, + const int *iu, + const double *abstol, + int *m, + int *nz, + double *w, + const double *orfac, + double *Z, + const int *iz, + const int *jz, + const int *descz, + double *work, + int *lwork, + double *rwork, + int *lrwork, + int *iwork, + int *liwork, + int *ifail, + int *iclustr, + double *gap, + int *info) { // double - pdsygvx_(itype, - jobz, - range, - uplo, - n, - A, - ia, - ja, - desca, - B, - ib, - jb, - descb, - vl, - vu, - il, - iu, - abstol, - m, - nz, - w, - orfac, - Z, - iz, - jz, - descz, - work, - lwork, - iwork, - liwork, - ifail, - iclustr, - gap, - info); + pdsygvx_( + itype, + jobz, + range, + uplo, + n, + A, + ia, + ja, + desca, + B, + ib, + jb, + descb, + vl, + vu, + il, + iu, + abstol, + m, + nz, + w, + orfac, + Z, + iz, + jz, + descz, + work, + lwork, + iwork, + liwork, + ifail, + iclustr, + gap, + info); } -void pxxxgvx_(const int* itype, - const char* jobz, - const char* range, - const char* uplo, - const int* n, - std::complex* A, - const int* ia, - const int* ja, - const int* desca, - std::complex* B, - const int* ib, - const int* jb, - const int* descb, - const double* vl, - const double* vu, - const int* il, - const int* iu, - const double* abstol, - int* m, - int* nz, - double* w, - const double* orfac, - std::complex* Z, - const int* iz, - const int* jz, - const int* descz, - std::complex* work, - int* lwork, - double* rwork, - int* lrwork, - int* iwork, - int* liwork, - int* ifail, - int* iclustr, - double* gap, - int* info) +void pxxxgvx_( + const int *itype, + const char *jobz, + const char *range, + const char *uplo, + const int *n, + std::complex *A, + const int *ia, + const int *ja, + const int *desca, + std::complex *B, + const int *ib, + const int *jb, + const int *descb, + const double *vl, + const double *vu, + const int *il, + const int *iu, + const double *abstol, + int *m, + int *nz, + double *w, + const double *orfac, + std::complex *Z, + const int *iz, + const int *jz, + const int *descz, + std::complex *work, + int *lwork, + double *rwork, + int *lrwork, + int *iwork, + int *liwork, + int *ifail, + int *iclustr, + double *gap, + int *info) { // std::complex - pzhegvx_(itype, - jobz, - range, - uplo, - n, - A, - ia, - ja, - desca, - B, - ib, - jb, - descb, - vl, - vu, - il, - iu, - abstol, - m, - nz, - w, - orfac, - Z, - iz, - jz, - descz, - work, - lwork, - rwork, - lrwork, - iwork, - liwork, - ifail, - iclustr, - gap, - info); + pzhegvx_( + itype, + jobz, + range, + uplo, + n, + A, + ia, + ja, + desca, + B, + ib, + jb, + descb, + vl, + vu, + il, + iu, + abstol, + m, + nz, + w, + orfac, + Z, + iz, + jz, + descz, + work, + lwork, + rwork, + lrwork, + iwork, + liwork, + ifail, + iclustr, + gap, + info); } -void pxxxgvx_(const int* itype, - const char* jobz, - const char* range, - const char* uplo, - const int* n, - float* A, - const int* ia, - const int* ja, - const int* desca, - float* B, - const int* ib, - const int* jb, - const int* descb, - const float* vl, - const float* vu, - const int* il, - const int* iu, - const float* abstol, - int* m, - int* nz, - float* w, - const float* orfac, - float* Z, - const int* iz, - const int* jz, - const int* descz, - float* work, - int* lwork, - float* rwork, - int* lrwork, - int* iwork, - int* liwork, - int* ifail, - int* iclustr, - float* gap, - int* info) +void pxxxgvx_( + const int *itype, + const char *jobz, + const char *range, + const char *uplo, + const int *n, + float *A, + const int *ia, + const int *ja, + const int *desca, + float *B, + const int *ib, + const int *jb, + const int *descb, + const float *vl, + const float *vu, + const int *il, + const int *iu, + const float *abstol, + int *m, + int *nz, + float *w, + const float *orfac, + float *Z, + const int *iz, + const int *jz, + const int *descz, + float *work, + int *lwork, + float *rwork, + int *lrwork, + int *iwork, + int *liwork, + int *ifail, + int *iclustr, + float *gap, + int *info) { // float - pssygvx_(itype, - jobz, - range, - uplo, - n, - A, - ia, - ja, - desca, - B, - ib, - jb, - descb, - vl, - vu, - il, - iu, - abstol, - m, - nz, - w, - orfac, - Z, - iz, - jz, - descz, - work, - lwork, - iwork, - liwork, - ifail, - iclustr, - gap, - info); + pssygvx_( + itype, + jobz, + range, + uplo, + n, + A, + ia, + ja, + desca, + B, + ib, + jb, + descb, + vl, + vu, + il, + iu, + abstol, + m, + nz, + w, + orfac, + Z, + iz, + jz, + descz, + work, + lwork, + iwork, + liwork, + ifail, + iclustr, + gap, + info); } -void pxxxgvx_(const int* itype, - const char* jobz, - const char* range, - const char* uplo, - const int* n, - std::complex* A, - const int* ia, - const int* ja, - const int* desca, - std::complex* B, - const int* ib, - const int* jb, - const int* descb, - const float* vl, - const float* vu, - const int* il, - const int* iu, - const float* abstol, - int* m, - int* nz, - float* w, - const float* orfac, - std::complex* Z, - const int* iz, - const int* jz, - const int* descz, - std::complex* work, - int* lwork, - float* rwork, - int* lrwork, - int* iwork, - int* liwork, - int* ifail, - int* iclustr, - float* gap, - int* info) +void pxxxgvx_( + const int *itype, + const char *jobz, + const char *range, + const char *uplo, + const int *n, + std::complex *A, + const int *ia, + const int *ja, + const int *desca, + std::complex *B, + const int *ib, + const int *jb, + const int *descb, + const float *vl, + const float *vu, + const int *il, + const int *iu, + const float *abstol, + int *m, + int *nz, + float *w, + const float *orfac, + std::complex *Z, + const int *iz, + const int *jz, + const int *descz, + std::complex *work, + int *lwork, + float *rwork, + int *lrwork, + int *iwork, + int *liwork, + int *ifail, + int *iclustr, + float *gap, + int *info) { // std::complex - pchegvx_(itype, - jobz, - range, - uplo, - n, - A, - ia, - ja, - desca, - B, - ib, - jb, - descb, - vl, - vu, - il, - iu, - abstol, - m, - nz, - w, - orfac, - Z, - iz, - jz, - descz, - work, - lwork, - rwork, - lrwork, - iwork, - liwork, - ifail, - iclustr, - gap, - info); + pchegvx_( + itype, + jobz, + range, + uplo, + n, + A, + ia, + ja, + desca, + B, + ib, + jb, + descb, + vl, + vu, + il, + iu, + abstol, + m, + nz, + w, + orfac, + Z, + iz, + jz, + descz, + work, + lwork, + rwork, + lrwork, + iwork, + liwork, + ifail, + iclustr, + gap, + info); } -void pxxxgvx_post_processing(const int info, - const std::vector& ifail, - const std::vector& iclustr, - const int M, - const int NZ, - const int nbands, - int& degeneracy_max) +void pxxxgvx_post_processing( + const int info, + const std::vector &ifail, + const std::vector &iclustr, + const int M, + const int NZ, + const int nbands, + int °eneracy_max) { const std::string str_info = "Scalapack diagonalization: \n info = " + std::to_string(info) + ".\n"; @@ -391,15 +400,15 @@ void pxxxgvx_post_processing(const int info, { const int info_negative = -info; const std::string str_index = (info_negative > 100) - ? std::to_string(info_negative / 100) + "-th argument " - + std::to_string(info_negative % 100) + "-entry is illegal.\n" + ? std::to_string(info_negative / 100) + "-th argument " + + std::to_string(info_negative % 100) + "-entry is illegal.\n" : std::to_string(info_negative) + "-th argument is illegal.\n"; throw std::runtime_error(str_info + str_index); } else if (info % 2) { std::string str_ifail = "ifail = "; - for (const int i: ifail) + for (const int i : ifail) { str_ifail += std::to_string(i) + " "; } @@ -443,41 +452,35 @@ void pxxxgvx_post_processing(const int info, } } -void get_lwork(int& lwork, std::vector& work) -{ - lwork = work[0]; -} +void get_lwork(int &lwork, std::vector &work) { lwork = work[0]; } -void get_lwork(int& lwork, std::vector& work) -{ - lwork = work[0]; -} +void get_lwork(int &lwork, std::vector &work) { lwork = work[0]; } -void get_lwork(int& lwork, std::vector>& work) -{ - lwork = work[0].real(); -} +void get_lwork(int &lwork, std::vector> &work) { lwork = work[0].real(); } -void get_lwork(int& lwork, std::vector>& work) -{ - lwork = work[0].real(); -} +void get_lwork(int &lwork, std::vector> &work) { lwork = work[0].real(); } template -void pxxxgvx_diag(const int* const desc, - const int ncol, - const int nrow, - const int nbands, - const T* const h_mat, - const T* const s_mat, - typename GetTypeReal::type* const ekb, - T* const wfc_2d) +void pxxxgvx_diag( + const int *const desc, + const int ncol, + const int nrow, + const int nbands, + const T *const h_mat, + const T *const s_mat, + typename GetTypeReal::type *const ekb, + T *const wfc_2d) { int nprow = 1; int npcol = 1; int myprow = 0; int mypcol = 0; - Cblacs_gridinfo(desc[1], &nprow, &npcol, &myprow, &mypcol); + Cblacs_gridinfo( + desc[1], + &nprow, + &npcol, + &myprow, + &mypcol); int dsize = nprow * npcol; int degeneracy_max = 12; // only used for complex and std::complex @@ -513,42 +516,43 @@ void pxxxgvx_diag(const int* const desc, std::vector iclustr(2 * dsize); std::vector::type> gap(dsize); - pxxxgvx_(&itype, - &jobz, - &range, - &uplo, - &ndim_global, - h_tmp.data(), - &one, - &one, - desc, - s_tmp.data(), - &one, - &one, - desc, - &vl, - &vu, - &il, - &iu, - &abstol, - &M, - &NZ, - ekb, - &orfac, - wfc_2d, - &one, - &one, - desc, - work.data(), - &lwork, // is not used for real data type - rwork.data(), // is not used for real data type - &lrwork, - iwork.data(), - &liwork, - ifail.data(), - iclustr.data(), - gap.data(), - &info); + pxxxgvx_( + &itype, + &jobz, + &range, + &uplo, + &ndim_global, + h_tmp.data(), + &one, + &one, + desc, + s_tmp.data(), + &one, + &one, + desc, + &vl, + &vu, + &il, + &iu, + &abstol, + &M, + &NZ, + ekb, + &orfac, + wfc_2d, + &one, + &one, + desc, + work.data(), + &lwork, // is not used for real data type + rwork.data(), // is not used for real data type + &lrwork, + iwork.data(), + &liwork, + ifail.data(), + iclustr.data(), + gap.data(), + &info); if (info) { @@ -573,48 +577,56 @@ void pxxxgvx_diag(const int* const desc, iwork.resize(liwork, 0); } - pxxxgvx_(&itype, - &jobz, - &range, - &uplo, - &ndim_global, - h_tmp.data(), - &one, - &one, - desc, - s_tmp.data(), - &one, - &one, - desc, - &vl, - &vu, - &il, - &iu, - &abstol, - &M, - &NZ, - ekb, - &orfac, - wfc_2d, - &one, - &one, - desc, - work.data(), - &lwork, - rwork.data(), // is not used for real data type - &lrwork, // is not used for real data type - iwork.data(), - &liwork, - ifail.data(), - iclustr.data(), - gap.data(), - &info); + pxxxgvx_( + &itype, + &jobz, + &range, + &uplo, + &ndim_global, + h_tmp.data(), + &one, + &one, + desc, + s_tmp.data(), + &one, + &one, + desc, + &vl, + &vu, + &il, + &iu, + &abstol, + &M, + &NZ, + ekb, + &orfac, + wfc_2d, + &one, + &one, + desc, + work.data(), + &lwork, + rwork.data(), // is not used for real data type + &lrwork, // is not used for real data type + iwork.data(), + &liwork, + ifail.data(), + iclustr.data(), + gap.data(), + &info); if (info == 0) { return; } - pxxxgvx_post_processing(info, ifail, iclustr, M, NZ, nbands, degeneracy_max); + pxxxgvx_post_processing( + info, + ifail, + iclustr, + M, + NZ, + nbands, + degeneracy_max); // break the loop for real data type if (std::is_same::value || std::is_same::value) @@ -626,44 +638,48 @@ void pxxxgvx_diag(const int* const desc, // template instantiation // double -template void pxxxgvx_diag(const int* const desc, - const int ncol, - const int nrow, - const int nbands, - const double* const h_mat, - const double* const s_mat, - double* const ekb, - double* const wfc_2d); +template void pxxxgvx_diag( + const int *const desc, + const int ncol, + const int nrow, + const int nbands, + const double *const h_mat, + const double *const s_mat, + double *const ekb, + double *const wfc_2d); -// std::complex -template void pxxxgvx_diag(const int* const desc, - const int ncol, - const int nrow, - const int nbands, - const std::complex* const h_mat, - const std::complex* const s_mat, - double* const ekb, - std::complex* const wfc_2d); +// std::complex +template void pxxxgvx_diag( + const int *const desc, + const int ncol, + const int nrow, + const int nbands, + const std::complex *const h_mat, + const std::complex *const s_mat, + double *const ekb, + std::complex *const wfc_2d); // float -template void pxxxgvx_diag(const int* const desc, - const int ncol, - const int nrow, - const int nbands, - const float* const h_mat, - const float* const s_mat, - float* const ekb, - float* const wfc_2d); +template void pxxxgvx_diag( + const int *const desc, + const int ncol, + const int nrow, + const int nbands, + const float *const h_mat, + const float *const s_mat, + float *const ekb, + float *const wfc_2d); // std::complex -template void pxxxgvx_diag(const int* const desc, - const int ncol, - const int nrow, - const int nbands, - const std::complex* const h_mat, - const std::complex* const s_mat, - float* const ekb, - std::complex* const wfc_2d); +template void pxxxgvx_diag( + const int *const desc, + const int ncol, + const int nrow, + const int nbands, + const std::complex *const h_mat, + const std::complex *const s_mat, + float *const ekb, + std::complex *const wfc_2d); #endif diff --git a/source/source_hsolver/diago_pxxxgvx.h b/source/source_hsolver/diago_pxxxgvx.h index 4b1e8486a0e..a7d8fb0c211 100644 --- a/source/source_hsolver/diago_pxxxgvx.h +++ b/source/source_hsolver/diago_pxxxgvx.h @@ -1,15 +1,16 @@ #ifndef HSOLVER_DIAGO_PXXXGVX_H #define HSOLVER_DIAGO_PXXXGVX_H -#include #include "source_base/macros.h" +#include + namespace hsolver { #ifdef __MPI /** * @brief Wrapper function for Scalapack's generalized eigensolver routines: pdsygvx_, pzhegvx_, pssygvx_, pchegvx_. - * + * * @param desc the descriptor of scalapack descriptor * @param ncol the number of columns of the H/S matrix in current processor * @param nrow the number of rows of the H/S matrix in current processor @@ -18,19 +19,20 @@ namespace hsolver * @param s_mat the overlap matrix * @param ekb the eigenvalues * @param wfc_2d the eigenvectors in 2D block cyclic distribution - * + * */ template -void pxxxgvx_diag(const int* const desc, - const int ncol, - const int nrow, - const int nbands, - const T* const h_mat, - const T* const s_mat, - typename GetTypeReal::type* const ekb, - T* const wfc_2d); -#endif +void pxxxgvx_diag( + const int *const desc, + const int ncol, + const int nrow, + const int nbands, + const T *const h_mat, + const T *const s_mat, + typename GetTypeReal::type *const ekb, + T *const wfc_2d); +#endif } // namespace hsolver diff --git a/source/source_hsolver/diago_scalapack.cpp b/source/source_hsolver/diago_scalapack.cpp index 644fb83f9a0..1e8c4e7a7b2 100644 --- a/source/source_hsolver/diago_scalapack.cpp +++ b/source/source_hsolver/diago_scalapack.cpp @@ -21,80 +21,138 @@ namespace { /// Number of processes in the BLACS process grid that `desc` lives on. /// p?sygvx requires iclustr(2*NPROW*NPCOL) and gap(NPROW*NPCOL). -int blacs_grid_size(const int* const desc) +int blacs_grid_size(const int *const desc) { int nprow = 0, npcol = 0, myprow = 0, mypcol = 0; - Cblacs_gridinfo(desc[1], &nprow, &npcol, &myprow, &mypcol); + Cblacs_gridinfo( + desc[1], + &nprow, + &npcol, + &myprow, + &mypcol); return nprow * npcol; } } // namespace - template<> - void DiagoScalapack::diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in) +template <> +void DiagoScalapack::diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in) { ModuleBase::TITLE("DiagoScalapack", "diag"); assert(h_mat.col == s_mat.col && h_mat.row == s_mat.row && h_mat.desc == s_mat.desc); std::vector eigen(this->nlocal, 0.0); - this->pdsygvx_diag(h_mat.desc, h_mat.col, h_mat.row, h_mat.p, s_mat.p, eigen.data(), psi); + this->pdsygvx_diag( + h_mat.desc, + h_mat.col, + h_mat.row, + h_mat.p, + s_mat.p, + eigen.data(), + psi); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); } - template<> - void DiagoScalapack>::diag(ModuleBase::MatrixBlock>& h_mat, - ModuleBase::MatrixBlock>& s_mat, - psi::Psi>& psi, - Real* eigenvalue_in) +template <> +void DiagoScalapack>::diag( + ModuleBase::MatrixBlock> &h_mat, + ModuleBase::MatrixBlock> &s_mat, + psi::Psi> &psi, + Real *eigenvalue_in) { ModuleBase::TITLE("DiagoScalapack", "diag"); assert(h_mat.col == s_mat.col && h_mat.row == s_mat.row && h_mat.desc == s_mat.desc); std::vector eigen(this->nlocal, 0.0); - this->pzhegvx_diag(h_mat.desc, h_mat.col, h_mat.row, h_mat.p, s_mat.p, eigen.data(), psi); + this->pzhegvx_diag( + h_mat.desc, + h_mat.col, + h_mat.row, + h_mat.p, + s_mat.p, + eigen.data(), + psi); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); } #ifdef __MPI - template<> - void DiagoScalapack::diag_pool(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in, - MPI_Comm& comm) +template <> +void DiagoScalapack::diag_pool( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in, + MPI_Comm &comm) { ModuleBase::TITLE("DiagoScalapack", "diag_pool"); assert(h_mat.col == s_mat.col && h_mat.row == s_mat.row && h_mat.desc == s_mat.desc); std::vector eigen(this->nlocal, 0.0); - this->pdsygvx_diag(h_mat.desc, h_mat.col, h_mat.row, h_mat.p, s_mat.p, eigen.data(), psi); + this->pdsygvx_diag( + h_mat.desc, + h_mat.col, + h_mat.row, + h_mat.p, + s_mat.p, + eigen.data(), + psi); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); } - template<> - void DiagoScalapack>::diag_pool(ModuleBase::MatrixBlock>& h_mat, - ModuleBase::MatrixBlock>& s_mat, - psi::Psi>& psi, - Real* eigenvalue_in, - MPI_Comm& comm) +template <> +void DiagoScalapack>::diag_pool( + ModuleBase::MatrixBlock> &h_mat, + ModuleBase::MatrixBlock> &s_mat, + psi::Psi> &psi, + Real *eigenvalue_in, + MPI_Comm &comm) { ModuleBase::TITLE("DiagoScalapack", "diag_pool"); assert(h_mat.col == s_mat.col && h_mat.row == s_mat.row && h_mat.desc == s_mat.desc); std::vector eigen(this->nlocal, 0.0); - this->pzhegvx_diag(h_mat.desc, h_mat.col, h_mat.row, h_mat.p, s_mat.p, eigen.data(), psi); + this->pzhegvx_diag( + h_mat.desc, + h_mat.col, + h_mat.row, + h_mat.p, + s_mat.p, + eigen.data(), + psi); const int inc = 1; - BlasConnector::copy(this->nbands, eigen.data(), inc, eigenvalue_in, inc); + BlasConnector::copy( + this->nbands, + eigen.data(), + inc, + eigenvalue_in, + inc); } #endif - template - std::pair> DiagoScalapack::pdsygvx_once(const int* const desc, - const int ncol, - const int nrow, - const double *const h_mat, - const double *const s_mat, - double *const ekb, - psi::Psi &wfc_2d) const +template +std::pair> DiagoScalapack::pdsygvx_once( + const int *const desc, + const int ncol, + const int nrow, + const double *const h_mat, + const double *const s_mat, + double *const ekb, + psi::Psi &wfc_2d) const { ModuleBase::matrix h_tmp(ncol, nrow, false); memcpy(h_tmp.c, h_mat, sizeof(double) * ncol * nrow); @@ -113,112 +171,129 @@ int blacs_grid_size(const int* const desc) std::vector iclustr(2 * ngrid); std::vector gap(ngrid); - pdsygvx_(&itype, - &jobz, - &range, - &uplo, - &this->nlocal, - h_tmp.c, - &one, - &one, - desc, - s_tmp.c, - &one, - &one, - desc, - &vl, - &vu, - &il, - &iu, - &abstol, - &M, - &NZ, - ekb, - &orfac, - wfc_2d.get_pointer(), - &one, - &one, - desc, - work.data(), - &lwork, - iwork.data(), - &liwork, - ifail.data(), - iclustr.data(), - gap.data(), - &info); - if (info) { - throw std::runtime_error("info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" - + std::string(__FILE__) + " line " - + std::to_string(__LINE__)); -} + pdsygvx_( + &itype, + &jobz, + &range, + &uplo, + &this->nlocal, + h_tmp.c, + &one, + &one, + desc, + s_tmp.c, + &one, + &one, + desc, + &vl, + &vu, + &il, + &iu, + &abstol, + &M, + &NZ, + ekb, + &orfac, + wfc_2d.get_pointer(), + &one, + &one, + desc, + work.data(), + &lwork, + iwork.data(), + &liwork, + ifail.data(), + iclustr.data(), + gap.data(), + &info); + if (info) + { + throw std::runtime_error( + "info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" + std::string(__FILE__) + " line " + + std::to_string(__LINE__)); + } lwork = work[0]; - work.resize(std::max(lwork,3), 0); + work.resize(std::max(lwork, 3), 0); liwork = iwork[0]; iwork.resize(liwork, 0); - pdsygvx_(&itype, - &jobz, - &range, - &uplo, - &this->nlocal, - h_tmp.c, - &one, - &one, - desc, - s_tmp.c, - &one, - &one, - desc, - &vl, - &vu, - &il, - &iu, - &abstol, - &M, - &NZ, - ekb, - &orfac, - wfc_2d.get_pointer(), - &one, - &one, - desc, - work.data(), - &lwork, - iwork.data(), - &liwork, - ifail.data(), - iclustr.data(), - gap.data(), - &info); + pdsygvx_( + &itype, + &jobz, + &range, + &uplo, + &this->nlocal, + h_tmp.c, + &one, + &one, + desc, + s_tmp.c, + &one, + &one, + desc, + &vl, + &vu, + &il, + &iu, + &abstol, + &M, + &NZ, + ekb, + &orfac, + wfc_2d.get_pointer(), + &one, + &one, + desc, + work.data(), + &lwork, + iwork.data(), + &liwork, + ifail.data(), + iclustr.data(), + gap.data(), + &info); - if (info == 0) { + if (info == 0) + { return std::make_pair(info, std::vector{}); - } else if (info < 0) { + } + else if (info < 0) + { return std::make_pair(info, std::vector{}); - } else if (info % 2) { + } + else if (info % 2) + { return std::make_pair(info, ifail); - } else if (info / 2 % 2) { + } + else if (info / 2 % 2) + { return std::make_pair(info, iclustr); - } else if (info / 4 % 2) { + } + else if (info / 4 % 2) + { return std::make_pair(info, std::vector{M, NZ}); - } else if (info / 16 % 2) { + } + else if (info / 16 % 2) + { return std::make_pair(info, ifail); - } else { - throw std::runtime_error("info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" - + std::string(__FILE__) + " line " - + std::to_string(__LINE__)); -} + } + else + { + throw std::runtime_error( + "info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" + std::string(__FILE__) + " line " + + std::to_string(__LINE__)); + } } - template - std::pair> DiagoScalapack::pzhegvx_once(const int* const desc, - const int ncol, - const int nrow, - const std::complex *const h_mat, - const std::complex *const s_mat, - double *const ekb, - psi::Psi> &wfc_2d) const +template +std::pair> DiagoScalapack::pzhegvx_once( + const int *const desc, + const int ncol, + const int nrow, + const std::complex *const h_mat, + const std::complex *const s_mat, + double *const ekb, + psi::Psi> &wfc_2d) const { ModuleBase::ComplexMatrix h_tmp(ncol, nrow, false); memcpy(h_tmp.c, h_mat, sizeof(std::complex) * ncol * nrow); @@ -229,9 +304,9 @@ int blacs_grid_size(const int* const desc) const int itype = 1, il = 1, iu = this->nbands, one = 1; int M = 0, NZ = 0, lwork = -1, lrwork = -1, liwork = -1, info = 0; const double abstol = SCALAPACK_ABSTOL, orfac = SCALAPACK_ORFAC; - //Note: pzhegvx_ has a bug - // We must give vl,vu a value, although we do not use range 'V' - // We must give rwork at least a memory of sizeof(double) * 3 + // Note: pzhegvx_ has a bug + // We must give vl,vu a value, although we do not use range 'V' + // We must give rwork at least a memory of sizeof(double) * 3 const double vl = 0, vu = 0; std::vector> work(1, 0); std::vector rwork(3, 0); @@ -241,155 +316,187 @@ int blacs_grid_size(const int* const desc) std::vector iclustr(2 * ngrid); std::vector gap(ngrid); - pzhegvx_(&itype, - &jobz, - &range, - &uplo, - &this->nlocal, - h_tmp.c, - &one, - &one, - desc, - s_tmp.c, - &one, - &one, - desc, - &vl, - &vu, - &il, - &iu, - &abstol, - &M, - &NZ, - ekb, - &orfac, - wfc_2d.get_pointer(), - &one, - &one, - desc, - work.data(), - &lwork, - rwork.data(), - &lrwork, - iwork.data(), - &liwork, - ifail.data(), - iclustr.data(), - gap.data(), - &info); - if (info) { - throw std::runtime_error("info=" + ModuleBase::GlobalFunc::TO_STRING(info) + ". " - + std::string(__FILE__) + " line " - + std::to_string(__LINE__)); -} + pzhegvx_( + &itype, + &jobz, + &range, + &uplo, + &this->nlocal, + h_tmp.c, + &one, + &one, + desc, + s_tmp.c, + &one, + &one, + desc, + &vl, + &vu, + &il, + &iu, + &abstol, + &M, + &NZ, + ekb, + &orfac, + wfc_2d.get_pointer(), + &one, + &one, + desc, + work.data(), + &lwork, + rwork.data(), + &lrwork, + iwork.data(), + &liwork, + ifail.data(), + iclustr.data(), + gap.data(), + &info); + if (info) + { + throw std::runtime_error( + "info=" + ModuleBase::GlobalFunc::TO_STRING(info) + ". " + std::string(__FILE__) + " line " + + std::to_string(__LINE__)); + } lwork = work[0].real(); work.resize(lwork, 0); lrwork = rwork[0] + this->degeneracy_max * this->nlocal; - int maxlrwork = std::max(lrwork,3); + int maxlrwork = std::max(lrwork, 3); rwork.resize(maxlrwork, 0); liwork = iwork[0]; iwork.resize(liwork, 0); - pzhegvx_(&itype, - &jobz, - &range, - &uplo, - &this->nlocal, - h_tmp.c, - &one, - &one, - desc, - s_tmp.c, - &one, - &one, - desc, - &vl, - &vu, - &il, - &iu, - &abstol, - &M, - &NZ, - ekb, - &orfac, - wfc_2d.get_pointer(), - &one, - &one, - desc, - work.data(), - &lwork, - rwork.data(), - &lrwork, - iwork.data(), - &liwork, - ifail.data(), - iclustr.data(), - gap.data(), - &info); + pzhegvx_( + &itype, + &jobz, + &range, + &uplo, + &this->nlocal, + h_tmp.c, + &one, + &one, + desc, + s_tmp.c, + &one, + &one, + desc, + &vl, + &vu, + &il, + &iu, + &abstol, + &M, + &NZ, + ekb, + &orfac, + wfc_2d.get_pointer(), + &one, + &one, + desc, + work.data(), + &lwork, + rwork.data(), + &lrwork, + iwork.data(), + &liwork, + ifail.data(), + iclustr.data(), + gap.data(), + &info); - if (info == 0) { + if (info == 0) + { return std::make_pair(info, std::vector{}); - } else if (info < 0) { + } + else if (info < 0) + { return std::make_pair(info, std::vector{}); - } else if (info % 2) { + } + else if (info % 2) + { return std::make_pair(info, ifail); - } else if (info / 2 % 2) { + } + else if (info / 2 % 2) + { return std::make_pair(info, iclustr); - } else if (info / 4 % 2) { + } + else if (info / 4 % 2) + { return std::make_pair(info, std::vector{M, NZ}); - } else if (info / 16 % 2) { + } + else if (info / 16 % 2) + { return std::make_pair(info, ifail); - } else { - throw std::runtime_error("info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" - + std::string(__FILE__) + " line " - + std::to_string(__LINE__)); -} + } + else + { + throw std::runtime_error( + "info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n" + std::string(__FILE__) + " line " + + std::to_string(__LINE__)); + } } - template - void DiagoScalapack::pdsygvx_diag(const int* const desc, - const int ncol, - const int nrow, - const double *const h_mat, - const double *const s_mat, - double *const ekb, - psi::Psi &wfc_2d) +template +void DiagoScalapack::pdsygvx_diag( + const int *const desc, + const int ncol, + const int nrow, + const double *const h_mat, + const double *const s_mat, + double *const ekb, + psi::Psi &wfc_2d) { while (true) { - const std::pair> info_vec = pdsygvx_once(desc, ncol, nrow, h_mat, s_mat, ekb, wfc_2d); + const std::pair> info_vec = pdsygvx_once( + desc, + ncol, + nrow, + h_mat, + s_mat, + ekb, + wfc_2d); post_processing(info_vec.first, info_vec.second); - if (info_vec.first == 0) { + if (info_vec.first == 0) + { break; -} + } } } - template - void DiagoScalapack ::pzhegvx_diag(const int* const desc, - const int ncol, - const int nrow, - const std::complex *const h_mat, - const std::complex *const s_mat, - double *const ekb, - psi::Psi> &wfc_2d) +template +void DiagoScalapack::pzhegvx_diag( + const int *const desc, + const int ncol, + const int nrow, + const std::complex *const h_mat, + const std::complex *const s_mat, + double *const ekb, + psi::Psi> &wfc_2d) { while (true) { - const std::pair> info_vec = pzhegvx_once(desc, ncol, nrow, h_mat, s_mat, ekb, wfc_2d); + const std::pair> info_vec = pzhegvx_once( + desc, + ncol, + nrow, + h_mat, + s_mat, + ekb, + wfc_2d); post_processing(info_vec.first, info_vec.second); - if (info_vec.first == 0) { + if (info_vec.first == 0) + { break; -} + } } } - template - void DiagoScalapack::post_processing(const int info, const std::vector& vec) +template void DiagoScalapack::post_processing(const int info, const std::vector &vec) { const std::string str_info = "info = " + ModuleBase::GlobalFunc::TO_STRING(info) + ".\n"; - const std::string str_FILE - = std::string(__FILE__) + " line " + std::to_string(__LINE__) + ".\n"; + const std::string str_FILE = std::string(__FILE__) + " line " + std::to_string(__LINE__) + ".\n"; const std::string str_info_FILE = str_info + str_FILE; if (info == 0) @@ -399,33 +506,34 @@ int blacs_grid_size(const int* const desc) else if (info < 0) { const int info_negative = -info; - const std::string str_index - = (info_negative > 100) - ? ModuleBase::GlobalFunc::TO_STRING(info_negative / 100) + "-th argument " - + ModuleBase::GlobalFunc::TO_STRING(info_negative % 100) + "-entry is illegal.\n" - : ModuleBase::GlobalFunc::TO_STRING(info_negative) + "-th argument is illegal.\n"; + const std::string str_index = + (info_negative > 100) ? ModuleBase::GlobalFunc::TO_STRING(info_negative / 100) + "-th argument " + + ModuleBase::GlobalFunc::TO_STRING(info_negative % 100) + "-entry is illegal.\n" + : ModuleBase::GlobalFunc::TO_STRING(info_negative) + "-th argument is illegal.\n"; throw std::runtime_error(str_info_FILE + str_index); } else if (info % 2) { std::string str_ifail = "ifail = "; - for (const int i: vec) { + for (const int i : vec) + { str_ifail += ModuleBase::GlobalFunc::TO_STRING(i) + " "; -} + } throw std::runtime_error(str_info_FILE + str_ifail); } else if (info / 2 % 2) { int degeneracy_need = 0; // `vec` is iclustr, sized 2*NPROW*NPCOL by the caller - for (std::size_t irank = 0; 2 * irank + 1 < vec.size(); ++irank) { + for (std::size_t irank = 0; 2 * irank + 1 < vec.size(); ++irank) + { degeneracy_need = std::max(degeneracy_need, vec[2 * irank + 1] - vec[2 * irank]); } if (degeneracy_need <= this->degeneracy_max) { throw std::runtime_error( - str_info_FILE + "degeneracy_need = " + ModuleBase::GlobalFunc::TO_STRING(degeneracy_need) + ".\n" - + "degeneracy_saved = " + ModuleBase::GlobalFunc::TO_STRING(this->degeneracy_max) + ".\n"); + str_info_FILE + "degeneracy_need = " + ModuleBase::GlobalFunc::TO_STRING(degeneracy_need) + ".\n" + + "degeneracy_saved = " + ModuleBase::GlobalFunc::TO_STRING(this->degeneracy_max) + ".\n"); } else { @@ -437,8 +545,7 @@ int blacs_grid_size(const int* const desc) { const std::string str_M = "M = " + ModuleBase::GlobalFunc::TO_STRING(vec[0]) + ".\n"; const std::string str_NZ = "NZ = " + ModuleBase::GlobalFunc::TO_STRING(vec[1]) + ".\n"; - const std::string str_NBANDS - = "nbands = " + ModuleBase::GlobalFunc::TO_STRING(this->nbands) + ".\n"; + const std::string str_NBANDS = "nbands = " + ModuleBase::GlobalFunc::TO_STRING(this->nbands) + ".\n"; throw std::runtime_error(str_info_FILE + str_M + str_NZ + str_NBANDS); } else if (info / 16 % 2) diff --git a/source/source_hsolver/diago_scalapack.h b/source/source_hsolver/diago_scalapack.h index eb05928dc29..7877065b814 100644 --- a/source/source_hsolver/diago_scalapack.h +++ b/source/source_hsolver/diago_scalapack.h @@ -8,68 +8,78 @@ #ifndef DIAGO_SCALAPACK_H #define DIAGO_SCALAPACK_H -#include -#include -#include - -#include "source_base/macros.h" // GetRealType -#include "source_base/matrix_block.h" -#include "source_psi/psi.h" #include "source_base/complexmatrix.h" +#include "source_base/macros.h" // GetRealType #include "source_base/matrix.h" +#include "source_base/matrix_block.h" #include "source_basis/module_ao/parallel_orbitals.h" +#include "source_psi/psi.h" + +#include +#include +#include namespace hsolver { - template - class DiagoScalapack +template class DiagoScalapack { -private: + private: using Real = typename GetTypeReal::type; + public: /// @param nlocal_in global dimension of the NAO Hamiltonian /// @param nbands_in number of lowest eigenpairs to compute DiagoScalapack(const int nlocal_in, const int nbands_in) : nlocal(nlocal_in), nbands(nbands_in) {}; - void diag(ModuleBase::MatrixBlock& h_mat, - ModuleBase::MatrixBlock& s_mat, - psi::Psi& psi, - Real* eigenvalue_in); + void diag( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in); #ifdef __MPI // diagnolization used in parallel-k case - void diag_pool(ModuleBase::MatrixBlock& h_mat, ModuleBase::MatrixBlock& s_mat, psi::Psi& psi, Real* eigenvalue_in, MPI_Comm& comm); + void diag_pool( + ModuleBase::MatrixBlock &h_mat, + ModuleBase::MatrixBlock &s_mat, + psi::Psi &psi, + Real *eigenvalue_in, + MPI_Comm &comm); #endif private: - void pdsygvx_diag(const int *const desc, - const int ncol, - const int nrow, - const double *const h_mat, - const double *const s_mat, - double *const ekb, - psi::Psi &wfc_2d); - void pzhegvx_diag(const int *const desc, - const int ncol, - const int nrow, - const std::complex *const h_mat, - const std::complex *const s_mat, - double *const ekb, - psi::Psi> &wfc_2d); + void pdsygvx_diag( + const int *const desc, + const int ncol, + const int nrow, + const double *const h_mat, + const double *const s_mat, + double *const ekb, + psi::Psi &wfc_2d); + void pzhegvx_diag( + const int *const desc, + const int ncol, + const int nrow, + const std::complex *const h_mat, + const std::complex *const s_mat, + double *const ekb, + psi::Psi> &wfc_2d); - std::pair> pdsygvx_once(const int *const desc, - const int ncol, - const int nrow, - const double *const h_mat, - const double *const s_mat, - double *const ekb, - psi::Psi &wfc_2d) const; - std::pair> pzhegvx_once(const int *const desc, - const int ncol, - const int nrow, - const std::complex *const h_mat, - const std::complex *const s_mat, - double *const ekb, - psi::Psi> &wfc_2d) const; + std::pair> pdsygvx_once( + const int *const desc, + const int ncol, + const int nrow, + const double *const h_mat, + const double *const s_mat, + double *const ekb, + psi::Psi &wfc_2d) const; + std::pair> pzhegvx_once( + const int *const desc, + const int ncol, + const int nrow, + const std::complex *const h_mat, + const std::complex *const s_mat, + double *const ekb, + psi::Psi> &wfc_2d) const; int degeneracy_max = 12; // For reorthogonalized memory. 12 followes siesta. diff --git a/source/source_hsolver/hs_matrix.h b/source/source_hsolver/hs_matrix.h index 46591bf71bc..0bf468f0387 100644 --- a/source/source_hsolver/hs_matrix.h +++ b/source/source_hsolver/hs_matrix.h @@ -14,15 +14,14 @@ namespace hsolver * (hamilt::HamiltHSMatrix in source_hamilt/hamilt_hs_adapter.h); HSolverLCAO * and Parallel_K2D only ever ask it for the two matrix views. */ -template -class HSMatrix +template class HSMatrix { public: virtual ~HSMatrix() = default; /// H(k) and S(k) of k point ik as non-owning views. The memory they point /// to belongs to the implementer and stays valid until the next call. - virtual void hs_at_k(const int ik, ModuleBase::MatrixBlock& hk, ModuleBase::MatrixBlock& sk) = 0; + virtual void hs_at_k(const int ik, ModuleBase::MatrixBlock &hk, ModuleBase::MatrixBlock &sk) = 0; }; } // namespace hsolver diff --git a/source/source_hsolver/hs_operator.h b/source/source_hsolver/hs_operator.h index 30547b2336b..9c75210d89e 100644 --- a/source/source_hsolver/hs_operator.h +++ b/source/source_hsolver/hs_operator.h @@ -19,8 +19,7 @@ namespace hsolver * Block vectors are column major: vector i occupies x[i*ld, i*ld + ld). Only * the first npw rows carry data, ld is the leading dimension. */ -template -class HSOperator +template class HSOperator { public: virtual ~HSOperator() = default; @@ -29,24 +28,28 @@ class HSOperator virtual void update_k(const int ik) = 0; /// hx[:, 0:nvec) = H * x[:, 0:nvec), both with leading dimension ld - virtual void hpsi(const T* x, T* hx, const int ld, const int nvec) const = 0; + virtual void hpsi( + const T *x, + T *hx, + const int ld, + const int nvec) const = 0; /// sx[:, 0:nvec) = S * x[:, 0:nvec), both with leading dimension ld - virtual void spsi(const T* x, T* sx, const int ld, const int nvec) const = 0; + virtual void spsi( + const T *x, + T *sx, + const int ld, + const int nvec) const = 0; /// Hook used by DiagoIterAssist::diag_subspace_init: hcc is the n*n /// subspace Hamiltonian (column major, ld n). A Hamiltonian carrying a /// term that hpsi() does not cover (EXX in lcao_in_pw) adds it here. - virtual void add_to_subspace_h(T* hcc, const int n) const - { - } + virtual void add_to_subspace_h(T *hcc, const int n) const {} /// Hook used by DiagoIterAssist::diag_subspace_init: vcc is the n*nband /// matrix of subspace eigenvectors, handed out right after the subspace /// diagonalization. - virtual void export_subspace_vec(const T* vcc, const int n, const int nband) const - { - } + virtual void export_subspace_vec(const T *vcc, const int n, const int nband) const {} }; } // namespace hsolver diff --git a/source/source_hsolver/hsolver.cpp b/source/source_hsolver/hsolver.cpp index 66e3db65c25..6c23c128929 100644 --- a/source/source_hsolver/hsolver.cpp +++ b/source/source_hsolver/hsolver.cpp @@ -8,18 +8,19 @@ namespace hsolver { -double set_diagethr_ks(const std::string basis_type, - const std::string esolver_type, - const std::string calculation_in, - const std::string init_chg_in, - const std::string precision_flag_in, - const int istep, - const int iter, - const double drho, - const double pw_diag_thr_init, - const double diag_ethr_in, - const double nelec_in, - const double scf_thr_in) +double set_diagethr_ks( + const std::string basis_type, + const std::string esolver_type, + const std::string calculation_in, + const std::string init_chg_in, + const std::string precision_flag_in, + const int istep, + const int iter, + const double drho, + const double pw_diag_thr_init, + const double diag_ethr_in, + const double nelec_in, + const double scf_thr_in) { double res_diag_ethr = diag_ethr_in; @@ -68,8 +69,9 @@ double set_diagethr_ks(const std::string basis_type, { res_diag_ethr = 1.e-2; } - res_diag_ethr = std::min(res_diag_ethr, - static_cast(0.1) * drho / std::max(static_cast(1.0), static_cast(nelec_in))); + res_diag_ethr = std::min( + res_diag_ethr, + static_cast(0.1) * drho / std::max(static_cast(1.0), static_cast(nelec_in))); } // Keep the single-precision diagonalization threshold at or above the floor @@ -87,19 +89,20 @@ double set_diagethr_ks(const std::string basis_type, return res_diag_ethr; } -double set_diagethr_sdft(const std::string basis_type, - const std::string esolver_type, - const std::string calculation_in, - const std::string init_chg_in, - const int istep, - const int iter, - const double drho, - const double pw_diag_thr_init, - const double diag_ethr_in, - const int nband_in, - const double stoiter_ks_ne_in, - const double nelec_in, - const double scf_thr_in) +double set_diagethr_sdft( + const std::string basis_type, + const std::string esolver_type, + const std::string calculation_in, + const std::string init_chg_in, + const int istep, + const int iter, + const double drho, + const double pw_diag_thr_init, + const double diag_ethr_in, + const int nband_in, + const double stoiter_ks_ne_in, + const double nelec_in, + const double scf_thr_in) { double res_diag_ethr = diag_ethr_in; @@ -144,14 +147,15 @@ double set_diagethr_sdft(const std::string basis_type, return res_diag_ethr; } -double reset_diag_ethr(std::ofstream& ofs_running, - const std::string basis_type, - const std::string esolver_type, - const std::string precision_flag_in, - const double hsover_error, - const double drho_in, - const double diag_ethr_in, - const double nelec_in) +double reset_diag_ethr( + std::ofstream &ofs_running, + const std::string basis_type, + const std::string esolver_type, + const std::string precision_flag_in, + const double hsover_error, + const double drho_in, + const double diag_ethr_in, + const double nelec_in) { double new_diag_ethr = 0.0; @@ -183,7 +187,11 @@ double reset_diag_ethr(std::ofstream& ofs_running, return new_diag_ethr; }; -double cal_hsolve_error(const std::string basis_type, const std::string esolver_type, const double diag_ethr_in, const double nelec_in) +double cal_hsolve_error( + const std::string basis_type, + const std::string esolver_type, + const double diag_ethr_in, + const double nelec_in) { if (basis_type == "pw" && esolver_type == "ksdft") { diff --git a/source/source_hsolver/hsolver.h b/source/source_hsolver/hsolver.h index e2dc3c7fbcf..ba16d2a07ca 100644 --- a/source/source_hsolver/hsolver.h +++ b/source/source_hsolver/hsolver.h @@ -7,49 +7,52 @@ namespace hsolver { -double set_diagethr_ks(const std::string basis_type, - const std::string esolver_type, - const std::string calculation_in, - const std::string init_chg_in, - const std::string precision_flag_in, - const int istep, - const int iter, - const double drho, - const double pw_diag_thr_init, - const double diag_ethr_in, - const double nelec_in, - const double scf_thr_in); - -double set_diagethr_sdft(const std::string basis_type, - const std::string esolver_type, - const std::string calculation_in, - const std::string init_chg_in, - const int istep, - const int iter, - const double drho, - const double pw_diag_thr_init, - const double diag_ethr_in, - const int nband_in, - const double stoiter_ks_ne_in, - const double nelec_in, - const double scf_thr_in); - +double set_diagethr_ks( + const std::string basis_type, + const std::string esolver_type, + const std::string calculation_in, + const std::string init_chg_in, + const std::string precision_flag_in, + const int istep, + const int iter, + const double drho, + const double pw_diag_thr_init, + const double diag_ethr_in, + const double nelec_in, + const double scf_thr_in); + +double set_diagethr_sdft( + const std::string basis_type, + const std::string esolver_type, + const std::string calculation_in, + const std::string init_chg_in, + const int istep, + const int iter, + const double drho, + const double pw_diag_thr_init, + const double diag_ethr_in, + const int nband_in, + const double stoiter_ks_ne_in, + const double nelec_in, + const double scf_thr_in); // reset diagethr according to drho and hsolver_error -double reset_diag_ethr(std::ofstream& ofs_running, - const std::string basis_type, - const std::string esolver_type, - const std::string precision_flag_in, - const double hsover_error, - const double drho_in, - const double diag_ethr_in, - const double nelec_in); +double reset_diag_ethr( + std::ofstream &ofs_running, + const std::string basis_type, + const std::string esolver_type, + const std::string precision_flag_in, + const double hsover_error, + const double drho_in, + const double diag_ethr_in, + const double nelec_in); // calculate hsolver_error (for sdft, lcao and lcao-in-pw, we suppose the error is zero) -double cal_hsolve_error(const std::string basis_type, - const std::string esolver_type, - const double diag_ethr_in, - const double nelec_in); +double cal_hsolve_error( + const std::string basis_type, + const std::string esolver_type, + const double diag_ethr_in, + const double nelec_in); } // namespace hsolver #endif \ No newline at end of file diff --git a/source/source_hsolver/hsolver_lcao.cpp b/source/source_hsolver/hsolver_lcao.cpp index bef8e0befad..fe6284d184d 100644 --- a/source/source_hsolver/hsolver_lcao.cpp +++ b/source/source_hsolver/hsolver_lcao.cpp @@ -26,47 +26,55 @@ #include "diago_pexsi.h" #endif -#include "source_base/module_device/device.h" -#include "source_estate/elecstate_tools.h" #include "source_base/memory_recorder.h" +#include "source_base/module_device/device.h" #include "source_base/timer.h" #include "source_estate/elecstate_lcao.h" -#include "source_estate/module_dm/dm_from_psi.h" +#include "source_estate/elecstate_tools.h" #include "source_estate/module_dm/density_matrix.h" +#include "source_estate/module_dm/dm_from_psi.h" #include "source_hsolver/parallel_k2d.h" namespace hsolver { template -void HSolverLCAO::solve(HSMatrix& hs, - psi::Psi& psi, - elecstate::ElecState* pes, - module_dm::DensityMatrix& dm, // mohan add 2025-11-03 - Charge &chr, - const int nspin, - const double omega, - const bool skip_charge) +void HSolverLCAO::solve( + HSMatrix &hs, + psi::Psi &psi, + elecstate::ElecState *pes, + module_dm::DensityMatrix &dm, // mohan add 2025-11-03 + Charge &chr, + const int nspin, + const double omega, + const bool skip_charge) { ModuleBase::TITLE("HSolverLCAO", "solve"); ModuleBase::timer::start("HSolverLCAO", "solve"); if (this->method != "pexsi") { - #ifdef __MPI - #ifdef __CUDA +#ifdef __MPI +#ifdef __CUDA if (this->method == "cusolver" && this->world_nproc > 1) { this->parakSolve_cusolver(hs, psi, pes); - }else - #endif - if (this->kpar_lcao > 1 - && (this->method == "genelpa" || this->method == "elpa" || this->method == "scalapack_gvx" || this->method == "lapack")) + } + else +#endif + if (this->kpar_lcao > 1 && (this->method == "genelpa" || this->method == "elpa" || + this->method == "scalapack_gvx" || this->method == "lapack")) { - this->parakSolve(hs, psi, pes, this->kpar_lcao, nspin); - } else - #endif - if (this->kpar_lcao == 1) + this->parakSolve( + hs, + psi, + pes, + this->kpar_lcao, + nspin); + } + else +#endif + if (this->kpar_lcao == 1) { /// Loop over k points for solve Hamiltonian to eigenpairs(eigenvalues and eigenvectors). for (int ik = 0; ik < psi.get_nk(); ++ik) @@ -79,26 +87,35 @@ void HSolverLCAO::solve(HSMatrix& hs, psi.fix_k(ik); /// solve eigenvector and eigenvalue for H(k) - this->hamiltSolvePsiK(hk, sk, psi, &(pes->ekb(ik, 0))); + this->hamiltSolvePsiK( + hk, + sk, + psi, + &(pes->ekb(ik, 0))); } } else { - ModuleBase::WARNING_QUIT("HSolverLCAO::solve", - "This method and KPAR setting is not supported for lcao basis in ABACUS!"); + ModuleBase::WARNING_QUIT( + "HSolverLCAO::solve", "This method and KPAR setting is not supported for lcao basis in ABACUS!"); } - elecstate::calculate_weights(pes->ekb, - pes->wg, - pes->klist, - pes->eferm, - pes->f_en, - pes->nelec_spin, - this->nbands, - pes->skip_weights); + elecstate::calculate_weights( + pes->ekb, + pes->wg, + pes->klist, + pes->eferm, + pes->f_en, + pes->nelec_spin, + this->nbands, + pes->skip_weights); elecstate::calEBand(pes->ekb, pes->wg, pes->f_en); - module_dm::dm_from_psi(this->ParaV, pes->wg, psi, dm); + module_dm::dm_from_psi( + this->ParaV, + pes->wg, + psi, + dm); dm.cal_dmr(-1); if (!skip_charge) @@ -106,7 +123,11 @@ void HSolverLCAO::solve(HSMatrix& hs, // compute charge density from density matrix, mohan update 20251024 // delegate to ElecStateLCAO to keep the source_lcao dependency out of // source_hsolver (mirrors the pexsi branch below and the PW psiToRho path) - dynamic_cast*>(pes)->dmToRho(dm.get_dmr_vec(), nspin, &chr, omega); + dynamic_cast *>(pes)->dmToRho( + dm.get_dmr_vec(), + nspin, + &chr, + omega); } else { @@ -116,19 +137,32 @@ void HSolverLCAO::solve(HSMatrix& hs, else if (this->method == "pexsi") { #ifdef __PEXSI // other purification methods should follow this routine - DiagoPexsi pe(ParaV, nspin, this->nlocal, this->nelec, this->world_nproc); + DiagoPexsi pe( + ParaV, + nspin, + this->nlocal, + this->nelec, + this->world_nproc); for (int ik = 0; ik < psi.get_nk(); ++ik) { psi.fix_k(ik); ModuleBase::MatrixBlock hk, sk; hs.hs_at_k(ik, hk, sk); // solve eigenvector and eigenvalue for H(k) - pe.diag(hk, sk, psi, nullptr); + pe.diag( + hk, + sk, + psi, + nullptr); } - auto _pes = dynamic_cast*>(pes); + auto _pes = dynamic_cast *>(pes); pes->f_en.eband = pe.totalFreeEnergy; // maybe eferm could be dealt with in the future - _pes->dm2rho(pe.DM, pe.EDM, &dm, omega); + _pes->dm2rho( + pe.DM, + pe.EDM, + &dm, + omega); #endif } @@ -137,10 +171,11 @@ void HSolverLCAO::solve(HSMatrix& hs, } template -void HSolverLCAO::hamiltSolvePsiK(ModuleBase::MatrixBlock& hk, - ModuleBase::MatrixBlock& sk, - psi::Psi& psi, - double* eigenvalue) +void HSolverLCAO::hamiltSolvePsiK( + ModuleBase::MatrixBlock &hk, + ModuleBase::MatrixBlock &sk, + psi::Psi &psi, + double *eigenvalue) { ModuleBase::TITLE("HSolverLCAO", "hamiltSolvePsiK"); ModuleBase::timer::start("HSolverLCAO", "hamiltSolvePsiK"); @@ -149,19 +184,31 @@ void HSolverLCAO::hamiltSolvePsiK(ModuleBase::MatrixBlock& hk, { #ifdef __MPI DiagoScalapack sa(this->nlocal, this->nbands); - sa.diag(hk, sk, psi, eigenvalue); + sa.diag( + hk, + sk, + psi, + eigenvalue); #endif } #ifdef __ELPA else if (this->method == "genelpa") { DiagoElpa el(this->nlocal, this->nbands); - el.diag(hk, sk, psi, eigenvalue); + el.diag( + hk, + sk, + psi, + eigenvalue); } else if (this->method == "elpa") { DiagoElpaNative el(this->nlocal, this->nbands, this->use_gpu); - el.diag(hk, sk, psi, eigenvalue); + el.diag( + hk, + sk, + psi, + eigenvalue); } #endif #ifdef __CUDA @@ -169,20 +216,32 @@ void HSolverLCAO::hamiltSolvePsiK(ModuleBase::MatrixBlock& hk, { // Note: This branch will only be executed in the single-process case DiagoCusolver cu(this->nlocal, this->nbands); - cu.diag(hk, sk, psi, eigenvalue); + cu.diag( + hk, + sk, + psi, + eigenvalue); } #ifdef __CUSOLVERMP else if (this->method == "cusolvermp") { DiagoCusolverMP cm(this->nlocal, this->nbands); - cm.diag(hk, sk, psi, eigenvalue); + cm.diag( + hk, + sk, + psi, + eigenvalue); } #endif #endif else if (this->method == "lapack") // only for single core { DiagoLapack la(this->nlocal, this->nbands); - la.diag(hk, sk, psi, eigenvalue); + la.diag( + hk, + sk, + psi, + eigenvalue); } else { @@ -193,11 +252,12 @@ void HSolverLCAO::hamiltSolvePsiK(ModuleBase::MatrixBlock& hk, } template -void HSolverLCAO::parakSolve(HSMatrix& hs, - psi::Psi& psi, - elecstate::ElecState* pes, - const int kpar, - const int nspin) +void HSolverLCAO::parakSolve( + HSMatrix &hs, + psi::Psi &psi, + elecstate::ElecState *pes, + const int kpar, + const int nspin) { #ifdef __MPI ModuleBase::timer::start("HSolverLCAO", "parakSolve"); @@ -207,12 +267,22 @@ void HSolverLCAO::parakSolve(HSMatrix& hs, int nks = psi.get_nk(); int nrow = this->ParaV->get_global_row_size(); int nb2d = this->ParaV->get_block_size(); - k2d.set_para_env(psi.get_nk(), nrow, nb2d, this->world_nproc, this->world_rank, nspin); + k2d.set_para_env( + psi.get_nk(), + nrow, + nb2d, + this->world_nproc, + this->world_rank, + nspin); /// set psi_pool const int zero = 0; int coord_col = k2d.get_p2D_pool()->get_coord_col(); - int ncol_bands_pool - = numroc_(&(nbands), &(nb2d), &coord_col, &zero, &(k2d.get_p2D_pool()->dim1)); + int ncol_bands_pool = numroc_( + &(nbands), + &(nb2d), + &coord_col, + &zero, + &(k2d.get_p2D_pool()->dim1)); /// Loop over k points for solve Hamiltonian to charge density for (int ik = 0; ik < k2d.get_pKpoints()->get_max_nks_pool(); ++ik) { @@ -241,47 +311,71 @@ void HSolverLCAO::parakSolve(HSMatrix& hs, k2d.distribute_hsk(hs, ik_kpar, nrow); /// global index of k point int ik_global = ik + k2d.get_pKpoints()->startk_pool[k2d.get_my_pool()]; - auto psi_pool = psi::Psi(1, ncol_bands_pool, k2d.get_p2D_pool()->nrow, k2d.get_p2D_pool()->nrow, true); + auto psi_pool = psi::Psi( + 1, + ncol_bands_pool, + k2d.get_p2D_pool()->nrow, + k2d.get_p2D_pool()->nrow, + true); ModuleBase::Memory::record("HSolverLCAO::psi_pool", nrow * ncol_bands_pool * sizeof(T)); if (ik_global < psi.get_nk() && ik < k2d.get_pKpoints()->nks_pool[k2d.get_my_pool()]) { /// local psi in pool psi_pool.fix_k(0); ModuleBase::MatrixBlock hk_pool = ModuleBase::MatrixBlock{k2d.hk_pool.data(), - (size_t)k2d.get_p2D_pool()->get_row_size(), - (size_t)k2d.get_p2D_pool()->get_col_size(), - k2d.get_p2D_pool()->desc}; + (size_t)k2d.get_p2D_pool()->get_row_size(), (size_t)k2d.get_p2D_pool()->get_col_size(), + k2d.get_p2D_pool()->desc}; ModuleBase::MatrixBlock sk_pool = ModuleBase::MatrixBlock{k2d.sk_pool.data(), - (size_t)k2d.get_p2D_pool()->get_row_size(), - (size_t)k2d.get_p2D_pool()->get_col_size(), - k2d.get_p2D_pool()->desc}; + (size_t)k2d.get_p2D_pool()->get_row_size(), (size_t)k2d.get_p2D_pool()->get_col_size(), + k2d.get_p2D_pool()->desc}; /// solve eigenvector and eigenvalue for H(k) if (this->method == "scalapack_gvx") { DiagoScalapack sa(this->nlocal, this->nbands); - sa.diag_pool(hk_pool, sk_pool, psi_pool, &(pes->ekb(ik_global, 0)), k2d.POOL_WORLD_K2D); + sa.diag_pool( + hk_pool, + sk_pool, + psi_pool, + &(pes->ekb(ik_global, 0)), + k2d.POOL_WORLD_K2D); } else if (this->method == "lapack") { DiagoLapack la(this->nlocal, this->nbands); - la.diag_pool(hk_pool, sk_pool, psi_pool, &(pes->ekb(ik_global, 0)), k2d.POOL_WORLD_K2D); + la.diag_pool( + hk_pool, + sk_pool, + psi_pool, + &(pes->ekb(ik_global, 0)), + k2d.POOL_WORLD_K2D); } #ifdef __ELPA else if (this->method == "genelpa") { DiagoElpa el(this->nlocal, this->nbands); - el.diag_pool(hk_pool, sk_pool, psi_pool, &(pes->ekb(ik_global, 0)), k2d.POOL_WORLD_K2D); + el.diag_pool( + hk_pool, + sk_pool, + psi_pool, + &(pes->ekb(ik_global, 0)), + k2d.POOL_WORLD_K2D); } else if (this->method == "elpa") { DiagoElpaNative el(this->nlocal, this->nbands, this->use_gpu); - el.diag_pool(hk_pool, sk_pool, psi_pool, &(pes->ekb(ik_global, 0)), k2d.POOL_WORLD_K2D); + el.diag_pool( + hk_pool, + sk_pool, + psi_pool, + &(pes->ekb(ik_global, 0)), + k2d.POOL_WORLD_K2D); } #endif else { - ModuleBase::WARNING_QUIT("HSolverLCAO::solve", - "This type of eigensolver for k-parallelism diagnolization is not supported!"); + ModuleBase::WARNING_QUIT( + "HSolverLCAO::solve", + "This type of eigensolver for k-parallelism diagnolization is not supported!"); } } MPI_Barrier(MPI_COMM_WORLD); @@ -289,7 +383,12 @@ void HSolverLCAO::parakSolve(HSMatrix& hs, for (int ipool = 0; ipool < ik_kpar.size(); ++ipool) { int source = k2d.get_pKpoints()->get_startpro_pool(ipool); - MPI_Bcast(&(pes->ekb(ik_kpar[ipool], 0)), nbands, MPI_DOUBLE, source, MPI_COMM_WORLD); + MPI_Bcast( + &(pes->ekb(ik_kpar[ipool], 0)), + nbands, + MPI_DOUBLE, + source, + MPI_COMM_WORLD); int desc_pool[9]; std::copy(k2d.get_p2D_pool()->desc, k2d.get_p2D_pool()->desc + 9, desc_pool); if (k2d.get_my_pool() != ipool) @@ -297,17 +396,18 @@ void HSolverLCAO::parakSolve(HSMatrix& hs, desc_pool[1] = -1; } psi.fix_k(ik_kpar[ipool]); - Cpxgemr2d(nrow, - nbands, - psi_pool.get_pointer(), - 1, - 1, - desc_pool, - psi.get_pointer(), - 1, - 1, - k2d.get_p2D_global()->desc, - k2d.get_p2D_global()->blacs_ctxt); + Cpxgemr2d( + nrow, + nbands, + psi_pool.get_pointer(), + 1, + 1, + desc_pool, + psi.get_pointer(), + 1, + 1, + k2d.get_p2D_global()->desc, + k2d.get_p2D_global()->blacs_ctxt); } MPI_Barrier(MPI_COMM_WORLD); ModuleBase::timer::end("HSolverLCAO", "collect_psi"); @@ -317,15 +417,13 @@ void HSolverLCAO::parakSolve(HSMatrix& hs, #endif } -#if defined (__MPI) && defined (__CUDA) +#if defined(__MPI) && defined(__CUDA) template -void HSolverLCAO::parakSolve_cusolver(HSMatrix& hs, - psi::Psi& psi, - elecstate::ElecState* pes) +void HSolverLCAO::parakSolve_cusolver(HSMatrix &hs, psi::Psi &psi, elecstate::ElecState *pes) { ModuleBase::timer::start("HSolverLCAO", "parakSolve"); // GPU device is already bound by DeviceContext::init() in read_input.cpp - auto& dev_ctx = base_device::DeviceContext::instance(); + auto &dev_ctx = base_device::DeviceContext::instance(); const int local_rank = dev_ctx.get_local_rank(); const int device_count = dev_ctx.get_device_count(); @@ -340,17 +438,34 @@ void HSolverLCAO::parakSolve_cusolver(HSMatrix& hs, std::vector all_local_ranks(world_size); std::vector all_is_active(world_size); - MPI_Allgather(&local_rank, 1, MPI_INT, - all_local_ranks.data(), 1, MPI_INT, MPI_COMM_WORLD); - MPI_Allgather(&is_active_int, 1, MPI_INT, - all_is_active.data(), 1, MPI_INT, MPI_COMM_WORLD); + MPI_Allgather( + &local_rank, + 1, + MPI_INT, + all_local_ranks.data(), + 1, + MPI_INT, + MPI_COMM_WORLD); + MPI_Allgather( + &is_active_int, + 1, + MPI_INT, + all_is_active.data(), + 1, + MPI_INT, + MPI_COMM_WORLD); int total_active_ranks = 0; int max_local_rank = 0; - for (int i = 0; i < world_size; ++i) { - if (all_is_active[i]) { + for (int i = 0; i < world_size; ++i) + { + if (all_is_active[i]) + { total_active_ranks++; - if(all_local_ranks[i] > max_local_rank) max_local_rank = all_local_ranks[i]; + if (all_local_ranks[i] > max_local_rank) + { + max_local_rank = all_local_ranks[i]; + } } } @@ -359,44 +474,64 @@ void HSolverLCAO::parakSolve_cusolver(HSMatrix& hs, // The k-points will be distributed among these ranks in a round-robin fashion. // The purpose of setting the order is to ensure load balancing among nodes as much as possible std::vector active_ranks; - for(int r = 0; r <= max_local_rank; r++) + for (int r = 0; r <= max_local_rank; r++) { - for(int j = 0; j < world_size; j++) + for (int j = 0; j < world_size; j++) { - if(all_is_active[j] && all_local_ranks[j] == r) + if (all_is_active[j] && all_local_ranks[j] == r) { active_ranks.push_back(j); } } } - - const int nks = psi.get_nk(); // total number of k points + + const int nks = psi.get_nk(); // total number of k points const int nbands = this->ParaV->get_nbands(); // Set the parallel storage scheme for the matrix and psi - Parallel_2D mat_para_global; // store the info about how the origin matrix is distributed in parallel - Parallel_2D mat_para_local; // store the info about how the matrix is distributed after collected from all processes - Parallel_2D psi_para_global; // store the info about how the psi is distributed in parallel - Parallel_2D psi_para_local; // store the info about how the psi is distributed before distributing to all processes - - MPI_Comm self_comm; // the communicator that only contains the current process itself - MPI_Comm_split(MPI_COMM_WORLD, world_rank, 0, &self_comm); + Parallel_2D mat_para_global; // store the info about how the origin matrix is distributed in parallel + Parallel_2D mat_para_local; // store the info about how the matrix is distributed after collected from all processes + Parallel_2D psi_para_global; // store the info about how the psi is distributed in parallel + Parallel_2D psi_para_local; // store the info about how the psi is distributed before distributing to all processes + + MPI_Comm self_comm; // the communicator that only contains the current process itself + MPI_Comm_split( + MPI_COMM_WORLD, + world_rank, + 0, + &self_comm); int nrow = this->ParaV->get_global_row_size(); // number of rows in the global matrix int ncol = nrow; - int nb2d = this->ParaV->get_block_size(); // block size for the 2D matrix distribution - mat_para_global.init(nrow, ncol, nb2d, MPI_COMM_WORLD); - psi_para_global.init(nrow, nbands, nb2d, MPI_COMM_WORLD); - mat_para_local.init(nrow, ncol, nb2d, self_comm); - psi_para_local.init(nrow, ncol, nb2d, self_comm); + int nb2d = this->ParaV->get_block_size(); // block size for the 2D matrix distribution + mat_para_global.init( + nrow, + ncol, + nb2d, + MPI_COMM_WORLD); + psi_para_global.init( + nrow, + nbands, + nb2d, + MPI_COMM_WORLD); + mat_para_local.init( + nrow, + ncol, + nb2d, + self_comm); + psi_para_local.init( + nrow, + ncol, + nb2d, + self_comm); std::vector hk_mat; // temporary storage for H(k) matrix collected from all processes std::vector sk_mat; // temporary storage for S(k) matrix collected from all processes // In each iteration, we process total_active_ranks k-points. - for(int ik_start = 0; ik_start < nks; ik_start += total_active_ranks) + for (int ik_start = 0; ik_start < nks; ik_start += total_active_ranks) { int kpt_assigned = -1; // the k-point assigned to the current MPI process in this iteration // Compute and gather the hk and sk matrices distributed across different processes in parallel, // preparing for subsequent transfer to the GPU for computation. - for(int ik = ik_start; ik < ik_start + total_active_ranks && ik < nks; ik++) + for (int ik = ik_start; ik < ik_start + total_active_ranks && ik < nks; ik++) { // `is_active` indicates whether this MPI process is assigned to compute the current k-point bool is_active = world_rank == active_ranks[ik % total_active_ranks]; @@ -409,64 +544,91 @@ void HSolverLCAO::parakSolve_cusolver(HSMatrix& hs, ModuleBase::MatrixBlock hk_2D, sk_2D; hs.hs_at_k(ik, hk_2D, sk_2D); int desc_tmp[9]; - T* hk_local_ptr = hk_mat.data(); - T* sk_local_ptr = sk_mat.data(); + T *hk_local_ptr = hk_mat.data(); + T *sk_local_ptr = sk_mat.data(); std::copy(mat_para_local.desc, mat_para_local.desc + 9, desc_tmp); - if( !is_active) + if (!is_active) { desc_tmp[1] = -1; } - Cpxgemr2d(nrow, ncol, hk_2D.p, 1, 1, mat_para_global.desc, - hk_local_ptr, 1, 1, desc_tmp, - mat_para_global.blacs_ctxt); - Cpxgemr2d(nrow, ncol, sk_2D.p, 1, 1, mat_para_global.desc, - sk_local_ptr, 1, 1, desc_tmp, - mat_para_global.blacs_ctxt); + Cpxgemr2d( + nrow, + ncol, + hk_2D.p, + 1, + 1, + mat_para_global.desc, + hk_local_ptr, + 1, + 1, + desc_tmp, + mat_para_global.blacs_ctxt); + Cpxgemr2d( + nrow, + ncol, + sk_2D.p, + 1, + 1, + mat_para_global.desc, + sk_local_ptr, + 1, + 1, + desc_tmp, + mat_para_global.blacs_ctxt); } // diagonalize the Hamiltonian matrix using cusolver psi::Psi psi_local{}; - if(kpt_assigned != -1) + if (kpt_assigned != -1) { psi_local.resize(1, ncol, nrow); DiagoCusolver cu(this->nlocal, this->nbands); - ModuleBase::MatrixBlock hk_local = ModuleBase::MatrixBlock{ - hk_mat.data(), (size_t)nrow, (size_t)ncol, - mat_para_local.desc}; - ModuleBase::MatrixBlock sk_local = ModuleBase::MatrixBlock{ - sk_mat.data(), (size_t)nrow, (size_t)ncol, - mat_para_local.desc}; - cu.diag(hk_local, sk_local, psi_local, &(pes->ekb(kpt_assigned, 0))); + ModuleBase::MatrixBlock hk_local = ModuleBase::MatrixBlock{hk_mat.data(), (size_t)nrow, (size_t)ncol, + mat_para_local.desc}; + ModuleBase::MatrixBlock sk_local = ModuleBase::MatrixBlock{sk_mat.data(), (size_t)nrow, (size_t)ncol, + mat_para_local.desc}; + cu.diag( + hk_local, + sk_local, + psi_local, + &(pes->ekb(kpt_assigned, 0))); } // transfer the eigenvectors and eigenvalues to all processes - for(int ik = ik_start; ik < ik_start + total_active_ranks && ik < nks; ik++) + for (int ik = ik_start; ik < ik_start + total_active_ranks && ik < nks; ik++) { int root = active_ranks[ik % total_active_ranks]; - MPI_Bcast(&(pes->ekb(ik, 0)), nbands, MPI_DOUBLE, root, MPI_COMM_WORLD); + MPI_Bcast( + &(pes->ekb(ik, 0)), + nbands, + MPI_DOUBLE, + root, + MPI_COMM_WORLD); int desc_pool[9]; std::copy(psi_para_local.desc, psi_para_local.desc + 9, desc_pool); - T* psi_local_ptr = nullptr; + T *psi_local_ptr = nullptr; if (world_rank != root) { desc_pool[1] = -1; - }else + } + else { psi_local_ptr = psi_local.get_pointer(); } psi.fix_k(ik); - Cpxgemr2d(nrow, - nbands, - psi_local_ptr, - 1, - 1, - desc_pool, - psi.get_pointer(), - 1, - 1, - psi_para_global.desc, - psi_para_global.blacs_ctxt); + Cpxgemr2d( + nrow, + nbands, + psi_local_ptr, + 1, + 1, + desc_pool, + psi.get_pointer(), + 1, + 1, + psi_para_global.desc, + psi_para_global.blacs_ctxt); } } @@ -475,7 +637,6 @@ void HSolverLCAO::parakSolve_cusolver(HSMatrix& hs, } #endif - template class HSolverLCAO; template class HSolverLCAO>; diff --git a/source/source_hsolver/hsolver_lcao.h b/source/source_hsolver/hsolver_lcao.h index 06702f8bc7c..38ddc998033 100644 --- a/source/source_hsolver/hsolver_lcao.h +++ b/source/source_hsolver/hsolver_lcao.h @@ -1,60 +1,60 @@ #ifndef HSOLVERLCAO_H #define HSOLVERLCAO_H -#include "source_estate/elecstate.h" -#include "source_hsolver/hs_matrix.h" #include "source_basis/module_ao/parallel_orbitals.h" - -#include "source_estate/module_charge/charge.h" // mohan add 20251024 +#include "source_estate/elecstate.h" +#include "source_estate/module_charge/charge.h" // mohan add 20251024 #include "source_estate/module_dm/density_matrix.h" // mohan add 20251103 +#include "source_hsolver/hs_matrix.h" namespace hsolver { -template -class HSolverLCAO +template class HSolverLCAO { public: - HSolverLCAO(const Parallel_Orbitals* ParaV_in, - const std::string method_in, - const int kpar_lcao_in, - const int nlocal_in, - const int nbands_in, - const double nelec_in, - const bool use_gpu_in, - const int world_nproc_in, - const int world_rank_in) + HSolverLCAO( + const Parallel_Orbitals *ParaV_in, + const std::string method_in, + const int kpar_lcao_in, + const int nlocal_in, + const int nbands_in, + const double nelec_in, + const bool use_gpu_in, + const int world_nproc_in, + const int world_rank_in) : ParaV(ParaV_in), method(method_in), kpar_lcao(kpar_lcao_in), nlocal(nlocal_in), nbands(nbands_in), - nelec(nelec_in), use_gpu(use_gpu_in), world_nproc(world_nproc_in), world_rank(world_rank_in){}; + nelec(nelec_in), use_gpu(use_gpu_in), world_nproc(world_nproc_in), world_rank(world_rank_in) {}; /// @param hs supplies H(k) and S(k) of the Hamiltonian for each k point - void solve(HSMatrix& hs, - psi::Psi& psi, - elecstate::ElecState* pes, - module_dm::DensityMatrix& dm, // mohan add 2025-11-03 - Charge &chr, // charge density - const int nspin, - const double omega, // current cell volume (ucell.omega), NOT rhopw->omega - const bool skip_charge); + void solve( + HSMatrix &hs, + psi::Psi &psi, + elecstate::ElecState *pes, + module_dm::DensityMatrix &dm, // mohan add 2025-11-03 + Charge &chr, // charge density + const int nspin, + const double omega, // current cell volume (ucell.omega), NOT rhopw->omega + const bool skip_charge); private: - void hamiltSolvePsiK(ModuleBase::MatrixBlock& hk, - ModuleBase::MatrixBlock& sk, - psi::Psi& psi, - double* eigenvalue); // for kpar_lcao == 1 + void hamiltSolvePsiK( + ModuleBase::MatrixBlock &hk, + ModuleBase::MatrixBlock &sk, + psi::Psi &psi, + double *eigenvalue); // for kpar_lcao == 1 - void parakSolve(HSMatrix& hs, - psi::Psi& psi, - elecstate::ElecState* pes, - const int kpar, - const int nspin); // for kpar_lcao > 1 + void parakSolve( + HSMatrix &hs, + psi::Psi &psi, + elecstate::ElecState *pes, + const int kpar, + const int nspin); // for kpar_lcao > 1 // The solving algorithm using cusolver is different from others, so a separate function is needed - void parakSolve_cusolver(HSMatrix& hs, - psi::Psi& psi, - elecstate::ElecState* pes); + void parakSolve_cusolver(HSMatrix &hs, psi::Psi &psi, elecstate::ElecState *pes); - const Parallel_Orbitals* ParaV = nullptr; + const Parallel_Orbitals *ParaV = nullptr; const std::string method; diff --git a/source/source_hsolver/hsolver_lcaopw.cpp b/source/source_hsolver/hsolver_lcaopw.cpp index 7c324c22296..dc382edf794 100644 --- a/source/source_hsolver/hsolver_lcaopw.cpp +++ b/source/source_hsolver/hsolver_lcaopw.cpp @@ -17,13 +17,14 @@ namespace hsolver lcao_in_pw */ template -void HSolverLIP::solve(HSOperator& op, // ESolver_KS_PW::p_hamilt behind the operator interface - psi::Psi& psi, // ESolver_KS_PW::kspw_psi - elecstate::ElecState* pes, // ESolver_KS_PW::pes - psi::Psi& transform, - const diag_comm_info& diag_comm, - std::ostream& log, - const bool skip_charge) +void HSolverLIP::solve( + HSOperator &op, // ESolver_KS_PW::p_hamilt behind the operator interface + psi::Psi &psi, // ESolver_KS_PW::kspw_psi + elecstate::ElecState *pes, // ESolver_KS_PW::pes + psi::Psi &transform, + const diag_comm_info &diag_comm, + std::ostream &log, + const bool skip_charge) { ModuleBase::TITLE("HSolverLIP", "solve"); ModuleBase::timer::start("HSolverLIP", "solve"); @@ -37,15 +38,16 @@ void HSolverLIP::solve(HSOperator& op, // ESolver_KS_PW::p_hamilt be transform.fix_k(ik); /// solve eigenvector and eigenvalue for H(k) - hsolver::DiagoIterAssist::diag_subspace_init(op, - transform.get_pointer(), // transform matrix between lcao and pw - transform.get_nbands(), - transform.get_nbasis(), - psi, // psi in pw basis - eigenvalues.data() + ik * pes->ekb.nc, // eigenvalues - this->basis_type, - this->calculation, - diag_comm); + hsolver::DiagoIterAssist::diag_subspace_init( + op, + transform.get_pointer(), // transform matrix between lcao and pw + transform.get_nbands(), + transform.get_nbasis(), + psi, // psi in pw basis + eigenvalues.data() + ik * pes->ekb.nc, // eigenvalues + this->basis_type, + this->calculation, + diag_comm); if (skip_charge) { @@ -57,29 +59,28 @@ void HSolverLIP::solve(HSOperator& op, // ESolver_KS_PW::p_hamilt be /// calculate the contribution of Psi for charge density rho } base_device::memory::cast_memory_op()( - pes->ekb.c, - eigenvalues.data(), - pes->ekb.nr * pes->ekb.nc); + pes->ekb.c, eigenvalues.data(), pes->ekb.nr * pes->ekb.nc); - elecstate::calculate_weights(pes->ekb, - pes->wg, - pes->klist, - pes->eferm, - pes->f_en, - pes->nelec_spin, - this->global_nbands, - pes->skip_weights); - elecstate::calEBand(pes->ekb,pes->wg,pes->f_en); + elecstate::calculate_weights( + pes->ekb, + pes->wg, + pes->klist, + pes->eferm, + pes->f_en, + pes->nelec_spin, + this->global_nbands, + pes->skip_weights); + elecstate::calEBand(pes->ekb, pes->wg, pes->f_en); if (skip_charge) { if (this->use_uspp) { - reinterpret_cast*>(pes)->cal_becsum(psi); + reinterpret_cast *>(pes)->cal_becsum(psi); } ModuleBase::timer::end("HSolverLIP", "solve"); return; } - reinterpret_cast*>(pes)->psiToRho(psi); + reinterpret_cast *>(pes)->psiToRho(psi); ModuleBase::timer::end("HSolverLIP", "solve"); return; diff --git a/source/source_hsolver/hsolver_lcaopw.h b/source/source_hsolver/hsolver_lcaopw.h index 0e5fac0a9fc..5a8072c5363 100644 --- a/source/source_hsolver/hsolver_lcaopw.h +++ b/source/source_hsolver/hsolver_lcaopw.h @@ -4,6 +4,7 @@ #include "source_base/macros.h" #include "source_estate/elecstate.h" #include "source_hsolver/hs_operator.h" + #include namespace hsolver @@ -12,8 +13,7 @@ namespace hsolver struct diag_comm_info; // LCAO-in-PW does not support GPU now. -template -class HSolverLIP +template class HSolverLIP { private: // Note GetTypeReal::type will @@ -22,11 +22,12 @@ class HSolverLIP using Real = typename GetTypeReal::type; public: - HSolverLIP(ModulePW::PW_Basis_K* wfc_basis_in, - const bool use_uspp_in, - const std::string basis_type_in, - const std::string calculation_in, - const int global_nbands_in) + HSolverLIP( + ModulePW::PW_Basis_K *wfc_basis_in, + const bool use_uspp_in, + const std::string basis_type_in, + const std::string calculation_in, + const int global_nbands_in) : wfc_basis(wfc_basis_in), use_uspp(use_uspp_in), basis_type(basis_type_in), calculation(calculation_in), global_nbands(global_nbands_in) {}; @@ -36,16 +37,17 @@ class HSolverLIP /// @param pes interface to elecstate /// @param transform transformation matrix between lcao and pw /// @param skip_charge - void solve(HSOperator& op, - psi::Psi& psi, - elecstate::ElecState* pes, - psi::Psi& transform, - const diag_comm_info& diag_comm, - std::ostream& log, - const bool skip_charge); + void solve( + HSOperator &op, + psi::Psi &psi, + elecstate::ElecState *pes, + psi::Psi &transform, + const diag_comm_info &diag_comm, + std::ostream &log, + const bool skip_charge); private: - ModulePW::PW_Basis_K* wfc_basis = nullptr; + ModulePW::PW_Basis_K *wfc_basis = nullptr; const bool use_uspp; // true if ultrasoft pseudopotentials are in use diff --git a/source/source_hsolver/hsolver_pw.cpp b/source/source_hsolver/hsolver_pw.cpp index 1b316b567e0..21e63573fbf 100644 --- a/source/source_hsolver/hsolver_pw.cpp +++ b/source/source_hsolver/hsolver_pw.cpp @@ -21,10 +21,11 @@ namespace hsolver { template -void HSolverPW::cal_smooth_ethr(const double& wk, - const double* wg, - const double& ethr, - std::vector& ethrs) +void HSolverPW::cal_smooth_ethr( + const double &wk, + const double *wg, + const double ðr, + std::vector ðrs) { // threshold for classifying occupied and unoccupied bands const double occ_threshold = 1e-2; @@ -63,14 +64,15 @@ void HSolverPW::cal_smooth_ethr(const double& wk, } template -void HSolverPW::solve(HSOperator& op, - psi::Psi& psi, - elecstate::ElecState* pes, - double* out_eigenvalues, - const int rank_in_pool_in, - const int nproc_in_pool_in, - std::ostream& log, - const bool skip_charge) +void HSolverPW::solve( + HSOperator &op, + psi::Psi &psi, + elecstate::ElecState *pes, + double *out_eigenvalues, + const int rank_in_pool_in, + const int nproc_in_pool_in, + std::ostream &log, + const bool skip_charge) { ModuleBase::TITLE("HSolverPW", "solve"); ModuleBase::timer::start("HSolverPW", "solve"); @@ -79,7 +81,10 @@ void HSolverPW::solve(HSOperator& op, this->nproc_in_pool = nproc_in_pool_in; // report if the specified diagonalization method is not supported - const std::initializer_list _methods = {"cg", "dav", "dav_subspace", "bpcg"}; + const std::initializer_list _methods = {"cg", + "dav", + "dav_subspace", + "bpcg"}; if (std::find(std::begin(_methods), std::end(_methods), this->method) == std::end(_methods)) { ModuleBase::WARNING_QUIT("HSolverPW::solve", "This type of eigensolver is not supported!"); @@ -92,12 +97,14 @@ void HSolverPW::solve(HSOperator& op, // Initialize k-point continuity if enabled static int count = 0; - if (use_k_continuity) { + if (use_k_continuity) + { build_k_neighbors(); } // Loop over k points for solve Hamiltonian to charge density - if (use_k_continuity) { + if (use_k_continuity) + { // K-point continuity case for (int i = 0; i < this->wfc_basis->nks; ++i) { @@ -106,32 +113,39 @@ void HSolverPW::solve(HSOperator& op, // update H(k) for each k point op.update_k(ik); - - // update psi pointer for each k point psi.fix_k(ik); // If using k-point continuity and not first k-point, propagate from parent - if (ik > 0 && count == 0 && k_parent.find(ik) != k_parent.end()) { + if (ik > 0 && count == 0 && k_parent.find(ik) != k_parent.end()) + { propagate_psi(psi, k_parent[ik], ik); } // template add precondition calculating here - update_precondition(precondition, ik, this->wfc_basis->npwk[ik], Real(pes->pot->get_vl_of_0())); + update_precondition( + precondition, + ik, + this->wfc_basis->npwk[ik], + Real(pes->pot->get_vl_of_0())); // use smooth threshold for all iter methods if (this->diago_smooth_ethr == true) { - this->cal_smooth_ethr(pes->klist->wk[ik], - &pes->wg(ik, 0), - DiagoIterAssist::PW_DIAG_THR, - ethr_band); + this->cal_smooth_ethr( + pes->klist->wk[ik], + &pes->wg(ik, 0), + DiagoIterAssist::PW_DIAG_THR, + ethr_band); } - - // solve eigenvector and eigenvalue for H(k) - this->hamiltSolvePsiK(op, psi, precondition, eigenvalues.data() + ik * psi.get_nbands(), this->wfc_basis->nks); + this->hamiltSolvePsiK( + op, + psi, + precondition, + eigenvalues.data() + ik * psi.get_nbands(), + this->wfc_basis->nks); if (skip_charge) { @@ -142,34 +156,41 @@ void HSolverPW::solve(HSOperator& op, } } } // if (use_k_continuity) - else { + else + { // Original code without k-point continuity for (int ik = 0; ik < this->wfc_basis->nks; ++ik) { // update H(k) for each k point op.update_k(ik); - - // update psi pointer for each k point psi.fix_k(ik); // template add precondition calculating here - update_precondition(precondition, ik, this->wfc_basis->npwk[ik], Real(pes->pot->get_vl_of_0())); + update_precondition( + precondition, + ik, + this->wfc_basis->npwk[ik], + Real(pes->pot->get_vl_of_0())); // use smooth threshold for all iter methods if (this->diago_smooth_ethr == true) { - this->cal_smooth_ethr(pes->klist->wk[ik], - &pes->wg(ik, 0), - DiagoIterAssist::PW_DIAG_THR, - ethr_band); + this->cal_smooth_ethr( + pes->klist->wk[ik], + &pes->wg(ik, 0), + DiagoIterAssist::PW_DIAG_THR, + ethr_band); } - - // solve eigenvector and eigenvalue for H(k) - this->hamiltSolvePsiK(op, psi, precondition, eigenvalues.data() + ik * psi.get_nbands(), this->wfc_basis->nks); + this->hamiltSolvePsiK( + op, + psi, + precondition, + eigenvalues.data() + ik * psi.get_nbands(), + this->wfc_basis->nks); // output iteration information and reset avg_iter if (skip_charge) @@ -192,31 +213,30 @@ void HSolverPW::solve(HSOperator& op, // copy eigenvalues to ekb in ElecState base_device::memory::cast_memory_op()( - out_eigenvalues, - eigenvalues.data(), - this->wfc_basis->nks * psi.get_nbands()); - - auto _pes_pw = reinterpret_cast*>(pes); - elecstate::calculate_weights(_pes_pw->ekb, - _pes_pw->wg, - _pes_pw->klist, - _pes_pw->eferm, - _pes_pw->f_en, - _pes_pw->nelec_spin, - this->nbands, - _pes_pw->skip_weights); - - elecstate::calEBand(_pes_pw->ekb,_pes_pw->wg,_pes_pw->f_en); + out_eigenvalues, eigenvalues.data(), this->wfc_basis->nks * psi.get_nbands()); + + auto _pes_pw = reinterpret_cast *>(pes); + elecstate::calculate_weights( + _pes_pw->ekb, + _pes_pw->wg, + _pes_pw->klist, + _pes_pw->eferm, + _pes_pw->f_en, + _pes_pw->nelec_spin, + this->nbands, + _pes_pw->skip_weights); + + elecstate::calEBand(_pes_pw->ekb, _pes_pw->wg, _pes_pw->f_en); if (skip_charge) { if (this->use_uspp) { - reinterpret_cast*>(pes)->cal_becsum(psi); + reinterpret_cast *>(pes)->cal_becsum(psi); } } else { - reinterpret_cast*>(pes)->psiToRho(psi); + reinterpret_cast *>(pes)->psiToRho(psi); } ModuleBase::timer::end("HSolverPW", "solve"); @@ -224,11 +244,12 @@ void HSolverPW::solve(HSOperator& op, } template -void HSolverPW::hamiltSolvePsiK(const HSOperator& op, - psi::Psi& psi, - std::vector& pre_condition, - Real* eigenvalue, - const int& nk_nums) +void HSolverPW::hamiltSolvePsiK( + const HSOperator &op, + psi::Psi &psi, + std::vector &pre_condition, + Real *eigenvalue, + const int &nk_nums) { ModuleBase::timer::start("HSolverPW", "solve_psik"); #ifdef __MPI @@ -249,38 +270,43 @@ void HSolverPW::hamiltSolvePsiK(const HSOperator& op, #ifdef __MPI if (this->nproc_in_pool > 1) { - MPI_Allreduce(&cur_nbasis, &npw_total, 1, MPI_INT, MPI_SUM, POOL_WORLD); + MPI_Allreduce( + &cur_nbasis, + &npw_total, + 1, + MPI_INT, + MPI_SUM, + POOL_WORLD); } #endif if (npw_total < nbands) { - const std::string msg = "npw_total < nbands (" + std::to_string(npw_total) + " < " + std::to_string(nbands) - + "): the total number of plane waves across all MPI processes " - + "is less than the number of bands, " - + "which leads to a rank-deficient problem. " - + "Please increase ecutwfc or reduce nbands."; + const std::string msg = "npw_total < nbands (" + std::to_string(npw_total) + " < " + std::to_string(nbands) + + "): the total number of plane waves across all MPI processes " + + "is less than the number of bands, " + "which leads to a rank-deficient problem. " + + "Please increase ecutwfc or reduce nbands."; ModuleBase::WARNING_QUIT("HSolverPW::hamiltSolvePsiK", msg); } if (this->method == "cg") { - DiagoCG cg(this->basis_type, - this->calculation_type, - this->need_subspace, - comm_info, - this->diag_thr, - this->diag_iter_max); - - DiagoIterAssist::avg_iter += static_cast( - cg.diag(op, - psi.get_nbasis(), - psi.get_nbands(), - psi.get_current_ngk(), - psi.get_pointer(), - eigenvalue, - this->ethr_band, - pre_condition.data()) - ); + DiagoCG cg( + this->basis_type, + this->calculation_type, + this->need_subspace, + comm_info, + this->diag_thr, + this->diag_iter_max); + + DiagoIterAssist::avg_iter += static_cast(cg.diag( + op, + psi.get_nbasis(), + psi.get_nbands(), + psi.get_current_ngk(), + psi.get_pointer(), + eigenvalue, + this->ethr_band, + pre_condition.data())); // TODO: Double check tensormap's potential problem // ct::TensorMap(psi.get_pointer(), psi_tensor, {psi.get_nbands(), psi.get_nbasis()}).sync(psi_tensor); } @@ -290,31 +316,39 @@ void HSolverPW::hamiltSolvePsiK(const HSOperator& op, const int nbasis = psi.get_nbasis(); const int ndim = psi.get_current_ngk(); DiagoBPCG bpcg(pre_condition.data()); - bpcg.init_iter(this->nbands, nband_l, nbasis, ndim); - bpcg.diag(op, psi.get_pointer(), eigenvalue, this->ethr_band); + bpcg.init_iter( + this->nbands, + nband_l, + nbasis, + ndim); + bpcg.diag( + op, + psi.get_pointer(), + eigenvalue, + this->ethr_band); } else if (this->method == "dav_subspace") { bool scf = this->calculation_type == "nscf" ? false : true; - Diago_DavSubspace dav_subspace(pre_condition, - psi.get_nbands(), - psi.get_k_first() ? psi.get_current_ngk() - : psi.get_nk() * psi.get_nbasis(), - this->pw_diag_ndim, - this->diag_thr, - this->diag_iter_max, - comm_info, - this->diag_subspace, - this->nb2d); - - DiagoIterAssist::avg_iter += static_cast( - dav_subspace.diag(op, - psi.get_pointer(), - psi.get_nbasis(), - eigenvalue, - this->ethr_band, - scf)); + Diago_DavSubspace dav_subspace( + pre_condition, + psi.get_nbands(), + psi.get_k_first() ? psi.get_current_ngk() : psi.get_nk() * psi.get_nbasis(), + this->pw_diag_ndim, + this->diag_thr, + this->diag_iter_max, + comm_info, + this->diag_subspace, + this->nb2d); + + DiagoIterAssist::avg_iter += static_cast(dav_subspace.diag( + op, + psi.get_pointer(), + psi.get_nbasis(), + eigenvalue, + this->ethr_band, + scf)); } else if (this->method == "dav") { @@ -332,30 +366,36 @@ void HSolverPW::hamiltSolvePsiK(const HSOperator& op, // dimensions of matrix to be solved const int dim = psi.get_current_ngk(); /// dimension of matrix - const int nband = psi.get_nbands(); /// number of eigenpairs sought - const int ld_psi = psi.get_nbasis(); /// leading dimension of psi - - DiagoDavid david(pre_condition.data(), nband, dim, this->pw_diag_ndim, comm_info); + const int nband = psi.get_nbands(); /// number of eigenpairs sought + const int ld_psi = psi.get_nbasis(); /// leading dimension of psi + + DiagoDavid david( + pre_condition.data(), + nband, + dim, + this->pw_diag_ndim, + comm_info); // do diag and add davidson iteration counts up to avg_iter - DiagoIterAssist::avg_iter += static_cast( - david.diag(op, - ld_psi, - psi.get_pointer(), - eigenvalue, - this->ethr_band, - david_maxiter, - ntry_max, - notconv_max)); + DiagoIterAssist::avg_iter += static_cast(david.diag( + op, + ld_psi, + psi.get_pointer(), + eigenvalue, + this->ethr_band, + david_maxiter, + ntry_max, + notconv_max)); } ModuleBase::timer::end("HSolverPW", "solve_psik"); return; } template -void HSolverPW::update_precondition(std::vector& h_diag, - const int ik, - const int npw, - const Real vl_of_0) +void HSolverPW::update_precondition( + std::vector &h_diag, + const int ik, + const int npw, + const Real vl_of_0) { h_diag.assign(h_diag.size(), 1.0); int precondition_type = 2; @@ -400,8 +440,7 @@ void HSolverPW::update_precondition(std::vector& h_diag, } } -template -void HSolverPW::output_iterInfo(std::ostream& log) +template void HSolverPW::output_iterInfo(std::ostream &log) { // in PW base, average iteration steps for each band and k-point should be printing if (DiagoIterAssist::avg_iter > 0.0) @@ -414,50 +453,51 @@ void HSolverPW::output_iterInfo(std::ostream& log) } } -template -void HSolverPW::build_k_neighbors() { +template void HSolverPW::build_k_neighbors() +{ const int nk = this->wfc_basis->nks; kvecs_c.resize(nk); k_order.clear(); k_order.reserve(nk); // Store k-points and corresponding indices - struct KPoint { + struct KPoint + { ModuleBase::Vector3 kvec; int index = 0; double norm = 0.0; - KPoint(const ModuleBase::Vector3& v, int i) : - kvec(v), index(i), norm(v.norm()) {} + KPoint(const ModuleBase::Vector3 &v, int i) : kvec(v), index(i), norm(v.norm()) {} }; // Build k-point list std::vector klist; - for (int ik = 0; ik < nk; ++ik) { + for (int ik = 0; ik < nk; ++ik) + { kvecs_c[ik] = this->wfc_basis->kvec_c[ik]; klist.push_back(KPoint(kvecs_c[ik], ik)); } // Sort k-points by distance from origin - std::sort(klist.begin(), klist.end(), - [](const KPoint& a, const KPoint& b) { - return a.norm < b.norm; - }); + std::sort(klist.begin(), klist.end(), [](const KPoint &a, const KPoint &b) { return a.norm < b.norm; }); // Build parent-child relationships k_order.push_back(klist[0].index); // Find nearest processed k-point as parent for each k-point - for (int i = 1; i < nk; ++i) { + for (int i = 1; i < nk; ++i) + { int current_k = klist[i].index; double min_dist = 1e10; int parent = -1; // find the nearest k-point as parent - for (int j = 0; j < k_order.size(); ++j) { + for (int j = 0; j < k_order.size(); ++j) + { int processed_k = k_order[j]; double dist = (kvecs_c[current_k] - kvecs_c[processed_k]).norm2(); - if (dist < min_dist) { + if (dist < min_dist) + { min_dist = dist; parent = processed_k; } @@ -469,7 +509,8 @@ void HSolverPW::build_k_neighbors() { } template -void HSolverPW::propagate_psi(psi::Psi& psi, const int from_ik, const int to_ik) { +void HSolverPW::propagate_psi(psi::Psi &psi, const int from_ik, const int to_ik) +{ const int nbands = psi.get_nbands(); const int npwk = this->wfc_basis->npwk[to_ik]; @@ -477,7 +518,7 @@ void HSolverPW::propagate_psi(psi::Psi& psi, const int fro ModuleBase::Vector3 dk = kvecs_c[to_ik] - kvecs_c[from_ik]; // Allocate porter locally - T* porter = nullptr; + T *porter = nullptr; resmem_complex_op()(porter, this->wfc_basis->nmaxgr, "HSolverPW::porter"); // Process each band @@ -488,7 +529,11 @@ void HSolverPW::propagate_psi(psi::Psi& psi, const int fro // FFT to real space // this->wfc_basis->recip_to_real(this->ctx, psi.get_pointer(ib), porter, from_ik); - this->wfc_basis->recip_to_real(this->ctx, &psi(from_ik, ib, 0), porter, from_ik); + this->wfc_basis->recip_to_real( + this->ctx, + &psi(from_ik, ib, 0), + porter, + from_ik); // Apply phase factor // // TODO: Check how to get the r vector @@ -502,7 +547,11 @@ void HSolverPW::propagate_psi(psi::Psi& psi, const int fro // FFT back to reciprocal space // this->wfc_basis->real_to_recip(this->ctx, porter, psi.get_pointer(ib), to_ik, true); - this->wfc_basis->real_to_recip(this->ctx, porter, &psi(to_ik, ib, 0), to_ik); + this->wfc_basis->real_to_recip( + this->ctx, + porter, + &psi(to_ik, ib, 0), + to_ik); } // Clean up porter diff --git a/source/source_hsolver/hsolver_pw.h b/source/source_hsolver/hsolver_pw.h index db6764154c8..8d8436ade44 100644 --- a/source/source_hsolver/hsolver_pw.h +++ b/source/source_hsolver/hsolver_pw.h @@ -14,8 +14,7 @@ class TestHSolverPW; // unit-test fixture, see source_hsolver/unittests/ namespace hsolver { -template -class HSolverPW +template class HSolverPW { /// the unit test drives the protected hamiltSolvePsiK() and /// update_precondition() directly; the fixture is at global scope @@ -31,57 +30,63 @@ class HSolverPW using setmem_complex_op = base_device::memory::set_memory_op; public: - HSolverPW(ModulePW::PW_Basis_K* wfc_basis_in, - const std::string calculation_type_in, - const std::string basis_type_in, - const std::string method_in, - const bool use_uspp_in, - const int nspin_in, - const int scf_iter_in, - const int diag_iter_max_in, - const double diag_thr_in, - const bool need_subspace_in, - const int nbands_in, - const bool diago_smooth_ethr_in, - const int pw_diag_ndim_in, - const int diag_subspace_in, - const int nb2d_in, - const bool use_k_continuity_in = false) + HSolverPW( + ModulePW::PW_Basis_K *wfc_basis_in, + const std::string calculation_type_in, + const std::string basis_type_in, + const std::string method_in, + const bool use_uspp_in, + const int nspin_in, + const int scf_iter_in, + const int diag_iter_max_in, + const double diag_thr_in, + const bool need_subspace_in, + const int nbands_in, + const bool diago_smooth_ethr_in, + const int pw_diag_ndim_in, + const int diag_subspace_in, + const int nb2d_in, + const bool use_k_continuity_in = false) : wfc_basis(wfc_basis_in), calculation_type(calculation_type_in), basis_type(basis_type_in), method(method_in), - use_uspp(use_uspp_in), nspin(nspin_in), scf_iter(scf_iter_in), - diag_iter_max(diag_iter_max_in), diag_thr(diag_thr_in), need_subspace(need_subspace_in), - nbands(nbands_in), diago_smooth_ethr(diago_smooth_ethr_in), pw_diag_ndim(pw_diag_ndim_in), - diag_subspace(diag_subspace_in), nb2d(nb2d_in), - use_k_continuity(use_k_continuity_in) {}; + use_uspp(use_uspp_in), nspin(nspin_in), scf_iter(scf_iter_in), diag_iter_max(diag_iter_max_in), + diag_thr(diag_thr_in), need_subspace(need_subspace_in), nbands(nbands_in), + diago_smooth_ethr(diago_smooth_ethr_in), pw_diag_ndim(pw_diag_ndim_in), diag_subspace(diag_subspace_in), + nb2d(nb2d_in), use_k_continuity(use_k_continuity_in) {}; /// @brief solve function for pw /// @param op the H and S operator of the Hamiltonian, switched to each k point in turn /// @param psi reference to psi /// @param pes interface to elecstate /// @param skip_charge - void solve(HSOperator& op, - psi::Psi& psi, - elecstate::ElecState* pes, - double* out_eigenvalues, - const int rank_in_pool_in, - const int nproc_in_pool_in, - std::ostream& log, - const bool skip_charge); + void solve( + HSOperator &op, + psi::Psi &psi, + elecstate::ElecState *pes, + double *out_eigenvalues, + const int rank_in_pool_in, + const int nproc_in_pool_in, + std::ostream &log, + const bool skip_charge); protected: // diago caller - void hamiltSolvePsiK(const HSOperator& op, - psi::Psi& psi, - std::vector& pre_condition, - Real* eigenvalue, - const int& nk_nums); + void hamiltSolvePsiK( + const HSOperator &op, + psi::Psi &psi, + std::vector &pre_condition, + Real *eigenvalue, + const int &nk_nums); // calculate the precondition array for diagonalization in PW base - void update_precondition(std::vector& h_diag, const int ik, const int npw, const Real vl_of_0); + void update_precondition( + std::vector &h_diag, + const int ik, + const int npw, + const Real vl_of_0); - void output_iterInfo(std::ostream& log); + void output_iterInfo(std::ostream &log); - ModulePW::PW_Basis_K* wfc_basis = nullptr; + ModulePW::PW_Basis_K *wfc_basis = nullptr; const std::string calculation_type; const std::string basis_type; @@ -95,16 +100,16 @@ class HSolverPW const bool need_subspace; // for cg or dav_subspace - const int nbands; // global number of bands, may differ from psi.get_nbands() under band parallelism - const bool diago_smooth_ethr; // use a band-wise smoothed threshold for all iter methods - const int pw_diag_ndim; // dimension of the workspace for Davidson-type methods - const int diag_subspace; // subspace eigensolver for dav_subspace: 0 Lapack, 1 elpa, 2 scalapack - const int nb2d; // 2d block size used by the dav_subspace scalapack path + const int nbands; // global number of bands, may differ from psi.get_nbands() under band parallelism + const bool diago_smooth_ethr; // use a band-wise smoothed threshold for all iter methods + const int pw_diag_ndim; // dimension of the workspace for Davidson-type methods + const int diag_subspace; // subspace eigensolver for dav_subspace: 0 Lapack, 1 elpa, 2 scalapack + const int nb2d; // 2d block size used by the dav_subspace scalapack path const bool use_k_continuity; protected: - Device* ctx = {}; + Device *ctx = {}; int rank_in_pool = 0; int nproc_in_pool = 1; @@ -113,17 +118,19 @@ class HSolverPW private: /// @brief calculate the threshold for iterative-diagonalization for each band - void cal_smooth_ethr(const double& wk, const double* wg, const double& ethr, std::vector& ethrs); - - + void cal_smooth_ethr( + const double &wk, + const double *wg, + const double ðr, + std::vector ðrs); // K-point continuity related members std::vector k_order; std::unordered_map k_parent; std::vector> kvecs_c; - + void build_k_neighbors(); - void propagate_psi(psi::Psi& psi, const int from_ik, const int to_ik); + void propagate_psi(psi::Psi &psi, const int from_ik, const int to_ik); }; } // namespace hsolver diff --git a/source/source_hsolver/kernels/bpcg_kernel_op.cpp b/source/source_hsolver/kernels/bpcg_kernel_op.cpp index 3b04f1390e2..89115238bb8 100644 --- a/source/source_hsolver/kernels/bpcg_kernel_op.cpp +++ b/source/source_hsolver/kernels/bpcg_kernel_op.cpp @@ -1,23 +1,25 @@ #include "source_hsolver/kernels/bpcg_kernel_op.h" + #include "source_base/kernels/math_kernel_op.h" #include "source_base/parallel_reduce.h" + #include namespace hsolver { -template -struct line_minimize_with_block_op +template struct line_minimize_with_block_op { using Real = typename GetTypeReal::type; - void operator()(T* grad_out, - T* hgrad_out, - T* sgrad_out, - T* psi_out, - T* hpsi_out, - T* spsi_out, - const int& n_basis, - const int& n_basis_max, - const int& n_band) + void operator()( + T *grad_out, + T *hgrad_out, + T *sgrad_out, + T *psi_out, + T *hpsi_out, + T *spsi_out, + const int &n_basis, + const int &n_basis_max, + const int &n_band) { for (int band_idx = 0; band_idx < n_band; ++band_idx) { @@ -51,8 +53,8 @@ struct line_minimize_with_block_op Parallel_Reduce::reduce_pool(epsilo_2); Real theta = 0.5 * std::atan2(2 * epsilo_1, epsilo_0 - epsilo_2); // Choose the rotation associated with the lower Ritz value. - const Real energy_delta = (epsilo_0 - epsilo_2) * std::cos(2.0 * theta) - + 2.0 * epsilo_1 * std::sin(2.0 * theta); + const Real energy_delta = + (epsilo_0 - epsilo_2) * std::cos(2.0 * theta) + 2.0 * epsilo_1 * std::sin(2.0 * theta); const Real energy_1 = 0.5 * (epsilo_0 + epsilo_2 + energy_delta); const Real energy_2 = 0.5 * (epsilo_0 + epsilo_2 - energy_delta); if (energy_1 > energy_2) @@ -72,21 +74,21 @@ struct line_minimize_with_block_op } }; -template -struct calc_grad_with_block_op +template struct calc_grad_with_block_op { using Real = typename GetTypeReal::type; - void operator()(const Real* prec_in, - Real* err_out, - Real* beta_out, - T* psi_out, - T* hpsi_out, - T* spsi_out, - T* grad_out, - T* grad_old_out, - const int& n_basis, - const int& n_basis_max, - const int& n_band) + void operator()( + const Real *prec_in, + Real *err_out, + Real *beta_out, + T *psi_out, + T *hpsi_out, + T *spsi_out, + T *grad_out, + T *grad_old_out, + const int &n_basis, + const int &n_basis_max, + const int &n_band) { for (int band_idx = 0; band_idx < n_band; ++band_idx) { @@ -134,11 +136,16 @@ struct calc_grad_with_block_op } }; -template -struct apply_eigenvalues_op +template struct apply_eigenvalues_op { using Real = typename GetTypeReal::type; - void operator()(const int& nbase, const int& nbase_x, const int& notconv, T* result, const T* vectors, const Real* eigenvalues) + void operator()( + const int &nbase, + const int &nbase_x, + const int ¬conv, + T *result, + const T *vectors, + const Real *eigenvalues) { for (int m = 0; m < notconv; m++) { @@ -150,15 +157,16 @@ struct apply_eigenvalues_op } }; -template -struct precondition_op { +template struct precondition_op +{ using Real = typename GetTypeReal::type; - void operator()(const int& dim, - T* psi_iter, - const int& nbase, - const int& notconv, - const Real* precondition, - const Real* eigenvalues) + void operator()( + const int &dim, + T *psi_iter, + const int &nbase, + const int ¬conv, + const Real *precondition, + const Real *eigenvalues) { std::vector pre(dim, 0.0); for (int m = 0; m < notconv; m++) @@ -169,58 +177,60 @@ struct precondition_op { pre[i] = 0.5 * (1.0 + x + sqrt(1 + (x - 1.0) * (x - 1.0))); } ModuleBase::vector_div_vector_op()( - dim, - psi_iter + (nbase + m) * dim, - psi_iter + (nbase + m) * dim, - pre.data()); + dim, + psi_iter + (nbase + m) * dim, + psi_iter + (nbase + m) * dim, + pre.data()); } } }; -template -struct normalize_op { - void operator()(const int& dim, - T* psi_iter, - const int& nbase, - const int& notconv, - typename GetTypeReal::type* psi_norm) +template struct normalize_op +{ + void operator()( + const int &dim, + T *psi_iter, + const int &nbase, + const int ¬conv, + typename GetTypeReal::type *psi_norm) { using Real = typename GetTypeReal::type; for (int m = 0; m < notconv; m++) { // Calculate norm using dot_real_op Real psi_m_norm = ModuleBase::dot_real_op()( - dim, - psi_iter + (nbase + m) * dim, - psi_iter + (nbase + m) * dim, - true); + dim, + psi_iter + (nbase + m) * dim, + psi_iter + (nbase + m) * dim, + true); assert(psi_m_norm > 0.0); psi_m_norm = sqrt(psi_m_norm); // Normalize using vector_div_constant_op ModuleBase::vector_div_constant_op()( - dim, - psi_iter + (nbase + m) * dim, - psi_iter + (nbase + m) * dim, - psi_m_norm); - if (psi_norm) { + dim, + psi_iter + (nbase + m) * dim, + psi_iter + (nbase + m) * dim, + psi_m_norm); + if (psi_norm) + { psi_norm[m] = psi_m_norm; } } } }; -template -struct refresh_hcc_scc_vcc_op +template struct refresh_hcc_scc_vcc_op { using Real = typename GetTypeReal::type; - void operator()(const int &n, - T *hcc, - T *scc, - T *vcc, - const int &ldh, - const Real *eigenvalue, - const T &one) + void operator()( + const int &n, + T *hcc, + T *scc, + T *vcc, + const int &ldh, + const Real *eigenvalue, + const T &one) { #ifdef _OPENMP #pragma omp parallel for collapse(1) schedule(static) diff --git a/source/source_hsolver/kernels/bpcg_kernel_op.h b/source/source_hsolver/kernels/bpcg_kernel_op.h index 83e7ce0a1a7..6d9aceb2b21 100644 --- a/source/source_hsolver/kernels/bpcg_kernel_op.h +++ b/source/source_hsolver/kernels/bpcg_kernel_op.h @@ -5,8 +5,7 @@ namespace hsolver { -template -struct line_minimize_with_block_op +template struct line_minimize_with_block_op { /// @brief dot_real_op computes the dot product of the given complex /// arrays(treated as float arrays). And there's may have MPI communications @@ -21,19 +20,19 @@ struct line_minimize_with_block_op /// /// \return res : the result vector /// T : dot product result - void operator()(T* grad_out, - T* hgrad_out, - T* sgrad_out, - T* psi_out, - T* hpsi_out, - T* spsi_out, - const int& n_basis, - const int& n_basis_max, - const int& n_band); + void operator()( + T *grad_out, + T *hgrad_out, + T *sgrad_out, + T *psi_out, + T *hpsi_out, + T *spsi_out, + const int &n_basis, + const int &n_basis_max, + const int &n_band); }; -template -struct calc_grad_with_block_op +template struct calc_grad_with_block_op { /// @brief dot_real_op computes the dot product of the given complex /// arrays(treated as float arrays). And there's may have MPI communications @@ -49,139 +48,161 @@ struct calc_grad_with_block_op /// \return res : the result vector /// T : dot product result using Real = typename GetTypeReal::type; - void operator()(const Real* prec_in, - Real* err_out, - Real* beta_out, - T* psi_out, - T* hpsi_out, - T* spsi_out, - T* grad_out, - T* grad_old_out, - const int& n_basis, - const int& n_basis_max, - const int& n_band); + void operator()( + const Real *prec_in, + Real *err_out, + Real *beta_out, + T *psi_out, + T *hpsi_out, + T *spsi_out, + T *grad_out, + T *grad_old_out, + const int &n_basis, + const int &n_basis_max, + const int &n_band); }; -template -struct apply_eigenvalues_op +template struct apply_eigenvalues_op { using Real = typename GetTypeReal::type; - void operator()(const int& nbase, - const int& nbase_x, - const int& notconv, - T* result, - const T* vectors, - const Real* eigenvalues); + void operator()( + const int &nbase, + const int &nbase_x, + const int ¬conv, + T *result, + const T *vectors, + const Real *eigenvalues); }; -template -struct precondition_op { +template struct precondition_op +{ using Real = typename GetTypeReal::type; - void operator()(const int& dim, - T* psi_iter, - const int& nbase, - const int& notconv, - const Real* precondition, - const Real* eigenvalues); + void operator()( + const int &dim, + T *psi_iter, + const int &nbase, + const int ¬conv, + const Real *precondition, + const Real *eigenvalues); }; -template -struct normalize_op { +template struct normalize_op +{ using Real = typename GetTypeReal::type; - void operator()(const int& dim, - T* psi_iter, - const int& nbase, - const int& notconv, - Real* psi_norm = nullptr); + void operator()( + const int &dim, + T *psi_iter, + const int &nbase, + const int ¬conv, + Real *psi_norm = nullptr); }; -template struct refresh_hcc_scc_vcc_op { +template struct refresh_hcc_scc_vcc_op +{ using Real = typename GetTypeReal::type; - /// @brief refresh hcc scc vcc - /// - /// Input Parameters - /// \param n : first dimension of matrix - /// \param ldh : leading dimension of hcc, scc, vcc - /// \param nbase : matrix size - /// \param eigenvalue : input eigenvalue - /// \param one : constant one - /// - /// Output Parameters - /// \param hcc : output matrix hcc - /// \param scc : output matrix scc - /// \param vcc : output matrix vcc - void operator()(const int &n, - T *hcc, - T *scc, - T *vcc, - const int &ldh, - const Real *eigenvalue, - const T& one); + /// @brief refresh hcc scc vcc + /// + /// Input Parameters + /// \param n : first dimension of matrix + /// \param ldh : leading dimension of hcc, scc, vcc + /// \param nbase : matrix size + /// \param eigenvalue : input eigenvalue + /// \param one : constant one + /// + /// Output Parameters + /// \param hcc : output matrix hcc + /// \param scc : output matrix scc + /// \param vcc : output matrix vcc + void operator()( + const int &n, + T *hcc, + T *scc, + T *vcc, + const int &ldh, + const Real *eigenvalue, + const T &one); }; #if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM -template -struct line_minimize_with_block_op { - using Real = typename GetTypeReal::type; - void operator()(T *grad_out, T *hgrad_out, T *sgrad_out, - T *psi_out, T *hpsi_out, T *spsi_out, - const int &n_basis, const int &n_basis_max, - const int &n_band); +template struct line_minimize_with_block_op +{ + using Real = typename GetTypeReal::type; + void operator()( + T *grad_out, + T *hgrad_out, + T *sgrad_out, + T *psi_out, + T *hpsi_out, + T *spsi_out, + const int &n_basis, + const int &n_basis_max, + const int &n_band); }; -template -struct calc_grad_with_block_op { - using Real = typename GetTypeReal::type; - void operator()(const Real *prec_in, Real *err_out, Real *beta_out, - T *psi_out, T *hpsi_out, T *spsi_out, - T *grad_out, T *grad_old_out, - const int &n_basis, const int &n_basis_max, - const int &n_band); +template struct calc_grad_with_block_op +{ + using Real = typename GetTypeReal::type; + void operator()( + const Real *prec_in, + Real *err_out, + Real *beta_out, + T *psi_out, + T *hpsi_out, + T *spsi_out, + T *grad_out, + T *grad_old_out, + const int &n_basis, + const int &n_basis_max, + const int &n_band); }; -template -struct apply_eigenvalues_op { - using Real = typename GetTypeReal::type; - void operator()(const int& nbase, - const int& nbase_x, - const int& notconv, - T* result, - const T* vectors, - const Real* eigenvalues); +template struct apply_eigenvalues_op +{ + using Real = typename GetTypeReal::type; + void operator()( + const int &nbase, + const int &nbase_x, + const int ¬conv, + T *result, + const T *vectors, + const Real *eigenvalues); }; -template -struct precondition_op { - using Real = typename GetTypeReal::type; - void operator()(const int& dim, - T* psi_iter, - const int& nbase, - const int& notconv, - const Real* precondition, - const Real* eigenvalues); +template struct precondition_op +{ + using Real = typename GetTypeReal::type; + void operator()( + const int &dim, + T *psi_iter, + const int &nbase, + const int ¬conv, + const Real *precondition, + const Real *eigenvalues); }; -template -struct normalize_op { - using Real = typename GetTypeReal::type; - void operator()(const int& dim, - T* psi_iter, - const int& nbase, - const int& notconv, - Real* psi_norm = nullptr); +template struct normalize_op +{ + using Real = typename GetTypeReal::type; + void operator()( + const int &dim, + T *psi_iter, + const int &nbase, + const int ¬conv, + Real *psi_norm = nullptr); }; -template -struct refresh_hcc_scc_vcc_op { - using Real = typename GetTypeReal::type; - void operator()(const int &n, - T *hcc, - T *scc, - T *vcc, - const int &ldh, - const Real *eigenvalue, - const T& one); +template struct refresh_hcc_scc_vcc_op +{ + using Real = typename GetTypeReal::type; + void operator()( + const int &n, + T *hcc, + T *scc, + T *vcc, + const int &ldh, + const Real *eigenvalue, + const T &one); }; #endif } // namespace hsolver diff --git a/source/source_hsolver/kernels/cuda/bpcg_kernel_op.cu b/source/source_hsolver/kernels/cuda/bpcg_kernel_op.cu index 48113fe1d43..008fa18f990 100644 --- a/source/source_hsolver/kernels/cuda/bpcg_kernel_op.cu +++ b/source/source_hsolver/kernels/cuda/bpcg_kernel_op.cu @@ -12,17 +12,17 @@ const int thread_per_block = 256; template __global__ void line_minimize_with_block( - thrust::complex* grad, - thrust::complex* hgrad, - thrust::complex* sgrad, - thrust::complex* psi, - thrust::complex* hpsi, - thrust::complex* spsi, - const int n_basis, - const int n_basis_max) + thrust::complex *grad, + thrust::complex *hgrad, + thrust::complex *sgrad, + thrust::complex *psi, + thrust::complex *hpsi, + thrust::complex *spsi, + const int n_basis, + const int n_basis_max) { int band_idx = blockIdx.x; // band_idx - int tid = threadIdx.x; // basis_idx + int tid = threadIdx.x; // basis_idx int item = 0; Real epsilo_0 = 0.0, epsilo_1 = 0.0, epsilo_2 = 0.0; Real theta = 0.0, cos_theta = 0.0, sin_theta = 0.0; @@ -30,14 +30,17 @@ __global__ void line_minimize_with_block( data[tid] = 0; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) + { item = band_idx * n_basis_max + basis_idx; data[tid] += (sgrad[item] * thrust::conj(grad[item])).real(); } __syncthreads(); // just do some parallel reduction in shared memory - for (int ii = thread_per_block >> 1; ii > warp_size; ii >>= 1) { - if (tid < ii) { + for (int ii = thread_per_block >> 1; ii > warp_size; ii >>= 1) + { + if (tid < ii) + { data[tid] += data[tid + ii]; } __syncthreads(); @@ -46,18 +49,26 @@ __global__ void line_minimize_with_block( // For threads in the same warp, it is better that they process the same work // Also, __syncwarp() should be used instead of __syncthreads() // Therefore we unroll the loop and ensure that the threads does the same work - if (tid < warp_size) { - data[tid] += data[tid + 32]; __syncwarp(); - data[tid] += data[tid + 16]; __syncwarp(); - data[tid] += data[tid + 8]; __syncwarp(); - data[tid] += data[tid + 4]; __syncwarp(); - data[tid] += data[tid + 2]; __syncwarp(); - data[tid] += data[tid + 1]; __syncwarp(); + if (tid < warp_size) + { + data[tid] += data[tid + 32]; + __syncwarp(); + data[tid] += data[tid + 16]; + __syncwarp(); + data[tid] += data[tid + 8]; + __syncwarp(); + data[tid] += data[tid + 4]; + __syncwarp(); + data[tid] += data[tid + 2]; + __syncwarp(); + data[tid] += data[tid + 1]; + __syncwarp(); } __syncthreads(); - if (!(data[0] > 1.0e-20)) { + if (!(data[0] > 1.0e-20)) + { return; } Real norm = 1.0 / sqrt(data[0]); @@ -66,7 +77,8 @@ __global__ void line_minimize_with_block( data[tid] = 0; data[thread_per_block + tid] = 0; data[2 * thread_per_block + tid] = 0; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) + { item = band_idx * n_basis_max + basis_idx; grad[item] *= norm; hgrad[item] *= norm; @@ -78,33 +90,54 @@ __global__ void line_minimize_with_block( __syncthreads(); // just do some parallel reduction in shared memory - for (int ii = thread_per_block >> 1; ii > warp_size; ii >>= 1) { - if (tid < ii) { + for (int ii = thread_per_block >> 1; ii > warp_size; ii >>= 1) + { + if (tid < ii) + { data[tid] += data[tid + ii]; data[thread_per_block + tid] += data[thread_per_block + tid + ii]; data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + ii]; } __syncthreads(); } - if (tid < warp_size) { - data[tid] += data[tid + 32]; __syncwarp(); - data[tid] += data[tid + 16]; __syncwarp(); - data[tid] += data[tid + 8]; __syncwarp(); - data[tid] += data[tid + 4]; __syncwarp(); - data[tid] += data[tid + 2]; __syncwarp(); - data[tid] += data[tid + 1]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 32]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 16]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 8]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 4]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 2]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 1]; __syncwarp(); - data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 32]; __syncwarp(); - data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 16]; __syncwarp(); - data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 8]; __syncwarp(); - data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 4]; __syncwarp(); - data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 2]; __syncwarp(); - data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 1]; __syncwarp(); + if (tid < warp_size) + { + data[tid] += data[tid + 32]; + __syncwarp(); + data[tid] += data[tid + 16]; + __syncwarp(); + data[tid] += data[tid + 8]; + __syncwarp(); + data[tid] += data[tid + 4]; + __syncwarp(); + data[tid] += data[tid + 2]; + __syncwarp(); + data[tid] += data[tid + 1]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 32]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 16]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 8]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 4]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 2]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 1]; + __syncwarp(); + data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 32]; + __syncwarp(); + data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 16]; + __syncwarp(); + data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 8]; + __syncwarp(); + data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 4]; + __syncwarp(); + data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 2]; + __syncwarp(); + data[2 * thread_per_block + tid] += data[2 * thread_per_block + tid + 1]; + __syncwarp(); } __syncthreads(); epsilo_0 = data[0]; @@ -113,18 +146,19 @@ __global__ void line_minimize_with_block( theta = 0.5 * atan2(2 * epsilo_1, epsilo_0 - epsilo_2); // Choose the rotation associated with the lower Ritz value. - const Real energy_delta = (epsilo_0 - epsilo_2) * cos(2.0 * theta) - + 2.0 * epsilo_1 * sin(2.0 * theta); + const Real energy_delta = (epsilo_0 - epsilo_2) * cos(2.0 * theta) + 2.0 * epsilo_1 * sin(2.0 * theta); const Real energy_1 = 0.5 * (epsilo_0 + epsilo_2 + energy_delta); const Real energy_2 = 0.5 * (epsilo_0 + epsilo_2 - energy_delta); - if (energy_1 > energy_2) { + if (energy_1 > energy_2) + { theta += 2.0 * atan(1.0); } cos_theta = cos(theta); sin_theta = sin(theta); - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) + { item = band_idx * n_basis_max + basis_idx; - psi [item] = psi [item] * cos_theta + grad [item] * sin_theta; + psi[item] = psi[item] * cos_theta + grad[item] * sin_theta; hpsi[item] = hpsi[item] * cos_theta + hgrad[item] * sin_theta; spsi[item] = spsi[item] * cos_theta + sgrad[item] * sin_theta; } @@ -132,19 +166,19 @@ __global__ void line_minimize_with_block( template __global__ void calc_grad_with_block( - const Real* prec, - Real* err, - Real* beta, - thrust::complex* psi, - thrust::complex* hpsi, - thrust::complex* spsi, - thrust::complex* grad, - thrust::complex* grad_old, - const int n_basis, - const int n_basis_max) + const Real *prec, + Real *err, + Real *beta, + thrust::complex *psi, + thrust::complex *hpsi, + thrust::complex *spsi, + thrust::complex *grad, + thrust::complex *grad_old, + const int n_basis, + const int n_basis_max) { int band_idx = blockIdx.x; // band_idx - int tid = threadIdx.x; // basis_idx + int tid = threadIdx.x; // basis_idx int item = 0; Real err_st = 0.0; Real beta_st = 0.0; @@ -156,26 +190,36 @@ __global__ void calc_grad_with_block( // Init shared memory data[tid] = 0; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) + { item = band_idx * n_basis_max + basis_idx; data[tid] += (spsi[item] * thrust::conj(psi[item])).real(); } __syncthreads(); // just do some parallel reduction in shared memory - for (int ii = thread_per_block >> 1; ii > warp_size; ii >>= 1) { - if (tid < ii) { + for (int ii = thread_per_block >> 1; ii > warp_size; ii >>= 1) + { + if (tid < ii) + { data[tid] += data[tid + ii]; } __syncthreads(); } - if (tid < warp_size) { - data[tid] += data[tid + 32]; __syncwarp(); - data[tid] += data[tid + 16]; __syncwarp(); - data[tid] += data[tid + 8]; __syncwarp(); - data[tid] += data[tid + 4]; __syncwarp(); - data[tid] += data[tid + 2]; __syncwarp(); - data[tid] += data[tid + 1]; __syncwarp(); + if (tid < warp_size) + { + data[tid] += data[tid + 32]; + __syncwarp(); + data[tid] += data[tid + 16]; + __syncwarp(); + data[tid] += data[tid + 8]; + __syncwarp(); + data[tid] += data[tid + 4]; + __syncwarp(); + data[tid] += data[tid + 2]; + __syncwarp(); + data[tid] += data[tid + 1]; + __syncwarp(); } __syncthreads(); @@ -184,7 +228,8 @@ __global__ void calc_grad_with_block( __syncthreads(); data[tid] = 0; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) + { item = band_idx * n_basis_max + basis_idx; psi[item] *= norm; hpsi[item] *= norm; @@ -194,20 +239,29 @@ __global__ void calc_grad_with_block( __syncthreads(); // just do some parallel reduction in shared memory - for (int ii = thread_per_block >> 1; ii > warp_size; ii >>= 1) { - if (tid < ii) { + for (int ii = thread_per_block >> 1; ii > warp_size; ii >>= 1) + { + if (tid < ii) + { data[tid] += data[tid + ii]; } __syncthreads(); } - if (tid < warp_size) { - data[tid] += data[tid + 32]; __syncwarp(); - data[tid] += data[tid + 16]; __syncwarp(); - data[tid] += data[tid + 8]; __syncwarp(); - data[tid] += data[tid + 4]; __syncwarp(); - data[tid] += data[tid + 2]; __syncwarp(); - data[tid] += data[tid + 1]; __syncwarp(); + if (tid < warp_size) + { + data[tid] += data[tid + 32]; + __syncwarp(); + data[tid] += data[tid + 16]; + __syncwarp(); + data[tid] += data[tid + 8]; + __syncwarp(); + data[tid] += data[tid + 4]; + __syncwarp(); + data[tid] += data[tid + 2]; + __syncwarp(); + data[tid] += data[tid + 1]; + __syncwarp(); } __syncthreads(); @@ -216,7 +270,8 @@ __global__ void calc_grad_with_block( data[tid] = 0; data[thread_per_block + tid] = 0; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) + { item = band_idx * n_basis_max + basis_idx; grad_1 = hpsi[item] - epsilo * spsi[item]; grad_2 = thrust::norm(grad_1); @@ -226,40 +281,57 @@ __global__ void calc_grad_with_block( __syncthreads(); // just do some parallel reduction in shared memory - for (int ii = thread_per_block >> 1; ii > warp_size; ii >>= 1) { - if (tid < ii) { + for (int ii = thread_per_block >> 1; ii > warp_size; ii >>= 1) + { + if (tid < ii) + { data[tid] += data[tid + ii]; data[thread_per_block + tid] += data[thread_per_block + tid + ii]; } __syncthreads(); } - if (tid < warp_size) { - data[tid] += data[tid + 32]; __syncwarp(); - data[tid] += data[tid + 16]; __syncwarp(); - data[tid] += data[tid + 8]; __syncwarp(); - data[tid] += data[tid + 4]; __syncwarp(); - data[tid] += data[tid + 2]; __syncwarp(); - data[tid] += data[tid + 1]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 32]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 16]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 8]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 4]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 2]; __syncwarp(); - data[thread_per_block + tid] += data[thread_per_block + tid + 1]; __syncwarp(); + if (tid < warp_size) + { + data[tid] += data[tid + 32]; + __syncwarp(); + data[tid] += data[tid + 16]; + __syncwarp(); + data[tid] += data[tid + 8]; + __syncwarp(); + data[tid] += data[tid + 4]; + __syncwarp(); + data[tid] += data[tid + 2]; + __syncwarp(); + data[tid] += data[tid + 1]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 32]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 16]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 8]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 4]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 2]; + __syncwarp(); + data[thread_per_block + tid] += data[thread_per_block + tid + 1]; + __syncwarp(); } __syncthreads(); err_st = data[0]; beta_st = data[thread_per_block]; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += thread_per_block) + { item = band_idx * n_basis_max + basis_idx; grad_1 = hpsi[item] - epsilo * spsi[item]; grad[item] = -grad_1 / prec[basis_idx] + beta_st / beta[band_idx] * grad_old[item]; } __syncthreads(); - if (tid == 0) { + if (tid == 0) + { beta[band_idx] = beta_st; err[band_idx] = sqrt(err_st); } @@ -267,88 +339,98 @@ __global__ void calc_grad_with_block( template __global__ void apply_eigenvalues_kernel( - const thrust::complex* vectors, - thrust::complex* result, - const Real* eigenvalues, - const int nbase, - const int nbase_x, - const int notconv) + const thrust::complex *vectors, + thrust::complex *result, + const Real *eigenvalues, + const int nbase, + const int nbase_x, + const int notconv) { int m = blockIdx.x; int idx = threadIdx.x + blockIdx.y * blockDim.x; - if (m < notconv && idx < nbase) { + if (m < notconv && idx < nbase) + { result[m * nbase_x + idx] = eigenvalues[m] * vectors[m * nbase_x + idx]; } } template __global__ void precondition_kernel( - thrust::complex* psi_iter, - const Real* precondition, - const Real* eigenvalues, - const int dim, - const int nbase, - const int notconv) + thrust::complex *psi_iter, + const Real *precondition, + const Real *eigenvalues, + const int dim, + const int nbase, + const int notconv) { int m = blockIdx.x; int i = threadIdx.x + blockIdx.y * blockDim.x; - if (m < notconv && i < dim) { + if (m < notconv && i < dim) + { Real x = abs(precondition[i] - eigenvalues[m]); Real pre = 0.5 * (1.0 + x + sqrt(1 + (x - 1.0) * (x - 1.0))); psi_iter[(nbase + m) * dim + i] = psi_iter[(nbase + m) * dim + i] / pre; } } -template -__device__ Real warpReduceSum(Real val) { +template __device__ Real warpReduceSum(Real val) +{ for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) + { val += __shfl_down_sync(FULL_MASK, val, offset); + } return val; } -template -__device__ Real blockReduceSum(Real val, volatile Real* shared) { +template __device__ Real blockReduceSum(Real val, volatile Real *shared) +{ int lane = threadIdx.x % WARP_SIZE; - int wid = threadIdx.x / WARP_SIZE; + int wid = threadIdx.x / WARP_SIZE; val = warpReduceSum(val); if (lane == 0) + { shared[wid] = val; + } __syncthreads(); Real sum = 0.0; - if (wid == 0) { + if (wid == 0) + { sum = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0; sum = warpReduceSum(sum); - if (lane == 0) shared[0] = sum; + if (lane == 0) + { + shared[0] = sum; + } } __syncthreads(); return shared[0]; } - template __global__ void normalize_kernel( - thrust::complex* psi_iter, - Real* psi_norm, - const int dim, - const int nbase, - const int notconv) + thrust::complex *psi_iter, + Real *psi_norm, + const int dim, + const int nbase, + const int notconv) { int m = blockIdx.x; int tid = threadIdx.x; extern __shared__ char s_char[]; - Real* shared = reinterpret_cast(s_char); + Real *shared = reinterpret_cast(s_char); Real local_sum = 0.0; // Calculate the sum for normalization - for (int i = tid; i < dim; i += thread_per_block) { + for (int i = tid; i < dim; i += thread_per_block) + { auto val = psi_iter[(nbase + m) * dim + i]; local_sum += (val * thrust::conj(val)).real(); } @@ -357,25 +439,27 @@ __global__ void normalize_kernel( Real norm = sqrt(l2_sq); // Normalize the vector - for (int i = tid; i < dim; i += thread_per_block) { + for (int i = tid; i < dim; i += thread_per_block) + { psi_iter[(nbase + m) * dim + i] /= norm; } // Store the norm if needed - if (tid == 0 && psi_norm != nullptr) { + if (tid == 0 && psi_norm != nullptr) + { psi_norm[m] = norm; } } template __global__ void refresh_hcc_scc_vcc_kernel( - const int n, - T *hcc, - T *scc, - T *vcc, - const int ldh, - const Real *eigenvalue, - const T one) + const int n, + T *hcc, + T *scc, + T *vcc, + const int ldh, + const Real *eigenvalue, + const T one) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < n) @@ -387,164 +471,216 @@ __global__ void refresh_hcc_scc_vcc_kernel( } template -void line_minimize_with_block_op::operator()(T* grad_out, - T* hgrad_out, - T* sgrad_out, - T* psi_out, - T* hpsi_out, - T* spsi_out, - const int& n_basis, - const int& n_basis_max, - const int& n_band) +void line_minimize_with_block_op::operator()( + T *grad_out, + T *hgrad_out, + T *sgrad_out, + T *psi_out, + T *hpsi_out, + T *spsi_out, + const int &n_basis, + const int &n_basis_max, + const int &n_band) { - auto A = reinterpret_cast*>(grad_out); - auto B = reinterpret_cast*>(hgrad_out); - auto C = reinterpret_cast*>(sgrad_out); - auto D = reinterpret_cast*>(psi_out); - auto E = reinterpret_cast*>(hpsi_out); - auto F = reinterpret_cast*>(spsi_out); + auto A = reinterpret_cast *>(grad_out); + auto B = reinterpret_cast *>(hgrad_out); + auto C = reinterpret_cast *>(sgrad_out); + auto D = reinterpret_cast *>(psi_out); + auto E = reinterpret_cast *>(hpsi_out); + auto F = reinterpret_cast *>(spsi_out); line_minimize_with_block<<>>( - A, B, C, D, E, F, - n_basis, n_basis_max); + A, + B, + C, + D, + E, + F, + n_basis, + n_basis_max); CHECK_CUDA_SYNC(); } template -void calc_grad_with_block_op::operator()(const Real* prec_in, - Real* err_out, - Real* beta_out, - T* psi_out, - T* hpsi_out, - T* spsi_out, - T* grad_out, - T* grad_old_out, - const int& n_basis, - const int& n_basis_max, - const int& n_band) +void calc_grad_with_block_op::operator()( + const Real *prec_in, + Real *err_out, + Real *beta_out, + T *psi_out, + T *hpsi_out, + T *spsi_out, + T *grad_out, + T *grad_old_out, + const int &n_basis, + const int &n_basis_max, + const int &n_band) { - auto A = reinterpret_cast*>(psi_out); - auto B = reinterpret_cast*>(hpsi_out); - auto C = reinterpret_cast*>(spsi_out); - auto D = reinterpret_cast*>(grad_out); - auto E = reinterpret_cast*>(grad_old_out); + auto A = reinterpret_cast *>(psi_out); + auto B = reinterpret_cast *>(hpsi_out); + auto C = reinterpret_cast *>(spsi_out); + auto D = reinterpret_cast *>(grad_out); + auto E = reinterpret_cast *>(grad_old_out); calc_grad_with_block<<>>( - prec_in, err_out, beta_out, - A, B, C, D, E, - n_basis, n_basis_max); + prec_in, + err_out, + beta_out, + A, + B, + C, + D, + E, + n_basis, + n_basis_max); CHECK_CUDA_SYNC(); } template -void apply_eigenvalues_op::operator()(const int& nbase, - const int& nbase_x, - const int& notconv, - T* result, - const T* vectors, - const Real* eigenvalues) +void apply_eigenvalues_op::operator()( + const int &nbase, + const int &nbase_x, + const int ¬conv, + T *result, + const T *vectors, + const Real *eigenvalues) { const int threads_per_block = 256; const int blocks_per_grid_y = (nbase + threads_per_block - 1) / threads_per_block; dim3 grid(notconv, blocks_per_grid_y); - auto vec_complex = reinterpret_cast*>(vectors); - auto res_complex = reinterpret_cast*>(result); + auto vec_complex = reinterpret_cast *>(vectors); + auto res_complex = reinterpret_cast *>(result); apply_eigenvalues_kernel<<>>( - vec_complex, res_complex, eigenvalues, nbase, nbase_x, notconv); + vec_complex, + res_complex, + eigenvalues, + nbase, + nbase_x, + notconv); CHECK_CUDA_SYNC(); } template -void precondition_op::operator()(const int& dim, - T* psi_iter, - const int& nbase, - const int& notconv, - const Real* precondition, - const Real* eigenvalues) +void precondition_op::operator()( + const int &dim, + T *psi_iter, + const int &nbase, + const int ¬conv, + const Real *precondition, + const Real *eigenvalues) { const int threads_per_block = 256; const int blocks_per_grid_y = (dim + threads_per_block - 1) / threads_per_block; dim3 grid(notconv, blocks_per_grid_y); - auto psi_complex = reinterpret_cast*>(psi_iter); + auto psi_complex = reinterpret_cast *>(psi_iter); precondition_kernel<<>>( - psi_complex, precondition, eigenvalues, dim, nbase, notconv); + psi_complex, + precondition, + eigenvalues, + dim, + nbase, + notconv); CHECK_CUDA_SYNC(); } template -void normalize_op::operator()(const int& dim, - T* psi_iter, - const int& nbase, - const int& notconv, - Real* psi_norm) +void normalize_op::operator()( + const int &dim, + T *psi_iter, + const int &nbase, + const int ¬conv, + Real *psi_norm) { - auto psi_complex = reinterpret_cast*>(psi_iter); + auto psi_complex = reinterpret_cast *>(psi_iter); int sharedMemSize = (thread_per_block / WARP_SIZE) * sizeof(Real); normalize_kernel<<>>( - psi_complex, psi_norm, dim, nbase, notconv); + psi_complex, + psi_norm, + dim, + nbase, + notconv); CHECK_CUDA_SYNC(); } template <> -void refresh_hcc_scc_vcc_op::operator()(const int &n, - double *hcc, - double *scc, - double *vcc, - const int &ldh, - const double *eigenvalue, - const double& one) +void refresh_hcc_scc_vcc_op::operator()( + const int &n, + double *hcc, + double *scc, + double *vcc, + const int &ldh, + const double *eigenvalue, + const double &one) { int thread = 512; int block = (n + thread - 1) / thread; - refresh_hcc_scc_vcc_kernel <<>> (n, hcc, scc, vcc, ldh, eigenvalue, one); + refresh_hcc_scc_vcc_kernel<<>>( + n, + hcc, + scc, + vcc, + ldh, + eigenvalue, + one); CHECK_CUDA_SYNC(); } template <> -void refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>::operator()(const int &n, - std::complex *hcc, - std::complex *scc, - std::complex *vcc, - const int &ldh, - const float *eigenvalue, - const std::complex& one) +void refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>::operator()( + const int &n, + std::complex *hcc, + std::complex *scc, + std::complex *vcc, + const int &ldh, + const float *eigenvalue, + const std::complex &one) { int thread = 512; int block = (n + thread - 1) / thread; - refresh_hcc_scc_vcc_kernel, float> <<>> (n, reinterpret_cast*>(hcc), - reinterpret_cast*>(scc), reinterpret_cast*>(vcc), ldh, eigenvalue, - thrust::complex(one)); + refresh_hcc_scc_vcc_kernel, float><<>>( + n, + reinterpret_cast *>(hcc), + reinterpret_cast *>(scc), + reinterpret_cast *>(vcc), + ldh, + eigenvalue, + thrust::complex(one)); CHECK_CUDA_SYNC(); } template <> -void refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>::operator()(const int &n, - std::complex *hcc, - std::complex *scc, - std::complex *vcc, - const int &ldh, - const double *eigenvalue, - const std::complex& one) +void refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>::operator()( + const int &n, + std::complex *hcc, + std::complex *scc, + std::complex *vcc, + const int &ldh, + const double *eigenvalue, + const std::complex &one) { int thread = 512; int block = (n + thread - 1) / thread; - refresh_hcc_scc_vcc_kernel, double> <<>> (n, reinterpret_cast*>(hcc), - reinterpret_cast*>(scc), reinterpret_cast*>(vcc), ldh, eigenvalue, - thrust::complex(one)); + refresh_hcc_scc_vcc_kernel, double><<>>( + n, + reinterpret_cast *>(hcc), + reinterpret_cast *>(scc), + reinterpret_cast *>(vcc), + ldh, + eigenvalue, + thrust::complex(one)); CHECK_CUDA_SYNC(); } @@ -565,4 +701,4 @@ template struct normalize_op; template struct refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>; template struct refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>; template struct refresh_hcc_scc_vcc_op; -} +} // namespace hsolver diff --git a/source/source_hsolver/kernels/cuda/diag_cusolver.cu b/source/source_hsolver/kernels/cuda/diag_cusolver.cu index 52e206483a2..faf265a643c 100644 --- a/source/source_hsolver/kernels/cuda/diag_cusolver.cu +++ b/source/source_hsolver/kernels/cuda/diag_cusolver.cu @@ -1,13 +1,15 @@ -#include #include "diag_cusolver.cuh" #include "source_base/module_device/device_check.h" -Diag_Cusolver_gvd::Diag_Cusolver_gvd(){ -// step 1: create cusolver/cublas handle +#include + +Diag_Cusolver_gvd::Diag_Cusolver_gvd() +{ + // step 1: create cusolver/cublas handle cusolverH = NULL; - CHECK_CUSOLVER( cusolverDnCreate(&cusolverH) ); + CHECK_CUSOLVER(cusolverDnCreate(&cusolverH)); - itype = CUSOLVER_EIG_TYPE_1; // A*x = (lambda)*B*x + itype = CUSOLVER_EIG_TYPE_1; // A*x = (lambda)*B*x jobz = CUSOLVER_EIG_MODE_VECTOR; // compute eigenvalues and eigenvectors. // LCAO supplies the authoritative upper triangle of the Hermitian H/S matrices. uplo = CUBLAS_FILL_MODE_UPPER; @@ -28,156 +30,241 @@ Diag_Cusolver_gvd::Diag_Cusolver_gvd(){ is_init = 0; } -void Diag_Cusolver_gvd::finalize(){ -// free resources and destroy - if (d_A ) {CHECK_CUDA( cudaFree(d_A) ); d_A = NULL;} - if (d_B ) {CHECK_CUDA( cudaFree(d_B) ); d_B = NULL;} - if (d_A2 ) {CHECK_CUDA( cudaFree(d_A2) ); d_A2 = NULL;} - if (d_B2 ) {CHECK_CUDA( cudaFree(d_B2) ); d_B2 = NULL;} - if (d_W ) {CHECK_CUDA( cudaFree(d_W) ); d_W = NULL;} - if (devInfo ) {CHECK_CUDA( cudaFree(devInfo) ); devInfo = NULL;} +void Diag_Cusolver_gvd::finalize() +{ + // free resources and destroy + if (d_A) + { + CHECK_CUDA(cudaFree(d_A)); + d_A = NULL; + } + if (d_B) + { + CHECK_CUDA(cudaFree(d_B)); + d_B = NULL; + } + if (d_A2) + { + CHECK_CUDA(cudaFree(d_A2)); + d_A2 = NULL; + } + if (d_B2) + { + CHECK_CUDA(cudaFree(d_B2)); + d_B2 = NULL; + } + if (d_W) + { + CHECK_CUDA(cudaFree(d_W)); + d_W = NULL; + } + if (devInfo) + { + CHECK_CUDA(cudaFree(devInfo)); + devInfo = NULL; + } } -Diag_Cusolver_gvd::~Diag_Cusolver_gvd(){ +Diag_Cusolver_gvd::~Diag_Cusolver_gvd() +{ finalize(); - if (cusolverH) {CHECK_CUSOLVER( cusolverDnDestroy(cusolverH) ); cusolverH = NULL;} - //CHECK_CUDA( cudaDeviceReset() ); + if (cusolverH) + { + CHECK_CUSOLVER(cusolverDnDestroy(cusolverH)); + cusolverH = NULL; + } + // CHECK_CUDA( cudaDeviceReset() ); } - -void Diag_Cusolver_gvd::init_double(int N){ -// step 2: Malloc A and B on device +void Diag_Cusolver_gvd::init_double(int N) +{ + // step 2: Malloc A and B on device m = lda = N; - CHECK_CUDA( cudaMalloc ((void**)&d_A, sizeof(double) * lda * m) ); - CHECK_CUDA( cudaMalloc ((void**)&d_B, sizeof(double) * lda * m) ); - CHECK_CUDA( cudaMalloc ((void**)&d_W, sizeof(double) * m) ); - CHECK_CUDA( cudaMalloc ((void**)&devInfo, sizeof(int)) ); + CHECK_CUDA(cudaMalloc((void **)&d_A, sizeof(double) * lda * m)); + CHECK_CUDA(cudaMalloc((void **)&d_B, sizeof(double) * lda * m)); + CHECK_CUDA(cudaMalloc((void **)&d_W, sizeof(double) * m)); + CHECK_CUDA(cudaMalloc((void **)&devInfo, sizeof(int))); } -void Diag_Cusolver_gvd::init_complex(int N){ -// step 2: Malloc A and B on device +void Diag_Cusolver_gvd::init_complex(int N) +{ + // step 2: Malloc A and B on device m = lda = N; - CHECK_CUDA( cudaMalloc ((void**)&d_A2, sizeof(cuDoubleComplex) * lda * m) ); - CHECK_CUDA( cudaMalloc ((void**)&d_B2, sizeof(cuDoubleComplex) * lda * m) ); - CHECK_CUDA( cudaMalloc ((void**)&d_W, sizeof(double) * m) ); - CHECK_CUDA( cudaMalloc ((void**)&devInfo, sizeof(int)) ); + CHECK_CUDA(cudaMalloc((void **)&d_A2, sizeof(cuDoubleComplex) * lda * m)); + CHECK_CUDA(cudaMalloc((void **)&d_B2, sizeof(cuDoubleComplex) * lda * m)); + CHECK_CUDA(cudaMalloc((void **)&d_W, sizeof(double) * m)); + CHECK_CUDA(cudaMalloc((void **)&devInfo, sizeof(int))); } - -void Diag_Cusolver_gvd::Dngvd_double(int N, int M, double *A, double *B, double *W, double *V){ + +void Diag_Cusolver_gvd::Dngvd_double( + int N, + int M, + double *A, + double *B, + double *W, + double *V) +{ // copy A, B to the GPU - assert(N == M); - if (M != m) { - this->finalize(); - this->init_double(M); - } - CHECK_CUDA( cudaMemcpy(d_A, A, sizeof(double) * lda * m, cudaMemcpyHostToDevice) ); - CHECK_CUDA( cudaMemcpy(d_B, B, sizeof(double) * lda * m, cudaMemcpyHostToDevice) ); + assert(N == M); + if (M != m) + { + this->finalize(); + this->init_double(M); + } + CHECK_CUDA(cudaMemcpy( + d_A, + A, + sizeof(double) * lda * m, + cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy( + d_B, + B, + sizeof(double) * lda * m, + cudaMemcpyHostToDevice)); // Query working space of sygvd // The helper functions below can calculate the sizes needed for pre-allocated buffer. // The S and D data types are real valued single and double precision, respectively. // The C and Z data types are complex valued single and double precision, respectively. - CHECK_CUSOLVER(cusolverDnDsygvd_bufferSize( - cusolverH, - itype, - jobz, - uplo, - m, - d_A, - lda, - d_B, - lda, - d_W, - &lwork - )); - CHECK_CUDA( cudaMalloc((void**)&d_work, sizeof(double)*lwork) ); + CHECK_CUSOLVER(cusolverDnDsygvd_bufferSize( + cusolverH, + itype, + jobz, + uplo, + m, + d_A, + lda, + d_B, + lda, + d_W, + &lwork)); + CHECK_CUDA(cudaMalloc((void **)&d_work, sizeof(double) * lwork)); // compute spectrum of (A,B) - CHECK_CUSOLVER(cusolverDnDsygvd( - cusolverH, - itype, - jobz, - uplo, - m, - d_A, - lda, - d_B, - lda, - d_W, - d_work, - lwork, - devInfo - )); - CHECK_CUDA( cudaDeviceSynchronize() ); + CHECK_CUSOLVER(cusolverDnDsygvd( + cusolverH, + itype, + jobz, + uplo, + m, + d_A, + lda, + d_B, + lda, + d_W, + d_work, + lwork, + devInfo)); + CHECK_CUDA(cudaDeviceSynchronize()); // copy (W, V) to the cpu root - CHECK_CUDA( cudaMemcpy(W, d_W, sizeof(double)*m, cudaMemcpyDeviceToHost) ); - CHECK_CUDA( cudaMemcpy(V, d_A, sizeof(double)*lda*m, cudaMemcpyDeviceToHost) ); - CHECK_CUDA( cudaMemcpy(&info_gpu, devInfo, sizeof(int), cudaMemcpyDeviceToHost) ); - assert(0 == info_gpu); + CHECK_CUDA(cudaMemcpy( + W, + d_W, + sizeof(double) * m, + cudaMemcpyDeviceToHost)); + CHECK_CUDA(cudaMemcpy( + V, + d_A, + sizeof(double) * lda * m, + cudaMemcpyDeviceToHost)); + CHECK_CUDA(cudaMemcpy( + &info_gpu, + devInfo, + sizeof(int), + cudaMemcpyDeviceToHost)); + assert(0 == info_gpu); // free the buffer - if (d_work ) CHECK_CUDA( cudaFree(d_work) ); - + if (d_work) + { + CHECK_CUDA(cudaFree(d_work)); + } } - -void Diag_Cusolver_gvd::Dngvd_complex(int N, int M, std::complex *A, std::complex *B, double *W, std::complex *V){ +void Diag_Cusolver_gvd::Dngvd_complex( + int N, + int M, + std::complex *A, + std::complex *B, + double *W, + std::complex *V) +{ // copy A, B to the GPU - assert(N == M); - if (M != m) { - this->finalize(); - this->init_complex(M); - } - CHECK_CUDA( cudaMemcpy(d_A2, A, sizeof(cuDoubleComplex) * lda * m, cudaMemcpyHostToDevice) ); - CHECK_CUDA( cudaMemcpy(d_B2, B, sizeof(cuDoubleComplex) * lda * m, cudaMemcpyHostToDevice) ); + assert(N == M); + if (M != m) + { + this->finalize(); + this->init_complex(M); + } + CHECK_CUDA(cudaMemcpy( + d_A2, + A, + sizeof(cuDoubleComplex) * lda * m, + cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy( + d_B2, + B, + sizeof(cuDoubleComplex) * lda * m, + cudaMemcpyHostToDevice)); // Query working space of Zhegvd // The helper functions below can calculate the sizes needed for pre-allocated buffer. // The S and D data types are real valued single and double precision, respectively. // The C and Z data types are complex valued single and double precision, respectively. - CHECK_CUSOLVER( - cusolverDnZhegvd_bufferSize( - cusolverH, - itype, - jobz, - uplo, - m, - d_A2, - lda, - d_B2, - lda, - d_W, - &lwork) - ); - CHECK_CUDA( cudaMalloc((void**)&d_work2, sizeof(cuDoubleComplex)*lwork) ); + CHECK_CUSOLVER(cusolverDnZhegvd_bufferSize( + cusolverH, + itype, + jobz, + uplo, + m, + d_A2, + lda, + d_B2, + lda, + d_W, + &lwork)); + CHECK_CUDA(cudaMalloc((void **)&d_work2, sizeof(cuDoubleComplex) * lwork)); // compute spectrum of (A,B) - CHECK_CUSOLVER( - cusolverDnZhegvd( - cusolverH, - itype, - jobz, - uplo, - m, - d_A2, - lda, - d_B2, - lda, - d_W, - d_work2, - lwork, - devInfo) - ); - CHECK_CUDA( cudaDeviceSynchronize() ); + CHECK_CUSOLVER(cusolverDnZhegvd( + cusolverH, + itype, + jobz, + uplo, + m, + d_A2, + lda, + d_B2, + lda, + d_W, + d_work2, + lwork, + devInfo)); + CHECK_CUDA(cudaDeviceSynchronize()); // copy (W, V) to the cpu root - CHECK_CUDA( cudaMemcpy(W, d_W, sizeof(double)*m, cudaMemcpyDeviceToHost) ); - CHECK_CUDA( cudaMemcpy(V, d_A2, sizeof(std::complex)*lda*m, cudaMemcpyDeviceToHost) ); - CHECK_CUDA( cudaMemcpy(&info_gpu, devInfo, sizeof(int), cudaMemcpyDeviceToHost) ); - assert(0 == info_gpu); + CHECK_CUDA(cudaMemcpy( + W, + d_W, + sizeof(double) * m, + cudaMemcpyDeviceToHost)); + CHECK_CUDA(cudaMemcpy( + V, + d_A2, + sizeof(std::complex) * lda * m, + cudaMemcpyDeviceToHost)); + CHECK_CUDA(cudaMemcpy( + &info_gpu, + devInfo, + sizeof(int), + cudaMemcpyDeviceToHost)); + assert(0 == info_gpu); // free the buffer - if (d_work2 ) CHECK_CUDA( cudaFree(d_work2) ); + if (d_work2) + { + CHECK_CUDA(cudaFree(d_work2)); + } } diff --git a/source/source_hsolver/kernels/cuda/diag_cusolver.cuh b/source/source_hsolver/kernels/cuda/diag_cusolver.cuh index e47d43be504..e3c3bb57532 100644 --- a/source/source_hsolver/kernels/cuda/diag_cusolver.cuh +++ b/source/source_hsolver/kernels/cuda/diag_cusolver.cuh @@ -1,29 +1,29 @@ #ifndef DIAG_CUSOLVER_CUH #define DIAG_CUSOLVER_CUH -#include #include - +#include #include #include - /// -/// The diag_cusolver contains subroutines for generalized eigenvector problems of dense matrice in lcao on a single GPU. -/// Cusolver apis, i.e. cusolverDnDsygvd and cusolverDnZhegvd, are wrapped with respect to different data types for easy calling. -/// For details of low-level methods relevent, one can refer to [cuSOLVER documentation](https://docs.nvidia.com/cuda/cusolver/index.html#cuSolverDN-lt-t-gt-sygvd). +/// The diag_cusolver contains subroutines for generalized eigenvector problems of dense matrice in lcao on a single +/// GPU. Cusolver apis, i.e. cusolverDnDsygvd and cusolverDnZhegvd, are wrapped with respect to different data types for +/// easy calling. For details of low-level methods relevent, one can refer to [cuSOLVER +/// documentation](https://docs.nvidia.com/cuda/cusolver/index.html#cuSolverDN-lt-t-gt-sygvd). /// // Xu Shu add 2022-03-23 // -class Diag_Cusolver_gvd{ +class Diag_Cusolver_gvd +{ -//------------------- -// private variables -//------------------- + //------------------- + // private variables + //------------------- cusolverDnHandle_t cusolverH = nullptr; - cusolverEigType_t itype = CUSOLVER_EIG_TYPE_1; //problem type: A*x = (lambda)*B*x + cusolverEigType_t itype = CUSOLVER_EIG_TYPE_1; // problem type: A*x = (lambda)*B*x cusolverEigMode_t jobz = CUSOLVER_EIG_MODE_NOVECTOR; // compute eigenvalues and eigenvectors. cublasFillMode_t uplo = CUBLAS_FILL_MODE_LOWER; @@ -41,50 +41,84 @@ class Diag_Cusolver_gvd{ double *d_W = nullptr; int *devInfo = nullptr; - int lwork = 0; + int lwork = 0; int info_gpu = 0; -// subroutines that are related to initializing the class: -// - init_double : initializing relevant double type data structures and gpu apis' handle and memory -// - init_complex : initializing relevant complex type data structures and gpu apis' handle and memory -// Input Parameters -// N: the dimension of the matrix + // subroutines that are related to initializing the class: + // - init_double : initializing relevant double type data structures and gpu apis' handle and memory + // - init_complex : initializing relevant complex type data structures and gpu apis' handle and memory + // Input Parameters + // N: the dimension of the matrix void init_double(int N); void init_complex(int N); - void finalize(); // for recycling the usage of the static class Diag_Cusolver_gvd -public: - - int is_init = 0; // For expensive gpu initialization only once when using cusolver for lcao + void finalize(); // for recycling the usage of the static class Diag_Cusolver_gvd + public: + int is_init = 0; // For expensive gpu initialization only once when using cusolver for lcao Diag_Cusolver_gvd(); ~Diag_Cusolver_gvd(); -// subroutines that are related to calculating generalized eigenvalues and eigenvectors for dense matrix pairs: -// - Dngvd_double : dense double type matrix -// - Dngvd_complex : dense complex type matrix -// Input Parameters -// N: the number of rows of the matrix -// M: the number of cols of the matrix -// A: the hermitian matrix A in A x=lambda B (column major) -// B: the SPD matrix B in A x=lambda B (column major) -// Output Parameter -// W: generalized eigenvalues -// V: generalized eigenvectors (column major) - - void Dngvd_double(int N, int M, double *A, double *B, double *W, double *V); - void Dngvd_complex(int N, int M, std::complex *A, std::complex *B, double *W, std::complex *V); - - void Dngvd(int N, int M, double *A, double *B, double *W, double *V) + // subroutines that are related to calculating generalized eigenvalues and eigenvectors for dense matrix pairs: + // - Dngvd_double : dense double type matrix + // - Dngvd_complex : dense complex type matrix + // Input Parameters + // N: the number of rows of the matrix + // M: the number of cols of the matrix + // A: the hermitian matrix A in A x=lambda B (column major) + // B: the SPD matrix B in A x=lambda B (column major) + // Output Parameter + // W: generalized eigenvalues + // V: generalized eigenvectors (column major) + + void Dngvd_double( + int N, + int M, + double *A, + double *B, + double *W, + double *V); + void Dngvd_complex( + int N, + int M, + std::complex *A, + std::complex *B, + double *W, + std::complex *V); + + void Dngvd( + int N, + int M, + double *A, + double *B, + double *W, + double *V) { - return Dngvd_double(N, M, A, B, W, V); + return Dngvd_double( + N, + M, + A, + B, + W, + V); }; - void Dngvd(int N, int M, std::complex *A, std::complex *B, double *W, std::complex *V) + void Dngvd( + int N, + int M, + std::complex *A, + std::complex *B, + double *W, + std::complex *V) { - return Dngvd_complex(N, M, A, B, W, V); + return Dngvd_complex( + N, + M, + A, + B, + W, + V); }; - }; #endif diff --git a/source/source_hsolver/kernels/cuda/diag_cusolvermp.cu b/source/source_hsolver/kernels/cuda/diag_cusolvermp.cu index 34c97b20508..e9ab678c01e 100644 --- a/source/source_hsolver/kernels/cuda/diag_cusolvermp.cu +++ b/source/source_hsolver/kernels/cuda/diag_cusolvermp.cu @@ -1,7 +1,7 @@ #ifdef __CUSOLVERMP #include "diag_cusolvermp.cuh" -#include "source_base/module_device/device_check.h" #include "source_base/global_function.h" +#include "source_base/module_device/device_check.h" #include @@ -9,30 +9,43 @@ extern "C" { #include "source_hsolver/module_genelpa/cblacs.h" } -#include -#include #include "source_base/module_device/device.h" #include "source_base/module_device/device_check.h" +#include +#include + #ifdef __USE_CAL // ============================================================================ // CAL callback functions for MPI communication // ============================================================================ -static calError_t allgather(void* src_buf, void* recv_buf, size_t size, void* data, void** request) +static calError_t allgather( + void *src_buf, + void *recv_buf, + size_t size, + void *data, + void **request) { MPI_Request req; intptr_t ptr = reinterpret_cast(data); - int err = MPI_Iallgather(src_buf, size, MPI_BYTE, recv_buf, size, MPI_BYTE, (MPI_Comm)ptr, &req); + int err = MPI_Iallgather( + src_buf, + size, + MPI_BYTE, + recv_buf, + size, + MPI_BYTE, (MPI_Comm)ptr, + &req); if (err != MPI_SUCCESS) { return CAL_ERROR; } - *request = (void*)(req); + *request = (void *)(req); return CAL_OK; } -static calError_t request_test(void* request) +static calError_t request_test(void *request) { intptr_t ptr = reinterpret_cast(request); MPI_Request req = (MPI_Request)ptr; @@ -45,17 +58,15 @@ static calError_t request_test(void* request) return completed ? CAL_OK : CAL_ERROR_INPROGRESS; } -static calError_t request_free(void* request) -{ - return CAL_OK; -} +static calError_t request_free(void *request) { return CAL_OK; } #endif // __USE_CAL template -Diag_CusolverMP_gvd::Diag_CusolverMP_gvd(const MPI_Comm mpi_comm, - const int narows, - const int nacols, - const int* desc) +Diag_CusolverMP_gvd::Diag_CusolverMP_gvd( + const MPI_Comm mpi_comm, + const int narows, + const int nacols, + const int *desc) { /// constructor implementation this->cblacs_ctxt = desc[1]; @@ -78,7 +89,12 @@ Diag_CusolverMP_gvd::Diag_CusolverMP_gvd(const MPI_Comm mpi_comm, MPI_Comm_rank(mpi_comm, &(this->globalMpiRank)); // GPU device is already bound by DeviceContext::init() in read_input.cpp int local_device_id = base_device::DeviceContext::instance().get_device_id(); - Cblacs_gridinfo(this->cblacs_ctxt, &this->nprows, &this->npcols, &this->myprow, &this->mypcol); + Cblacs_gridinfo( + this->cblacs_ctxt, + &this->nprows, + &this->npcols, + &this->myprow, + &this->mypcol); #ifdef __USE_CAL // Initialize CAL communicator @@ -86,7 +102,7 @@ Diag_CusolverMP_gvd::Diag_CusolverMP_gvd(const MPI_Comm mpi_comm, params.allgather = allgather; params.req_test = request_test; params.req_free = request_free; - params.data = (void*)(mpi_comm); + params.data = (void *)(mpi_comm); params.rank = this->globalMpiRank; params.nranks = this->globalMpiSize; params.local_device = local_device_id; @@ -94,11 +110,21 @@ Diag_CusolverMP_gvd::Diag_CusolverMP_gvd(const MPI_Comm mpi_comm, #else // Initialize NCCL communicator ncclUniqueId ncclId; - if (this->globalMpiRank == 0) { + if (this->globalMpiRank == 0) + { CHECK_NCCL(ncclGetUniqueId(&ncclId)); } - MPI_Bcast(&ncclId, sizeof(ncclId), MPI_BYTE, 0, mpi_comm); - CHECK_NCCL(ncclCommInitRank(&this->ncclComm, this->globalMpiSize, ncclId, this->globalMpiRank)); + MPI_Bcast( + &ncclId, + sizeof(ncclId), + MPI_BYTE, + 0, + mpi_comm); + CHECK_NCCL(ncclCommInitRank( + &this->ncclComm, + this->globalMpiSize, + ncclId, + this->globalMpiRank)); #endif CHECK_CUDA(cudaStreamCreate(&this->localStream)); @@ -130,35 +156,36 @@ Diag_CusolverMP_gvd::Diag_CusolverMP_gvd(const MPI_Comm mpi_comm, // So, when we use cusolvermp, we must ensure that the number of processes is equal to the number of GPUs. // In a sense, the MPI usage strategy of ABACUS must be subject to the cusolvermp. // Use ROW_MAJOR to match BLACS grid initialization (order='R' in parallel_2d.cpp) - CHECK_CUSOLVER(cusolverMpCreateDeviceGrid(cusolverMpHandle, - &this->grid, + CHECK_CUSOLVER(cusolverMpCreateDeviceGrid( + cusolverMpHandle, + &this->grid, #ifdef __USE_CAL - this->cusolverCalComm, + this->cusolverCalComm, #else - this->ncclComm, + this->ncclComm, #endif - this->nprows, - this->npcols, - CUSOLVERMP_GRID_MAPPING_ROW_MAJOR)); + this->nprows, + this->npcols, + CUSOLVERMP_GRID_MAPPING_ROW_MAJOR)); // 20240529 zhanghaochong // Actually, there should be three matrix descriptors, A matrix, B matrix, and output eigenvector matrix. // But in ABACUS the three matrices descriptors are the same. // So, I only create one matrix descriptor and use it for the three matrices. - CHECK_CUSOLVER(cusolverMpCreateMatrixDesc(&this->desc_for_cusolvermp, - this->grid, - this->datatype, - nFull, - nFull, - mb, - nb, - rsrc, - csrc, - this->lda)); + CHECK_CUSOLVER(cusolverMpCreateMatrixDesc( + &this->desc_for_cusolvermp, + this->grid, + this->datatype, + nFull, + nFull, + mb, + nb, + rsrc, + csrc, + this->lda)); } -template -Diag_CusolverMP_gvd::~Diag_CusolverMP_gvd() +template Diag_CusolverMP_gvd::~Diag_CusolverMP_gvd() { #ifdef __USE_CAL CHECK_CAL(cal_comm_barrier(this->cusolverCalComm, this->localStream)); @@ -178,90 +205,104 @@ Diag_CusolverMP_gvd::~Diag_CusolverMP_gvd() #endif } - template -int Diag_CusolverMP_gvd::generalized_eigenvector(inputT* A, inputT* B, outputT* EigenValue, inputT* EigenVector) +int Diag_CusolverMP_gvd::generalized_eigenvector( + inputT *A, + inputT *B, + outputT *EigenValue, + inputT *EigenVector) { void *d_A = NULL; void *d_B = NULL; void *d_D = NULL; void *d_Z = NULL; - CHECK_CUDA(cudaMalloc((void**)&d_A, this->n_local * this->m_local * sizeof(inputT))); - CHECK_CUDA(cudaMalloc((void**)&d_B, this->n_local * this->m_local * sizeof(inputT))); - CHECK_CUDA(cudaMalloc((void**)&d_D, this->nFull * sizeof(outputT))); - CHECK_CUDA(cudaMalloc((void**)&d_Z, this->n_local * this->m_local * sizeof(inputT))); - - CHECK_CUDA( - cudaMemcpy(d_A, (void*)A, this->n_local * this->m_local * sizeof(inputT), cudaMemcpyHostToDevice)); - CHECK_CUDA( - cudaMemcpy(d_B, (void*)B, this->n_local * this->m_local * sizeof(inputT), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMalloc((void **)&d_A, this->n_local * this->m_local * sizeof(inputT))); + CHECK_CUDA(cudaMalloc((void **)&d_B, this->n_local * this->m_local * sizeof(inputT))); + CHECK_CUDA(cudaMalloc((void **)&d_D, this->nFull * sizeof(outputT))); + CHECK_CUDA(cudaMalloc((void **)&d_Z, this->n_local * this->m_local * sizeof(inputT))); + + CHECK_CUDA(cudaMemcpy( + d_A, (void *)A, + this->n_local * this->m_local * sizeof(inputT), + cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy( + d_B, (void *)B, + this->n_local * this->m_local * sizeof(inputT), + cudaMemcpyHostToDevice)); CHECK_CUDA(cudaStreamSynchronize(this->localStream)); size_t sygvdWorkspaceInBytesOnDevice = 0; size_t sygvdWorkspaceInBytesOnHost = 0; - CHECK_CUSOLVER(cusolverMpSygvd_bufferSize(cusolverMpHandle, - CUSOLVER_EIG_TYPE_1, - CUSOLVER_EIG_MODE_VECTOR, - CUBLAS_FILL_MODE_LOWER, - this->nFull, - this->matrix_i, - this->matrix_j, - this->desc_for_cusolvermp, - this->matrix_i, - this->matrix_j, - this->desc_for_cusolvermp, - this->matrix_i, - this->matrix_j, - this->desc_for_cusolvermp, - this->datatype, - &sygvdWorkspaceInBytesOnDevice, - &sygvdWorkspaceInBytesOnHost)); + CHECK_CUSOLVER(cusolverMpSygvd_bufferSize( + cusolverMpHandle, + CUSOLVER_EIG_TYPE_1, + CUSOLVER_EIG_MODE_VECTOR, + CUBLAS_FILL_MODE_LOWER, + this->nFull, + this->matrix_i, + this->matrix_j, + this->desc_for_cusolvermp, + this->matrix_i, + this->matrix_j, + this->desc_for_cusolvermp, + this->matrix_i, + this->matrix_j, + this->desc_for_cusolvermp, + this->datatype, + &sygvdWorkspaceInBytesOnDevice, + &sygvdWorkspaceInBytesOnHost)); /* Distributed host workspace */ - void* h_sygvdWork = NULL; - h_sygvdWork = (void*)malloc(sygvdWorkspaceInBytesOnHost); + void *h_sygvdWork = NULL; + h_sygvdWork = (void *)malloc(sygvdWorkspaceInBytesOnHost); /* Distributed device workspace */ - void* d_sygvdWork = NULL; - CHECK_CUDA(cudaMalloc((void**)&d_sygvdWork, sygvdWorkspaceInBytesOnDevice)); + void *d_sygvdWork = NULL; + CHECK_CUDA(cudaMalloc((void **)&d_sygvdWork, sygvdWorkspaceInBytesOnDevice)); CHECK_CUDA(cudaMemset(d_sygvdWork, 0, sygvdWorkspaceInBytesOnDevice)); - int* d_sygvdInfo = NULL; - CHECK_CUDA(cudaMalloc((void**)&d_sygvdInfo, sizeof(int))); + int *d_sygvdInfo = NULL; + CHECK_CUDA(cudaMalloc((void **)&d_sygvdInfo, sizeof(int))); CHECK_CUDA(cudaMemset(d_sygvdInfo, 0, sizeof(int))); /* sync wait for data to arrive to device */ CHECK_CUDA(cudaStreamSynchronize(this->localStream)); - CHECK_CUSOLVER(cusolverMpSygvd(cusolverMpHandle, - CUSOLVER_EIG_TYPE_1, - CUSOLVER_EIG_MODE_VECTOR, - CUBLAS_FILL_MODE_LOWER, - this->nFull, - d_A, - this->matrix_i, - this->matrix_j, - this->desc_for_cusolvermp, - d_B, - this->matrix_i, - this->matrix_j, - this->desc_for_cusolvermp, - d_D, - d_Z, - this->matrix_i, - this->matrix_j, - this->desc_for_cusolvermp, - this->datatype, - d_sygvdWork, - sygvdWorkspaceInBytesOnDevice, - h_sygvdWork, - sygvdWorkspaceInBytesOnHost, - d_sygvdInfo)); + CHECK_CUSOLVER(cusolverMpSygvd( + cusolverMpHandle, + CUSOLVER_EIG_TYPE_1, + CUSOLVER_EIG_MODE_VECTOR, + CUBLAS_FILL_MODE_LOWER, + this->nFull, + d_A, + this->matrix_i, + this->matrix_j, + this->desc_for_cusolvermp, + d_B, + this->matrix_i, + this->matrix_j, + this->desc_for_cusolvermp, + d_D, + d_Z, + this->matrix_i, + this->matrix_j, + this->desc_for_cusolvermp, + this->datatype, + d_sygvdWork, + sygvdWorkspaceInBytesOnDevice, + h_sygvdWork, + sygvdWorkspaceInBytesOnHost, + d_sygvdInfo)); int h_sygvdInfo = 0; - CHECK_CUDA(cudaMemcpyAsync(&h_sygvdInfo, d_sygvdInfo, sizeof(int), cudaMemcpyDeviceToHost, this->localStream)); + CHECK_CUDA(cudaMemcpyAsync( + &h_sygvdInfo, + d_sygvdInfo, + sizeof(int), + cudaMemcpyDeviceToHost, + this->localStream)); /* wait for d_sygvdInfo copy */ CHECK_CUDA(cudaStreamSynchronize(this->localStream)); if (h_sygvdInfo != 0) @@ -275,11 +316,16 @@ int Diag_CusolverMP_gvd::generalized_eigenvector(inputT* A, inputT* B, o free(h_sygvdWork); - CHECK_CUDA(cudaMemcpy((void*)EigenValue, d_D, this->nFull * sizeof(outputT), cudaMemcpyDeviceToHost)); - CHECK_CUDA(cudaMemcpy((void*)EigenVector, - d_Z, - this->n_local * this->m_local * sizeof(inputT), - cudaMemcpyDeviceToHost)); + CHECK_CUDA(cudaMemcpy( + (void *)EigenValue, + d_D, + this->nFull * sizeof(outputT), + cudaMemcpyDeviceToHost)); + CHECK_CUDA(cudaMemcpy( + (void *)EigenVector, + d_Z, + this->n_local * this->m_local * sizeof(inputT), + cudaMemcpyDeviceToHost)); // 20240529 zhanghaochong // I move the free operations from destructor to here. // Because I think it is more reasonable to free the memory in the function where it is allocated. diff --git a/source/source_hsolver/kernels/cuda/diag_cusolvermp.cuh b/source/source_hsolver/kernels/cuda/diag_cusolvermp.cuh index ec5253b2b94..6fd76b0fb60 100644 --- a/source/source_hsolver/kernels/cuda/diag_cusolvermp.cuh +++ b/source/source_hsolver/kernels/cuda/diag_cusolvermp.cuh @@ -2,12 +2,12 @@ #define ABACUS_SOURCE_HSOLVER_KERNELS_CUDA_DIAG_CUSOLVERMP_CUH #include "mpi.h" +#include "source_base/macros.h" #include +#include #include #include -#include -#include "source_base/macros.h" #ifdef __USE_CAL #include @@ -15,30 +15,30 @@ #include #endif -template -class Diag_CusolverMP_gvd +template class Diag_CusolverMP_gvd { private: using outputT = typename GetTypeReal::type; public: - Diag_CusolverMP_gvd(const MPI_Comm comm, - const int narows, - const int nacols, - const int* desc); - + Diag_CusolverMP_gvd( + const MPI_Comm comm, + const int narows, + const int nacols, + const int *desc); // 20240530 zhanghaochong // Here eigen is the output rather than the input. Don't be confused by inputT. // Here inputT is the data type, the data type of EigenVector should be the same as // the data type of the input matrix. - int generalized_eigenvector(inputT* A, - inputT* B, - outputT* EigenValue, - inputT* EigenVector); + int generalized_eigenvector( + inputT *A, + inputT *B, + outputT *EigenValue, + inputT *EigenVector); ~Diag_CusolverMP_gvd(); - private: + private: int nFull; int cblacs_ctxt; int lda; diff --git a/source/source_hsolver/kernels/cuda/hegvd_op.cu b/source/source_hsolver/kernels/cuda/hegvd_op.cu index 911a6b8567a..56ed23663cc 100644 --- a/source/source_hsolver/kernels/cuda/hegvd_op.cu +++ b/source/source_hsolver/kernels/cuda/hegvd_op.cu @@ -1,8 +1,7 @@ -#include "source_hsolver/kernels/hegvd_op.h" #include "source_base/module_device/device_check.h" +#include "source_hsolver/kernels/hegvd_op.h" #include - #include namespace hsolver @@ -28,31 +27,58 @@ void destroyGpuSolverHandle() } #ifdef __LCAO -static inline -void xhegvd_wrapper( - const cublasFillMode_t& uplo, - const int& n, - double* A, const int& lda, - double* B, const int& ldb, - double* W) +static inline void xhegvd_wrapper( + const cublasFillMode_t &uplo, + const int &n, + double *A, + const int &lda, + double *B, + const int &ldb, + double *W) { // prepare some values for cusolverDnZhegvd_bufferSize - int* devInfo = nullptr; + int *devInfo = nullptr; int lwork = 0, info_gpu = 0; - double* work = nullptr; - CHECK_CUDA(cudaMalloc((void**)&devInfo, sizeof(int))); + double *work = nullptr; + CHECK_CUDA(cudaMalloc((void **)&devInfo, sizeof(int))); // calculate the sizes needed for pre-allocated buffer. - CHECK_CUSOLVER(cusolverDnDsygvd_bufferSize(cusolver_H, CUSOLVER_EIG_TYPE_1, CUSOLVER_EIG_MODE_VECTOR, uplo, n, - A, lda, B, ldb, W, &lwork)); + CHECK_CUSOLVER(cusolverDnDsygvd_bufferSize( + cusolver_H, + CUSOLVER_EIG_TYPE_1, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + A, + lda, + B, + ldb, + W, + &lwork)); // allocate memery - CHECK_CUDA(cudaMalloc((void**)&work, sizeof(double) * lwork)); + CHECK_CUDA(cudaMalloc((void **)&work, sizeof(double) * lwork)); // compute eigenvalues and eigenvectors. - CHECK_CUSOLVER(cusolverDnDsygvd(cusolver_H, CUSOLVER_EIG_TYPE_1, CUSOLVER_EIG_MODE_VECTOR, uplo, n, - A, lda, B, ldb, W, work, lwork, devInfo)); - - CHECK_CUDA(cudaMemcpy(&info_gpu, devInfo, sizeof(int), cudaMemcpyDeviceToHost)); + CHECK_CUSOLVER(cusolverDnDsygvd( + cusolver_H, + CUSOLVER_EIG_TYPE_1, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + A, + lda, + B, + ldb, + W, + work, + lwork, + devInfo)); + + CHECK_CUDA(cudaMemcpy( + &info_gpu, + devInfo, + sizeof(int), + cudaMemcpyDeviceToHost)); assert(0 == info_gpu); // free the buffer CHECK_CUDA(cudaFree(work)); @@ -60,64 +86,116 @@ void xhegvd_wrapper( } #endif -static inline -void xhegvd_wrapper ( - const cublasFillMode_t& uplo, - const int& n, - std::complex * A, const int& lda, - std::complex * B, const int& ldb, - float * W) +static inline void xhegvd_wrapper( + const cublasFillMode_t &uplo, + const int &n, + std::complex *A, + const int &lda, + std::complex *B, + const int &ldb, + float *W) { // prepare some values for cusolverDnZhegvd_bufferSize - int * devInfo = nullptr; + int *devInfo = nullptr; int lwork = 0, info_gpu = 0; - float2 * work = nullptr; - CHECK_CUDA(cudaMalloc((void**)&devInfo, sizeof(int))); + float2 *work = nullptr; + CHECK_CUDA(cudaMalloc((void **)&devInfo, sizeof(int))); // calculate the sizes needed for pre-allocated buffer. - CHECK_CUSOLVER(cusolverDnChegvd_bufferSize(cusolver_H, CUSOLVER_EIG_TYPE_1, CUSOLVER_EIG_MODE_VECTOR, uplo, n, - reinterpret_cast(A), lda, - reinterpret_cast(B), ldb, W, &lwork)); + CHECK_CUSOLVER(cusolverDnChegvd_bufferSize( + cusolver_H, + CUSOLVER_EIG_TYPE_1, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + reinterpret_cast(A), + lda, + reinterpret_cast(B), + ldb, + W, + &lwork)); // allocate memery - CHECK_CUDA(cudaMalloc((void**)&work, sizeof(float2) * lwork)); + CHECK_CUDA(cudaMalloc((void **)&work, sizeof(float2) * lwork)); // compute eigenvalues and eigenvectors. - CHECK_CUSOLVER(cusolverDnChegvd(cusolver_H, CUSOLVER_EIG_TYPE_1, CUSOLVER_EIG_MODE_VECTOR, uplo, n, - reinterpret_cast(A), lda, reinterpret_cast(B), ldb, W, work, lwork, devInfo)); - - CHECK_CUDA(cudaMemcpy(&info_gpu, devInfo, sizeof(int), cudaMemcpyDeviceToHost)); + CHECK_CUSOLVER(cusolverDnChegvd( + cusolver_H, + CUSOLVER_EIG_TYPE_1, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + reinterpret_cast(A), + lda, + reinterpret_cast(B), + ldb, + W, + work, + lwork, + devInfo)); + + CHECK_CUDA(cudaMemcpy( + &info_gpu, + devInfo, + sizeof(int), + cudaMemcpyDeviceToHost)); assert(0 == info_gpu); // free the buffer CHECK_CUDA(cudaFree(work)); CHECK_CUDA(cudaFree(devInfo)); } -static inline -void xhegvd_wrapper ( - const cublasFillMode_t& uplo, - const int& n, - std::complex * A, const int& lda, - std::complex * B, const int& ldb, - double * W) +static inline void xhegvd_wrapper( + const cublasFillMode_t &uplo, + const int &n, + std::complex *A, + const int &lda, + std::complex *B, + const int &ldb, + double *W) { // prepare some values for cusolverDnZhegvd_bufferSize - int * devInfo = nullptr; + int *devInfo = nullptr; int lwork = 0, info_gpu = 0; - double2 * work = nullptr; - CHECK_CUDA(cudaMalloc((void**)&devInfo, sizeof(int))); + double2 *work = nullptr; + CHECK_CUDA(cudaMalloc((void **)&devInfo, sizeof(int))); // calculate the sizes needed for pre-allocated buffer. - CHECK_CUSOLVER(cusolverDnZhegvd_bufferSize(cusolver_H, CUSOLVER_EIG_TYPE_1, CUSOLVER_EIG_MODE_VECTOR, uplo, n, - reinterpret_cast(A), lda, - reinterpret_cast(B), ldb, W, &lwork)); + CHECK_CUSOLVER(cusolverDnZhegvd_bufferSize( + cusolver_H, + CUSOLVER_EIG_TYPE_1, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + reinterpret_cast(A), + lda, + reinterpret_cast(B), + ldb, + W, + &lwork)); // allocate memery - CHECK_CUDA(cudaMalloc((void**)&work, sizeof(double2) * lwork)); + CHECK_CUDA(cudaMalloc((void **)&work, sizeof(double2) * lwork)); // compute eigenvalues and eigenvectors. - CHECK_CUSOLVER(cusolverDnZhegvd(cusolver_H, CUSOLVER_EIG_TYPE_1, CUSOLVER_EIG_MODE_VECTOR, uplo, n, - reinterpret_cast(A), lda, reinterpret_cast(B), ldb, W, work, lwork, devInfo)); - - CHECK_CUDA(cudaMemcpy(&info_gpu, devInfo, sizeof(int), cudaMemcpyDeviceToHost)); + CHECK_CUSOLVER(cusolverDnZhegvd( + cusolver_H, + CUSOLVER_EIG_TYPE_1, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + reinterpret_cast(A), + lda, + reinterpret_cast(B), + ldb, + W, + work, + lwork, + devInfo)); + + CHECK_CUDA(cudaMemcpy( + &info_gpu, + devInfo, + sizeof(int), + cudaMemcpyDeviceToHost)); assert(0 == info_gpu); // free the buffer CHECK_CUDA(cudaFree(work)); @@ -125,142 +203,221 @@ void xhegvd_wrapper ( } #ifdef __LCAO -static inline -void xheevd_wrapper( - const cublasFillMode_t& uplo, - const int& n, - double* A, const int& lda, - double* W) +static inline void xheevd_wrapper( + const cublasFillMode_t &uplo, + const int &n, + double *A, + const int &lda, + double *W) { // prepare some values for cusolverDnZhegvd_bufferSize - int* devInfo = nullptr; + int *devInfo = nullptr; int lwork = 0, info_gpu = 0; - double* work = nullptr; - CHECK_CUDA(cudaMalloc((void**)&devInfo, sizeof(int))); + double *work = nullptr; + CHECK_CUDA(cudaMalloc((void **)&devInfo, sizeof(int))); // calculate the sizes needed for pre-allocated buffer. - CHECK_CUSOLVER(cusolverDnDsyevd_bufferSize(cusolver_H, CUSOLVER_EIG_MODE_VECTOR, uplo, n, - A, lda, W, &lwork)); + CHECK_CUSOLVER(cusolverDnDsyevd_bufferSize( + cusolver_H, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + A, + lda, + W, + &lwork)); // allocate memery - CHECK_CUDA(cudaMalloc((void**)&work, sizeof(double) * lwork)); + CHECK_CUDA(cudaMalloc((void **)&work, sizeof(double) * lwork)); // compute eigenvalues and eigenvectors. - CHECK_CUSOLVER(cusolverDnDsyevd(cusolver_H, CUSOLVER_EIG_MODE_VECTOR, uplo, n, A, lda, W, work, lwork, devInfo)); - - CHECK_CUDA(cudaMemcpy(&info_gpu, devInfo, sizeof(int), cudaMemcpyDeviceToHost)); + CHECK_CUSOLVER(cusolverDnDsyevd( + cusolver_H, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + A, + lda, + W, + work, + lwork, + devInfo)); + + CHECK_CUDA(cudaMemcpy( + &info_gpu, + devInfo, + sizeof(int), + cudaMemcpyDeviceToHost)); assert(0 == info_gpu); CHECK_CUDA(cudaFree(work)); CHECK_CUDA(cudaFree(devInfo)); } #endif -static inline -void xheevd_wrapper ( - const cublasFillMode_t& uplo, - const int& n, - std::complex * A, const int& lda, - float * W) +static inline void xheevd_wrapper( + const cublasFillMode_t &uplo, + const int &n, + std::complex *A, + const int &lda, + float *W) { // prepare some values for cusolverDnZhegvd_bufferSize - int * devInfo = nullptr; + int *devInfo = nullptr; int lwork = 0, info_gpu = 0; - float2 * work = nullptr; - CHECK_CUDA(cudaMalloc((void**)&devInfo, sizeof(int))); + float2 *work = nullptr; + CHECK_CUDA(cudaMalloc((void **)&devInfo, sizeof(int))); // calculate the sizes needed for pre-allocated buffer. - CHECK_CUSOLVER(cusolverDnCheevd_bufferSize(cusolver_H, CUSOLVER_EIG_MODE_VECTOR, uplo, n, - reinterpret_cast(A), lda, W, &lwork)); + CHECK_CUSOLVER(cusolverDnCheevd_bufferSize( + cusolver_H, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + reinterpret_cast(A), + lda, + W, + &lwork)); // allocate memery - CHECK_CUDA(cudaMalloc((void**)&work, sizeof(float2) * lwork)); + CHECK_CUDA(cudaMalloc((void **)&work, sizeof(float2) * lwork)); // compute eigenvalues and eigenvectors. - CHECK_CUSOLVER(cusolverDnCheevd(cusolver_H, CUSOLVER_EIG_MODE_VECTOR, uplo, n, reinterpret_cast(A), lda, W, work, lwork, devInfo)); - - CHECK_CUDA(cudaMemcpy(&info_gpu, devInfo, sizeof(int), cudaMemcpyDeviceToHost)); + CHECK_CUSOLVER(cusolverDnCheevd( + cusolver_H, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + reinterpret_cast(A), + lda, + W, + work, + lwork, + devInfo)); + + CHECK_CUDA(cudaMemcpy( + &info_gpu, + devInfo, + sizeof(int), + cudaMemcpyDeviceToHost)); assert(0 == info_gpu); CHECK_CUDA(cudaFree(work)); CHECK_CUDA(cudaFree(devInfo)); } -static inline -void xheevd_wrapper ( - const cublasFillMode_t& uplo, - const int& n, - std::complex * A, const int& lda, - double * W) +static inline void xheevd_wrapper( + const cublasFillMode_t &uplo, + const int &n, + std::complex *A, + const int &lda, + double *W) { // prepare some values for cusolverDnZhegvd_bufferSize - int * devInfo = nullptr; + int *devInfo = nullptr; int lwork = 0, info_gpu = 0; - double2 * work = nullptr; - CHECK_CUDA(cudaMalloc((void**)&devInfo, sizeof(int))); + double2 *work = nullptr; + CHECK_CUDA(cudaMalloc((void **)&devInfo, sizeof(int))); // calculate the sizes needed for pre-allocated buffer. - CHECK_CUSOLVER(cusolverDnZheevd_bufferSize(cusolver_H, CUSOLVER_EIG_MODE_VECTOR, uplo, n, - reinterpret_cast(A), lda, W, &lwork)); + CHECK_CUSOLVER(cusolverDnZheevd_bufferSize( + cusolver_H, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + reinterpret_cast(A), + lda, + W, + &lwork)); // allocate memery - CHECK_CUDA(cudaMalloc((void**)&work, sizeof(double2) * lwork)); + CHECK_CUDA(cudaMalloc((void **)&work, sizeof(double2) * lwork)); // compute eigenvalues and eigenvectors. - CHECK_CUSOLVER(cusolverDnZheevd(cusolver_H, CUSOLVER_EIG_MODE_VECTOR, uplo, n, - reinterpret_cast(A), lda, W, work, lwork, devInfo)); - - CHECK_CUDA(cudaMemcpy(&info_gpu, devInfo, sizeof(int), cudaMemcpyDeviceToHost)); + CHECK_CUSOLVER(cusolverDnZheevd( + cusolver_H, + CUSOLVER_EIG_MODE_VECTOR, + uplo, + n, + reinterpret_cast(A), + lda, + W, + work, + lwork, + devInfo)); + + CHECK_CUDA(cudaMemcpy( + &info_gpu, + devInfo, + sizeof(int), + cudaMemcpyDeviceToHost)); assert(0 == info_gpu); CHECK_CUDA(cudaFree(work)); CHECK_CUDA(cudaFree(devInfo)); } -template -struct hegvd_op +template struct hegvd_op { using Real = typename GetTypeReal::type; - void operator()(const base_device::DEVICE_GPU* d, - const int nstart, - const int ldh, - const T* A, // hcc - T* B, // scc - Real* W, // eigenvalue - T* V) + void operator()( + const base_device::DEVICE_GPU *d, + const int nstart, + const int ldh, + const T *A, // hcc + T *B, // scc + Real *W, // eigenvalue + T *V) { // assert(nstart == ldh); // A to V - CHECK_CUDA(cudaMemcpy(V, A, sizeof(T) * ldh * nstart, cudaMemcpyDeviceToDevice)); - xhegvd_wrapper(CUBLAS_FILL_MODE_UPPER, nstart, V, ldh, - (T*)B, ldh, W); + CHECK_CUDA(cudaMemcpy( + V, + A, + sizeof(T) * ldh * nstart, + cudaMemcpyDeviceToDevice)); + xhegvd_wrapper( + CUBLAS_FILL_MODE_UPPER, + nstart, + V, + ldh, (T *)B, + ldh, + W); } }; -template -struct heevx_op +template struct heevx_op { using Real = typename GetTypeReal::type; - void operator()(const base_device::DEVICE_GPU* d, - const int nstart, - const int ldh, - const T* A, // hcc - const int m, - Real* W, // eigenvalue - T* V) + void operator()( + const base_device::DEVICE_GPU *d, + const int nstart, + const int ldh, + const T *A, // hcc + const int m, + Real *W, // eigenvalue + T *V) { assert(nstart <= ldh); // A to V - CHECK_CUDA(cudaMemcpy(V, A, sizeof(T) * nstart * ldh, cudaMemcpyDeviceToDevice)); - xheevd_wrapper(CUBLAS_FILL_MODE_LOWER, nstart, V, ldh, W); + CHECK_CUDA(cudaMemcpy( + V, + A, + sizeof(T) * nstart * ldh, + cudaMemcpyDeviceToDevice)); + xheevd_wrapper( + CUBLAS_FILL_MODE_LOWER, + nstart, + V, + ldh, + W); } }; -template -struct hegvx_op +template struct hegvx_op { using Real = typename GetTypeReal::type; - void operator()(const base_device::DEVICE_GPU* d, - const int nbase, - const int ldh, - T* hcc, - T* scc, - const int m, - Real* eigenvalue, - T* vcc) + void operator()( + const base_device::DEVICE_GPU *d, + const int nbase, + const int ldh, + T *hcc, + T *scc, + const int m, + Real *eigenvalue, + T *vcc) { - } }; diff --git a/source/source_hsolver/kernels/hegvd_op.cpp b/source/source_hsolver/kernels/hegvd_op.cpp index 2c00a80c44c..0f996bdd700 100644 --- a/source/source_hsolver/kernels/hegvd_op.cpp +++ b/source/source_hsolver/kernels/hegvd_op.cpp @@ -1,4 +1,5 @@ #include "source_hsolver/kernels/hegvd_op.h" + #include "source_base/module_container/base/third_party/lapack.h" #include @@ -11,17 +12,17 @@ namespace lapackConnector = container::lapackConnector; // see "source_base/modu namespace hsolver { // hegvd and sygvd; dn for dense? -template -struct hegvd_op +template struct hegvd_op { using Real = typename GetTypeReal::type; - void operator()(const base_device::DEVICE_CPU* d, - const int nstart, - const int ldh, - const T* hcc, - T* scc, - Real* eigenvalue, - T* vcc) + void operator()( + const base_device::DEVICE_CPU *d, + const int nstart, + const int ldh, + const T *hcc, + T *scc, + Real *eigenvalue, + T *vcc) { for (int i = 0; i < nstart * ldh; i++) { @@ -31,36 +32,37 @@ struct hegvd_op // SYGVD needs more WORK than HEGVD when eigenvectors are requested. const bool real_problem = std::is_same::value; int lwork = real_problem ? 1 + 6 * nstart + 2 * nstart * nstart : 2 * nstart + nstart * nstart; - T* work = new T[lwork]; + T *work = new T[lwork]; std::fill(work, work + lwork, T(0)); int lrwork = 1 + 5 * nstart + 2 * nstart * nstart; - Real* rwork = new Real[lrwork]; + Real *rwork = new Real[lrwork]; std::fill(rwork, rwork + lrwork, Real(0)); int liwork = 3 + 5 * nstart; - int* iwork = new int[liwork]; + int *iwork = new int[liwork]; std::fill(iwork, iwork + liwork, 0); //=========================== // calculate all eigenvalues //=========================== - lapackConnector::hegvd(1, - 'V', - 'U', - nstart, - vcc, - ldh, - scc, - ldh, - eigenvalue, - work, - lwork, - rwork, - lrwork, - iwork, - liwork, - info); + lapackConnector::hegvd( + 1, + 'V', + 'U', + nstart, + vcc, + ldh, + scc, + ldh, + eigenvalue, + work, + lwork, + rwork, + lrwork, + iwork, + liwork, + info); if (info != 0) { @@ -153,19 +155,19 @@ struct hegvd_op * both heevx and syevx are implemented through the `evx` interface of LAPACK. * wrapped in LapackWrapper::xheevx */ -template -struct heevx_op +template struct heevx_op { using Real = typename GetTypeReal::type; - void operator()(const base_device::DEVICE_CPU* /*ctx*/, - const int nstart, - const int ldh, - const T* hcc, // hcc - const int nbands, // nbands - Real* eigenvalue, // eigenvalue - T* vcc) // vcc + void operator()( + const base_device::DEVICE_CPU * /*ctx*/, + const int nstart, + const int ldh, + const T *hcc, // hcc + const int nbands, // nbands + Real *eigenvalue, // eigenvalue + T *vcc) // vcc { - T* aux = new T[nstart * ldh]; + T *aux = new T[nstart * ldh]; for (int ii = 0; ii < nstart * ldh; ii++) { aux[ii] = hcc[ii]; @@ -173,10 +175,10 @@ struct heevx_op int info = 0; int lwork = -1; - T* work = new T[1]; - Real* rwork = new Real[7 * nstart]; - int* iwork = new int[5 * nstart]; - int* ifail = new int[nstart]; + T *work = new T[1]; + Real *rwork = new Real[7 * nstart]; + int *iwork = new int[5 * nstart]; + int *ifail = new int[nstart]; // When lwork = -1, the demension of work will be assumed // Assume the denmension of work by output work[0] @@ -246,18 +248,18 @@ struct heevx_op } }; -template -struct hegvx_op +template struct hegvx_op { using Real = typename GetTypeReal::type; - void operator()(const base_device::DEVICE_CPU* d, - const int nbase, - const int ldh, - T* hcc, - T* scc, - const int m, - Real* eigenvalue, - T* vcc) + void operator()( + const base_device::DEVICE_CPU *d, + const int nbase, + const int ldh, + T *hcc, + T *scc, + const int m, + Real *eigenvalue, + T *vcc) { int info = 0; @@ -266,10 +268,10 @@ struct hegvx_op int lwork = -1; - T* work = new T[1]; - Real* rwork = new Real[7 * nbase]; - int* iwork = new int[5 * nbase]; - int* ifail = new int[nbase]; + T *work = new T[1]; + Real *rwork = new Real[7 * nbase]; + int *iwork = new int[5 * nbase]; + int *ifail = new int[nbase]; lapackConnector::hegvx( 1, // ITYPE = 1: A*x = (lambda)*B*x @@ -301,30 +303,31 @@ struct hegvx_op delete[] work; work = new T[lwork]; - lapackConnector::hegvx(1, - 'V', - 'I', - 'U', - nbase, - hcc, - ldh, - scc, - ldh, - 0.0, - 0.0, - 1, - m, - 0.0, - mm, - eigenvalue, - vcc, - ldh, - work, - lwork, - rwork, - iwork, - ifail, - info); + lapackConnector::hegvx( + 1, + 'V', + 'I', + 'U', + nbase, + hcc, + ldh, + scc, + ldh, + 0.0, + 0.0, + 1, + m, + 0.0, + mm, + eigenvalue, + vcc, + ldh, + work, + lwork, + rwork, + iwork, + ifail, + info); delete[] work; delete[] rwork; diff --git a/source/source_hsolver/kernels/hegvd_op.h b/source/source_hsolver/kernels/hegvd_op.h index c88e51bc639..cedb89671e2 100644 --- a/source/source_hsolver/kernels/hegvd_op.h +++ b/source/source_hsolver/kernels/hegvd_op.h @@ -1,4 +1,4 @@ -/// This is the module for wrapper of +/// This is the module for wrapper of /// DeNse Generalized eigenValue (eXtended) /// HErmitian / SYmmetric @@ -7,8 +7,8 @@ #ifndef MODULE_HSOLVER_HEGVD_H #define MODULE_HSOLVER_HEGVD_H -#include "source_base/parallel_reduce.h" #include "source_base/module_device/types.h" +#include "source_base/parallel_reduce.h" #include @@ -43,9 +43,7 @@ inline double get_real(const double &x) { return x; } inline float get_real(const float &x) { return x; } - -template -struct hegvd_op +template struct hegvd_op { using Real = typename GetTypeReal::type; /// @brief HEGVD computes all the eigenvalues and eigenvectors of a complex generalized @@ -67,7 +65,14 @@ struct hegvd_op /// Output Parameter /// @param W : calculated eigenvalues /// @param V : calculated eigenvectors (col major) - void operator()(const Device* d, const int nstart, const int ldh, const T* A, T* B, Real* W, T* V); + void operator()( + const Device *d, + const int nstart, + const int ldh, + const T *A, + T *B, + Real *W, + T *V); }; // template @@ -87,8 +92,7 @@ struct hegvd_op // void operator()(const Device* d, const int nstart, const int ldh, const T* A, T* B, Real* W, T* V); // }; -template -struct hegvx_op +template struct hegvx_op { using Real = typename GetTypeReal::type; /// @brief HEGVX computes first m eigenvalues and eigenvectors of a complex generalized @@ -102,11 +106,18 @@ struct hegvx_op /// Output Parameter /// @param W : calculated eigenvalues /// @param V : calculated eigenvectors (col major) - void operator()(const Device* d, const int nstart, const int ldh, T* A, T* B, const int m, Real* W, T* V); + void operator()( + const Device *d, + const int nstart, + const int ldh, + T *A, + T *B, + const int m, + Real *W, + T *V); }; -template -struct heevx_op +template struct heevx_op { using Real = typename GetTypeReal::type; /// @brief heevx computes the first m eigenvalues and their corresponding eigenvectors of @@ -128,7 +139,14 @@ struct heevx_op /// Output Parameter /// @param w: calculated eigenvalues /// @param z: calculated eigenvectors - void operator()(const Device *d, const int ndim, const int lda, const T *A, const int neig, Real *w, T *z); + void operator()( + const Device *d, + const int ndim, + const int lda, + const T *A, + const int neig, + Real *w, + T *z); }; #if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM diff --git a/source/source_hsolver/kernels/rocm/bpcg_kernel_op.hip.cu b/source/source_hsolver/kernels/rocm/bpcg_kernel_op.hip.cu index ee81337e2b8..539177f94ab 100644 --- a/source/source_hsolver/kernels/rocm/bpcg_kernel_op.hip.cu +++ b/source/source_hsolver/kernels/rocm/bpcg_kernel_op.hip.cu @@ -9,17 +9,17 @@ namespace hsolver { template __global__ void line_minimize_with_block( - thrust::complex* grad, - thrust::complex* hgrad, - thrust::complex* sgrad, - thrust::complex* psi, - thrust::complex* hpsi, - thrust::complex* spsi, - const int n_basis, - const int n_basis_max) + thrust::complex *grad, + thrust::complex *hgrad, + thrust::complex *sgrad, + thrust::complex *psi, + thrust::complex *hpsi, + thrust::complex *spsi, + const int n_basis, + const int n_basis_max) { int band_idx = blockIdx.x; // band_idx - int tid = threadIdx.x; // basis_idx + int tid = threadIdx.x; // basis_idx int item = 0; Real epsilo_0 = 0.0, epsilo_1 = 0.0, epsilo_2 = 0.0; Real theta = 0.0, cos_theta = 0.0, sin_theta = 0.0; @@ -27,20 +27,24 @@ __global__ void line_minimize_with_block( data[tid] = 0; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) + { item = band_idx * n_basis_max + basis_idx; data[tid] += (sgrad[item] * thrust::conj(grad[item])).real(); } __syncthreads(); // just do some parallel reduction in shared memory - for (int ii = THREAD_PER_BLOCK >> 1; ii > 0; ii >>= 1) { - if (tid < ii) { + for (int ii = THREAD_PER_BLOCK >> 1; ii > 0; ii >>= 1) + { + if (tid < ii) + { data[tid] += data[tid + ii]; } __syncthreads(); } - if (!(data[0] > 1.0e-20)) { + if (!(data[0] > 1.0e-20)) + { return; } Real norm = 1.0 / sqrt(data[0]); @@ -49,7 +53,8 @@ __global__ void line_minimize_with_block( data[tid] = 0; data[THREAD_PER_BLOCK + tid] = 0; data[2 * THREAD_PER_BLOCK + tid] = 0; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) + { item = band_idx * n_basis_max + basis_idx; grad[item] *= norm; hgrad[item] *= norm; @@ -61,8 +66,10 @@ __global__ void line_minimize_with_block( __syncthreads(); // just do some parallel reduction in shared memory - for (int ii = THREAD_PER_BLOCK >> 1; ii > 0; ii >>= 1) { - if (tid < ii) { + for (int ii = THREAD_PER_BLOCK >> 1; ii > 0; ii >>= 1) + { + if (tid < ii) + { data[tid] += data[tid + ii]; data[THREAD_PER_BLOCK + tid] += data[THREAD_PER_BLOCK + tid + ii]; data[2 * THREAD_PER_BLOCK + tid] += data[2 * THREAD_PER_BLOCK + tid + ii]; @@ -75,18 +82,19 @@ __global__ void line_minimize_with_block( theta = 0.5 * atan2(2 * epsilo_1, epsilo_0 - epsilo_2); // Choose the rotation associated with the lower Ritz value. - const Real energy_delta = (epsilo_0 - epsilo_2) * cos(2.0 * theta) - + 2.0 * epsilo_1 * sin(2.0 * theta); + const Real energy_delta = (epsilo_0 - epsilo_2) * cos(2.0 * theta) + 2.0 * epsilo_1 * sin(2.0 * theta); const Real energy_1 = 0.5 * (epsilo_0 + epsilo_2 + energy_delta); const Real energy_2 = 0.5 * (epsilo_0 + epsilo_2 - energy_delta); - if (energy_1 > energy_2) { + if (energy_1 > energy_2) + { theta += 2.0 * atan(1.0); } cos_theta = cos(theta); sin_theta = sin(theta); - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) + { item = band_idx * n_basis_max + basis_idx; - psi [item] = psi [item] * cos_theta + grad [item] * sin_theta; + psi[item] = psi[item] * cos_theta + grad[item] * sin_theta; hpsi[item] = hpsi[item] * cos_theta + hgrad[item] * sin_theta; spsi[item] = spsi[item] * cos_theta + sgrad[item] * sin_theta; } @@ -94,19 +102,19 @@ __global__ void line_minimize_with_block( template __global__ void calc_grad_with_block( - const Real* prec, - Real* err, - Real* beta, - thrust::complex* psi, - thrust::complex* hpsi, - thrust::complex* spsi, - thrust::complex* grad, - thrust::complex* grad_old, - const int n_basis, - const int n_basis_max) + const Real *prec, + Real *err, + Real *beta, + thrust::complex *psi, + thrust::complex *hpsi, + thrust::complex *spsi, + thrust::complex *grad, + thrust::complex *grad_old, + const int n_basis, + const int n_basis_max) { int band_idx = blockIdx.x; // band_idx - int tid = threadIdx.x; // basis_idx + int tid = threadIdx.x; // basis_idx int item = 0; Real err_st = 0.0; Real beta_st = 0.0; @@ -118,14 +126,17 @@ __global__ void calc_grad_with_block( // Init shared memory data[tid] = 0; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) + { item = band_idx * n_basis_max + basis_idx; data[tid] += (spsi[item] * thrust::conj(psi[item])).real(); } __syncthreads(); // just do some parallel reduction in shared memory - for (int ii = THREAD_PER_BLOCK >> 1; ii > 0; ii >>= 1) { - if (tid < ii) { + for (int ii = THREAD_PER_BLOCK >> 1; ii > 0; ii >>= 1) + { + if (tid < ii) + { data[tid] += data[tid + ii]; } __syncthreads(); @@ -135,7 +146,8 @@ __global__ void calc_grad_with_block( __syncthreads(); data[tid] = 0; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) + { item = band_idx * n_basis_max + basis_idx; psi[item] *= norm; hpsi[item] *= norm; @@ -145,8 +157,10 @@ __global__ void calc_grad_with_block( __syncthreads(); // just do some parallel reduction in shared memory - for (int ii = THREAD_PER_BLOCK >> 1; ii > 0; ii >>= 1) { - if (tid < ii) { + for (int ii = THREAD_PER_BLOCK >> 1; ii > 0; ii >>= 1) + { + if (tid < ii) + { data[tid] += data[tid + ii]; } __syncthreads(); @@ -156,7 +170,8 @@ __global__ void calc_grad_with_block( data[tid] = 0; data[THREAD_PER_BLOCK + tid] = 0; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) + { item = band_idx * n_basis_max + basis_idx; grad_1 = hpsi[item] - epsilo * spsi[item]; grad_2 = thrust::norm(grad_1); @@ -166,8 +181,10 @@ __global__ void calc_grad_with_block( __syncthreads(); // just do some parallel reduction in shared memory - for (int ii = THREAD_PER_BLOCK >> 1; ii > 0; ii >>= 1) { - if (tid < ii) { + for (int ii = THREAD_PER_BLOCK >> 1; ii > 0; ii >>= 1) + { + if (tid < ii) + { data[tid] += data[tid + ii]; data[THREAD_PER_BLOCK + tid] += data[THREAD_PER_BLOCK + tid + ii]; } @@ -175,14 +192,16 @@ __global__ void calc_grad_with_block( } err_st = data[0]; beta_st = data[THREAD_PER_BLOCK]; - for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) { + for (int basis_idx = tid; basis_idx < n_basis; basis_idx += THREAD_PER_BLOCK) + { item = band_idx * n_basis_max + basis_idx; grad_1 = hpsi[item] - epsilo * spsi[item]; grad[item] = -grad_1 / prec[basis_idx] + beta_st / beta[band_idx] * grad_old[item]; } __syncthreads(); - if (tid == 0) { + if (tid == 0) + { beta[band_idx] = beta_st; err[band_idx] = sqrt(err_st); } @@ -190,34 +209,36 @@ __global__ void calc_grad_with_block( template __global__ void apply_eigenvalues_kernel( - const thrust::complex* vectors, - thrust::complex* result, - const Real* eigenvalues, - const int nbase, - const int nbase_x, - const int notconv) + const thrust::complex *vectors, + thrust::complex *result, + const Real *eigenvalues, + const int nbase, + const int nbase_x, + const int notconv) { int m = blockIdx.x; int idx = threadIdx.x + blockIdx.y * blockDim.x; - if (m < notconv && idx < nbase) { + if (m < notconv && idx < nbase) + { result[m * nbase_x + idx] = eigenvalues[m] * vectors[m * nbase_x + idx]; } } template __global__ void precondition_kernel( - thrust::complex* psi_iter, - const Real* precondition, - const Real* eigenvalues, - const int dim, - const int nbase, - const int notconv) + thrust::complex *psi_iter, + const Real *precondition, + const Real *eigenvalues, + const int dim, + const int nbase, + const int notconv) { int m = blockIdx.x; int i = threadIdx.x + blockIdx.y * blockDim.x; - if (m < notconv && i < dim) { + if (m < notconv && i < dim) + { Real x = abs(precondition[i] - eigenvalues[m]); Real pre = 0.5 * (1.0 + x + sqrt(1 + (x - 1.0) * (x - 1.0))); psi_iter[(nbase + m) * dim + i] = psi_iter[(nbase + m) * dim + i] / pre; @@ -226,11 +247,11 @@ __global__ void precondition_kernel( template __global__ void normalize_kernel( - thrust::complex* psi_iter, - Real* psi_norm, - const int dim, - const int nbase, - const int notconv) + thrust::complex *psi_iter, + Real *psi_norm, + const int dim, + const int nbase, + const int notconv) { int m = blockIdx.x; int tid = threadIdx.x; @@ -239,7 +260,8 @@ __global__ void normalize_kernel( sum[tid] = 0.0; // Calculate the sum for normalization - for (int i = tid; i < dim; i += THREAD_PER_BLOCK) { + for (int i = tid; i < dim; i += THREAD_PER_BLOCK) + { auto val = psi_iter[(nbase + m) * dim + i]; sum[tid] += (val * thrust::conj(val)).real(); } @@ -247,8 +269,10 @@ __global__ void normalize_kernel( __syncthreads(); // Parallel reduction in shared memory - for (int s = THREAD_PER_BLOCK/2; s > 0; s >>= 1) { - if (tid < s) { + for (int s = THREAD_PER_BLOCK / 2; s > 0; s >>= 1) + { + if (tid < s) + { sum[tid] += sum[tid + s]; } __syncthreads(); @@ -257,25 +281,27 @@ __global__ void normalize_kernel( Real norm = sqrt(sum[0]); // Normalize the vector - for (int i = tid; i < dim; i += THREAD_PER_BLOCK) { + for (int i = tid; i < dim; i += THREAD_PER_BLOCK) + { psi_iter[(nbase + m) * dim + i] /= norm; } // Store the norm if needed - if (tid == 0 && psi_norm != nullptr) { + if (tid == 0 && psi_norm != nullptr) + { psi_norm[m] = norm; } } template __global__ void refresh_hcc_scc_vcc_kernel( - const int n, - T *hcc, - T *scc, - T *vcc, - const int ldh, - const Real *eigenvalue, - const T one) + const int n, + T *hcc, + T *scc, + T *vcc, + const int ldh, + const Real *eigenvalue, + const T one) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < n) @@ -287,163 +313,215 @@ __global__ void refresh_hcc_scc_vcc_kernel( } template -void line_minimize_with_block_op::operator()(T* grad_out, - T* hgrad_out, - T* sgrad_out, - T* psi_out, - T* hpsi_out, - T* spsi_out, - const int& n_basis, - const int& n_basis_max, - const int& n_band) +void line_minimize_with_block_op::operator()( + T *grad_out, + T *hgrad_out, + T *sgrad_out, + T *psi_out, + T *hpsi_out, + T *spsi_out, + const int &n_basis, + const int &n_basis_max, + const int &n_band) { - auto A = reinterpret_cast*>(grad_out); - auto B = reinterpret_cast*>(hgrad_out); - auto C = reinterpret_cast*>(sgrad_out); - auto D = reinterpret_cast*>(psi_out); - auto E = reinterpret_cast*>(hpsi_out); - auto F = reinterpret_cast*>(spsi_out); + auto A = reinterpret_cast *>(grad_out); + auto B = reinterpret_cast *>(hgrad_out); + auto C = reinterpret_cast *>(sgrad_out); + auto D = reinterpret_cast *>(psi_out); + auto E = reinterpret_cast *>(hpsi_out); + auto F = reinterpret_cast *>(spsi_out); line_minimize_with_block<<>>( - A, B, C, D, E, F, - n_basis, n_basis_max); + A, + B, + C, + D, + E, + F, + n_basis, + n_basis_max); hipCheckOnDebug(); } template -void calc_grad_with_block_op::operator()(const Real* prec_in, - Real* err_out, - Real* beta_out, - T* psi_out, - T* hpsi_out, - T* spsi_out, - T* grad_out, - T* grad_old_out, - const int& n_basis, - const int& n_basis_max, - const int& n_band) +void calc_grad_with_block_op::operator()( + const Real *prec_in, + Real *err_out, + Real *beta_out, + T *psi_out, + T *hpsi_out, + T *spsi_out, + T *grad_out, + T *grad_old_out, + const int &n_basis, + const int &n_basis_max, + const int &n_band) { - auto A = reinterpret_cast*>(psi_out); - auto B = reinterpret_cast*>(hpsi_out); - auto C = reinterpret_cast*>(spsi_out); - auto D = reinterpret_cast*>(grad_out); - auto E = reinterpret_cast*>(grad_old_out); + auto A = reinterpret_cast *>(psi_out); + auto B = reinterpret_cast *>(hpsi_out); + auto C = reinterpret_cast *>(spsi_out); + auto D = reinterpret_cast *>(grad_out); + auto E = reinterpret_cast *>(grad_old_out); calc_grad_with_block<<>>( - prec_in, err_out, beta_out, - A, B, C, D, E, - n_basis, n_basis_max); + prec_in, + err_out, + beta_out, + A, + B, + C, + D, + E, + n_basis, + n_basis_max); hipCheckOnDebug(); } template -void apply_eigenvalues_op::operator()(const int& nbase, - const int& nbase_x, - const int& notconv, - T* result, - const T* vectors, - const Real* eigenvalues) +void apply_eigenvalues_op::operator()( + const int &nbase, + const int &nbase_x, + const int ¬conv, + T *result, + const T *vectors, + const Real *eigenvalues) { const int threads_per_block = 256; const int blocks_per_grid_y = (nbase + threads_per_block - 1) / threads_per_block; dim3 grid(notconv, blocks_per_grid_y); - auto vec_complex = reinterpret_cast*>(vectors); - auto res_complex = reinterpret_cast*>(result); + auto vec_complex = reinterpret_cast *>(vectors); + auto res_complex = reinterpret_cast *>(result); apply_eigenvalues_kernel<<>>( - vec_complex, res_complex, eigenvalues, nbase, nbase_x, notconv); + vec_complex, + res_complex, + eigenvalues, + nbase, + nbase_x, + notconv); hipCheckOnDebug(); } template -void precondition_op::operator()(const int& dim, - T* psi_iter, - const int& nbase, - const int& notconv, - const Real* precondition, - const Real* eigenvalues) +void precondition_op::operator()( + const int &dim, + T *psi_iter, + const int &nbase, + const int ¬conv, + const Real *precondition, + const Real *eigenvalues) { const int threads_per_block = 256; const int blocks_per_grid_y = (dim + threads_per_block - 1) / threads_per_block; dim3 grid(notconv, blocks_per_grid_y); - auto psi_complex = reinterpret_cast*>(psi_iter); + auto psi_complex = reinterpret_cast *>(psi_iter); precondition_kernel<<>>( - psi_complex, precondition, eigenvalues, dim, nbase, notconv); + psi_complex, + precondition, + eigenvalues, + dim, + nbase, + notconv); hipCheckOnDebug(); } template -void normalize_op::operator()(const int& dim, - T* psi_iter, - const int& nbase, - const int& notconv, - Real* psi_norm) +void normalize_op::operator()( + const int &dim, + T *psi_iter, + const int &nbase, + const int ¬conv, + Real *psi_norm) { - auto psi_complex = reinterpret_cast*>(psi_iter); + auto psi_complex = reinterpret_cast *>(psi_iter); normalize_kernel<<>>( - psi_complex, psi_norm, dim, nbase, notconv); + psi_complex, + psi_norm, + dim, + nbase, + notconv); hipCheckOnDebug(); } template <> -void refresh_hcc_scc_vcc_op::operator()(const int &n, - double *hcc, - double *scc, - double *vcc, - const int &ldh, - const double *eigenvalue, - const double& one) +void refresh_hcc_scc_vcc_op::operator()( + const int &n, + double *hcc, + double *scc, + double *vcc, + const int &ldh, + const double *eigenvalue, + const double &one) { int thread = 512; int block = (n + thread - 1) / thread; - refresh_hcc_scc_vcc_kernel <<>> (n, hcc, scc, vcc, ldh, eigenvalue, one); + refresh_hcc_scc_vcc_kernel<<>>( + n, + hcc, + scc, + vcc, + ldh, + eigenvalue, + one); hipCheckOnDebug(); } template <> -void refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>::operator()(const int &n, - std::complex *hcc, - std::complex *scc, - std::complex *vcc, - const int &ldh, - const float *eigenvalue, - const std::complex& one) +void refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>::operator()( + const int &n, + std::complex *hcc, + std::complex *scc, + std::complex *vcc, + const int &ldh, + const float *eigenvalue, + const std::complex &one) { int thread = 512; int block = (n + thread - 1) / thread; - refresh_hcc_scc_vcc_kernel, float> <<>> (n, reinterpret_cast*>(hcc), - reinterpret_cast*>(scc), reinterpret_cast*>(vcc), ldh, eigenvalue, - thrust::complex(one)); + refresh_hcc_scc_vcc_kernel, float><<>>( + n, + reinterpret_cast *>(hcc), + reinterpret_cast *>(scc), + reinterpret_cast *>(vcc), + ldh, + eigenvalue, + thrust::complex(one)); hipCheckOnDebug(); } template <> -void refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>::operator()(const int &n, - std::complex *hcc, - std::complex *scc, - std::complex *vcc, - const int &ldh, - const double *eigenvalue, - const std::complex& one) +void refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>::operator()( + const int &n, + std::complex *hcc, + std::complex *scc, + std::complex *vcc, + const int &ldh, + const double *eigenvalue, + const std::complex &one) { int thread = 512; int block = (n + thread - 1) / thread; - refresh_hcc_scc_vcc_kernel, double> <<>> (n, reinterpret_cast*>(hcc), - reinterpret_cast*>(scc), reinterpret_cast*>(vcc), ldh, eigenvalue, - thrust::complex(one)); + refresh_hcc_scc_vcc_kernel, double><<>>( + n, + reinterpret_cast *>(hcc), + reinterpret_cast *>(scc), + reinterpret_cast *>(vcc), + ldh, + eigenvalue, + thrust::complex(one)); hipCheckOnDebug(); } @@ -464,4 +542,4 @@ template struct normalize_op; template struct refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>; template struct refresh_hcc_scc_vcc_op, base_device::DEVICE_GPU>; template struct refresh_hcc_scc_vcc_op; -} +} // namespace hsolver diff --git a/source/source_hsolver/kernels/rocm/hegvd_op.hip.cu b/source/source_hsolver/kernels/rocm/hegvd_op.hip.cu index 93b3457af69..21da852f22a 100644 --- a/source/source_hsolver/kernels/rocm/hegvd_op.hip.cu +++ b/source/source_hsolver/kernels/rocm/hegvd_op.hip.cu @@ -1,9 +1,10 @@ #include "source_hsolver/kernels/hegvd_op.h" -#include #include +#include -namespace hsolver { +namespace hsolver +{ // NOTE: mimicked from ../cuda/hegvd_op.cu for three hegvd_op @@ -11,14 +12,16 @@ static hipsolverHandle_t hipsolver_H = nullptr; // Test on DCU platform. When nstart is greater than 234, code on DCU performs better. const int N_DCU = 234; -void createGpuSolverHandle() { +void createGpuSolverHandle() +{ if (hipsolver_H == nullptr) { hipsolverErrcheck(hipsolverCreate(&hipsolver_H)); } } -void destroyGpuSolverHandle() { +void destroyGpuSolverHandle() +{ if (hipsolver_H != nullptr) { hipsolverErrcheck(hipsolverDestroy(hipsolver_H)); @@ -28,339 +31,502 @@ void destroyGpuSolverHandle() { #ifdef __LCAO template <> -void hegvd_op::operator()(const base_device::DEVICE_GPU* ctx, - const int nstart, - const int ldh, - const double* _hcc, - double* _scc, - double* _eigenvalue, - double* _vcc) +void hegvd_op::operator()( + const base_device::DEVICE_GPU *ctx, + const int nstart, + const int ldh, + const double *_hcc, + double *_scc, + double *_eigenvalue, + double *_vcc) { // copied from ../cuda/hegvd_op.cu, "hegvd_op" assert(nstart == ldh); - if (nstart > N_DCU){ - hipErrcheck(hipMemcpy(_vcc, _hcc, sizeof(double) * ldh * nstart, hipMemcpyDeviceToDevice)); + if (nstart > N_DCU) + { + hipErrcheck(hipMemcpy( + _vcc, + _hcc, + sizeof(double) * ldh * nstart, + hipMemcpyDeviceToDevice)); // now vcc contains hcc // prepare some values for hipsolverDnZhegvd_bufferSize - int * devInfo = nullptr; + int *devInfo = nullptr; int lwork = 0, info_gpu = 0; - double * work = nullptr; - hipErrcheck(hipMalloc((void**)&devInfo, sizeof(int))); + double *work = nullptr; + hipErrcheck(hipMalloc((void **)&devInfo, sizeof(int))); hipsolverFillMode_t uplo = HIPSOLVER_FILL_MODE_UPPER; // calculate the sizes needed for pre-allocated buffer. hipsolverErrcheck(hipsolverDnDsygvd_bufferSize( - hipsolver_H, HIPSOLVER_EIG_TYPE_1, HIPSOLVER_EIG_MODE_VECTOR, uplo, + hipsolver_H, + HIPSOLVER_EIG_TYPE_1, + HIPSOLVER_EIG_MODE_VECTOR, + uplo, nstart, - _vcc, ldh, - _scc, ldh, + _vcc, + ldh, + _scc, + ldh, _eigenvalue, &lwork)); // allocate memery - hipErrcheck(hipMalloc((void**)&work, sizeof(double) * lwork)); + hipErrcheck(hipMalloc((void **)&work, sizeof(double) * lwork)); // compute eigenvalues and eigenvectors. hipsolverErrcheck(hipsolverDnDsygvd( - hipsolver_H, HIPSOLVER_EIG_TYPE_1, HIPSOLVER_EIG_MODE_VECTOR, uplo, + hipsolver_H, + HIPSOLVER_EIG_TYPE_1, + HIPSOLVER_EIG_MODE_VECTOR, + uplo, nstart, - _vcc, ldh, - const_cast(_scc), ldh, + _vcc, + ldh, + const_cast(_scc), + ldh, _eigenvalue, - work, lwork, devInfo)); + work, + lwork, + devInfo)); - hipErrcheck(hipMemcpy(&info_gpu, devInfo, sizeof(int), hipMemcpyDeviceToHost)); + hipErrcheck(hipMemcpy( + &info_gpu, + devInfo, + sizeof(int), + hipMemcpyDeviceToHost)); // free the buffer hipErrcheck(hipFree(work)); hipErrcheck(hipFree(devInfo)); } // if(fail_info != nullptr) *fail_info = info_gpu; - else{ + else + { std::vector hcc(nstart * nstart, 0.0); std::vector scc(nstart * nstart, 0.0); std::vector vcc(nstart * nstart, 0.0); std::vector eigenvalue(nstart, 0); - hipErrcheck(hipMemcpy(hcc.data(), _hcc, sizeof(double) * hcc.size(), hipMemcpyDeviceToHost)); - hipErrcheck(hipMemcpy(scc.data(), _scc, sizeof(double) * scc.size(), hipMemcpyDeviceToHost)); - base_device::DEVICE_CPU* cpu_ctx = {}; - hegvd_op()(cpu_ctx, - nstart, - ldh, - hcc.data(), - scc.data(), - eigenvalue.data(), - vcc.data()); - hipErrcheck(hipMemcpy(_vcc, vcc.data(), sizeof(double) * vcc.size(), hipMemcpyHostToDevice)); - hipErrcheck(hipMemcpy(_eigenvalue, eigenvalue.data(), sizeof(double) * eigenvalue.size(), hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + hcc.data(), + _hcc, + sizeof(double) * hcc.size(), + hipMemcpyDeviceToHost)); + hipErrcheck(hipMemcpy( + scc.data(), + _scc, + sizeof(double) * scc.size(), + hipMemcpyDeviceToHost)); + base_device::DEVICE_CPU *cpu_ctx = {}; + hegvd_op()( + cpu_ctx, + nstart, + ldh, + hcc.data(), + scc.data(), + eigenvalue.data(), + vcc.data()); + hipErrcheck(hipMemcpy( + _vcc, + vcc.data(), + sizeof(double) * vcc.size(), + hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + _eigenvalue, + eigenvalue.data(), + sizeof(double) * eigenvalue.size(), + hipMemcpyHostToDevice)); } - - } #endif // __LCAO template <> -void hegvd_op, base_device::DEVICE_GPU>::operator()(const base_device::DEVICE_GPU* ctx, - const int nstart, - const int ldh, - const std::complex* _hcc, - const std::complex* _scc, - float* _eigenvalue, - std::complex* _vcc) +void hegvd_op, base_device::DEVICE_GPU>::operator()( + const base_device::DEVICE_GPU *ctx, + const int nstart, + const int ldh, + const std::complex *_hcc, + const std::complex *_scc, + float *_eigenvalue, + std::complex *_vcc) { // copied from ../cuda/hegvd_op.cu, "hegvd_op" assert(nstart == ldh); - if (nstart > N_DCU){ - hipErrcheck(hipMemcpy(_vcc, _hcc, sizeof(std::complex) * ldh * nstart, hipMemcpyDeviceToDevice)); + if (nstart > N_DCU) + { + hipErrcheck(hipMemcpy( + _vcc, + _hcc, + sizeof(std::complex) * ldh * nstart, + hipMemcpyDeviceToDevice)); // now vcc contains hcc // prepare some values for hipsolverDnZhegvd_bufferSize - int * devInfo = nullptr; + int *devInfo = nullptr; int lwork = 0, info_gpu = 0; - float2 * work = nullptr; - hipErrcheck(hipMalloc((void**)&devInfo, sizeof(int))); + float2 *work = nullptr; + hipErrcheck(hipMalloc((void **)&devInfo, sizeof(int))); hipsolverFillMode_t uplo = HIPSOLVER_FILL_MODE_UPPER; // calculate the sizes needed for pre-allocated buffer. hipsolverErrcheck(hipsolverDnChegvd_bufferSize( - hipsolver_H, HIPSOLVER_EIG_TYPE_1, HIPSOLVER_EIG_MODE_VECTOR, uplo, + hipsolver_H, + HIPSOLVER_EIG_TYPE_1, + HIPSOLVER_EIG_MODE_VECTOR, + uplo, nstart, - reinterpret_cast(_vcc), ldh, - reinterpret_cast(_scc), ldh, + reinterpret_cast(_vcc), + ldh, + reinterpret_cast(_scc), + ldh, _eigenvalue, &lwork)); // allocate memery - hipErrcheck(hipMalloc((void**)&work, sizeof(float2) * lwork)); + hipErrcheck(hipMalloc((void **)&work, sizeof(float2) * lwork)); // compute eigenvalues and eigenvectors. hipsolverErrcheck(hipsolverDnChegvd( - hipsolver_H, HIPSOLVER_EIG_TYPE_1, HIPSOLVER_EIG_MODE_VECTOR, uplo, + hipsolver_H, + HIPSOLVER_EIG_TYPE_1, + HIPSOLVER_EIG_MODE_VECTOR, + uplo, nstart, - reinterpret_cast(_vcc), ldh, - const_cast(reinterpret_cast(_scc)), ldh, + reinterpret_cast(_vcc), + ldh, + const_cast(reinterpret_cast(_scc)), + ldh, _eigenvalue, - work, lwork, devInfo)); - - hipErrcheck(hipMemcpy(&info_gpu, devInfo, sizeof(int), hipMemcpyDeviceToHost)); + work, + lwork, + devInfo)); + + hipErrcheck(hipMemcpy( + &info_gpu, + devInfo, + sizeof(int), + hipMemcpyDeviceToHost)); // free the buffer hipErrcheck(hipFree(work)); hipErrcheck(hipFree(devInfo)); } // if(fail_info != nullptr) *fail_info = info_gpu; - else{ + else + { std::vector> hcc(nstart * nstart, {0, 0}); std::vector> scc(nstart * nstart, {0, 0}); std::vector> vcc(nstart * nstart, {0, 0}); std::vector eigenvalue(nstart, 0); - hipErrcheck(hipMemcpy(hcc.data(), _hcc, sizeof(std::complex) * hcc.size(), hipMemcpyDeviceToHost)); - hipErrcheck(hipMemcpy(scc.data(), _scc, sizeof(std::complex) * scc.size(), hipMemcpyDeviceToHost)); - base_device::DEVICE_CPU* cpu_ctx = {}; - hegvd_op, base_device::DEVICE_CPU>()(cpu_ctx, - nstart, - ldh, - hcc.data(), - scc.data(), - eigenvalue.data(), - vcc.data()); - hipErrcheck(hipMemcpy(_vcc, vcc.data(), sizeof(std::complex) * vcc.size(), hipMemcpyHostToDevice)); - hipErrcheck(hipMemcpy(_eigenvalue, eigenvalue.data(), sizeof(float) * eigenvalue.size(), hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + hcc.data(), + _hcc, + sizeof(std::complex) * hcc.size(), + hipMemcpyDeviceToHost)); + hipErrcheck(hipMemcpy( + scc.data(), + _scc, + sizeof(std::complex) * scc.size(), + hipMemcpyDeviceToHost)); + base_device::DEVICE_CPU *cpu_ctx = {}; + hegvd_op, base_device::DEVICE_CPU>()( + cpu_ctx, + nstart, + ldh, + hcc.data(), + scc.data(), + eigenvalue.data(), + vcc.data()); + hipErrcheck(hipMemcpy( + _vcc, + vcc.data(), + sizeof(std::complex) * vcc.size(), + hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + _eigenvalue, + eigenvalue.data(), + sizeof(float) * eigenvalue.size(), + hipMemcpyHostToDevice)); } - - } template <> -void hegvd_op, base_device::DEVICE_GPU>::operator()(const base_device::DEVICE_GPU* ctx, - const int nstart, - const int ldh, - const std::complex* _hcc, - const std::complex* _scc, - double* _eigenvalue, - std::complex* _vcc - ) +void hegvd_op, base_device::DEVICE_GPU>::operator()( + const base_device::DEVICE_GPU *ctx, + const int nstart, + const int ldh, + const std::complex *_hcc, + const std::complex *_scc, + double *_eigenvalue, + std::complex *_vcc) { // copied from ../cuda/hegvd_op.cu, "hegvd_op" // assert(nstart == ldh); // save a copy of scc in case the diagonalization fails - if (nstart > N_DCU){ + if (nstart > N_DCU) + { std::vector> scc(nstart * nstart, {0, 0}); - hipErrcheck(hipMemcpy(scc.data(), _scc, sizeof(std::complex) * scc.size(), hipMemcpyDeviceToHost)); - - hipErrcheck(hipMemcpy(_vcc, _hcc, sizeof(std::complex) * ldh * nstart, hipMemcpyDeviceToDevice)); + hipErrcheck(hipMemcpy( + scc.data(), + _scc, + sizeof(std::complex) * scc.size(), + hipMemcpyDeviceToHost)); + + hipErrcheck(hipMemcpy( + _vcc, + _hcc, + sizeof(std::complex) * ldh * nstart, + hipMemcpyDeviceToDevice)); // now vcc contains hcc // prepare some values for hipsolverDnZhegvd_bufferSize - int * devInfo = nullptr; + int *devInfo = nullptr; int lwork = 0, info_gpu = 0; - double2 * work = nullptr; - hipErrcheck(hipMalloc((void**)&devInfo, sizeof(int))); + double2 *work = nullptr; + hipErrcheck(hipMalloc((void **)&devInfo, sizeof(int))); hipsolverFillMode_t uplo = HIPSOLVER_FILL_MODE_UPPER; // calculate the sizes needed for pre-allocated buffer. hipsolverErrcheck(hipsolverDnZhegvd_bufferSize( - hipsolver_H, HIPSOLVER_EIG_TYPE_1, HIPSOLVER_EIG_MODE_VECTOR, uplo, + hipsolver_H, + HIPSOLVER_EIG_TYPE_1, + HIPSOLVER_EIG_MODE_VECTOR, + uplo, nstart, - reinterpret_cast(_vcc), ldh, - reinterpret_cast(_scc), ldh, + reinterpret_cast(_vcc), + ldh, + reinterpret_cast(_scc), + ldh, _eigenvalue, &lwork)); // allocate memery - hipErrcheck(hipMalloc((void**)&work, sizeof(double2) * lwork)); + hipErrcheck(hipMalloc((void **)&work, sizeof(double2) * lwork)); // compute eigenvalues and eigenvectors. hipsolverErrcheck(hipsolverDnZhegvd( - hipsolver_H, HIPSOLVER_EIG_TYPE_1, HIPSOLVER_EIG_MODE_VECTOR, uplo, + hipsolver_H, + HIPSOLVER_EIG_TYPE_1, + HIPSOLVER_EIG_MODE_VECTOR, + uplo, nstart, - reinterpret_cast(_vcc), ldh, - const_cast(reinterpret_cast(_scc)), ldh, + reinterpret_cast(_vcc), + ldh, + const_cast(reinterpret_cast(_scc)), + ldh, _eigenvalue, - work, lwork, devInfo)); - - hipErrcheck(hipMemcpy(&info_gpu, devInfo, sizeof(int), hipMemcpyDeviceToHost)); + work, + lwork, + devInfo)); + + hipErrcheck(hipMemcpy( + &info_gpu, + devInfo, + sizeof(int), + hipMemcpyDeviceToHost)); // free the buffer hipErrcheck(hipFree(work)); hipErrcheck(hipFree(devInfo)); } // if(fail_info != nullptr) *fail_info = info_gpu; - else{ + else + { std::vector> hcc(nstart * nstart, {0, 0}); std::vector> scc(nstart * nstart, {0, 0}); std::vector> vcc(nstart * nstart, {0, 0}); std::vector eigenvalue(nstart, 0); - hipErrcheck(hipMemcpy(hcc.data(), _hcc, sizeof(std::complex) * hcc.size(), hipMemcpyDeviceToHost)); - hipErrcheck(hipMemcpy(scc.data(), _scc, sizeof(std::complex) * scc.size(), hipMemcpyDeviceToHost)); - base_device::DEVICE_CPU* cpu_ctx = {}; - hegvd_op, base_device::DEVICE_CPU>()(cpu_ctx, - nstart, - ldh, - hcc.data(), - scc.data(), - eigenvalue.data(), - vcc.data()); - hipErrcheck(hipMemcpy(_vcc, vcc.data(), sizeof(std::complex) * vcc.size(), hipMemcpyHostToDevice)); - hipErrcheck(hipMemcpy(_eigenvalue, eigenvalue.data(), sizeof(double) * eigenvalue.size(), hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + hcc.data(), + _hcc, + sizeof(std::complex) * hcc.size(), + hipMemcpyDeviceToHost)); + hipErrcheck(hipMemcpy( + scc.data(), + _scc, + sizeof(std::complex) * scc.size(), + hipMemcpyDeviceToHost)); + base_device::DEVICE_CPU *cpu_ctx = {}; + hegvd_op, base_device::DEVICE_CPU>()( + cpu_ctx, + nstart, + ldh, + hcc.data(), + scc.data(), + eigenvalue.data(), + vcc.data()); + hipErrcheck(hipMemcpy( + _vcc, + vcc.data(), + sizeof(std::complex) * vcc.size(), + hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + _eigenvalue, + eigenvalue.data(), + sizeof(double) * eigenvalue.size(), + hipMemcpyHostToDevice)); } - - - - - - - } #ifdef __LCAO template <> -void heevx_op::operator()(const base_device::DEVICE_GPU* ctx, - const int nstart, - const int ldh, - const double* _hcc, - const int m, - double* _eigenvalue, - double* _vcc) +void heevx_op::operator()( + const base_device::DEVICE_GPU *ctx, + const int nstart, + const int ldh, + const double *_hcc, + const int m, + double *_eigenvalue, + double *_vcc) { std::vector hcc(ldh * ldh, 0.0); std::vector vcc(ldh * ldh, 0.0); std::vector eigenvalue(ldh, 0); - hipErrcheck(hipMemcpy(hcc.data(), _hcc, sizeof(double) * hcc.size(), hipMemcpyDeviceToHost)); - base_device::DEVICE_CPU* cpu_ctx = {}; - heevx_op()(cpu_ctx, nstart, ldh, hcc.data(), m, eigenvalue.data(), vcc.data()); - hipErrcheck(hipMemcpy(_vcc, vcc.data(), sizeof(double) * vcc.size(), hipMemcpyHostToDevice)); - hipErrcheck(hipMemcpy(_eigenvalue, eigenvalue.data(), sizeof(double) * eigenvalue.size(), hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + hcc.data(), + _hcc, + sizeof(double) * hcc.size(), + hipMemcpyDeviceToHost)); + base_device::DEVICE_CPU *cpu_ctx = {}; + heevx_op()( + cpu_ctx, + nstart, + ldh, + hcc.data(), + m, + eigenvalue.data(), + vcc.data()); + hipErrcheck(hipMemcpy( + _vcc, + vcc.data(), + sizeof(double) * vcc.size(), + hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + _eigenvalue, + eigenvalue.data(), + sizeof(double) * eigenvalue.size(), + hipMemcpyHostToDevice)); } #endif // __LCAO template <> -void heevx_op, base_device::DEVICE_GPU>::operator()(const base_device::DEVICE_GPU* ctx, - const int nstart, - const int ldh, - const std::complex* _hcc, - const int m, - float* _eigenvalue, - std::complex* _vcc) +void heevx_op, base_device::DEVICE_GPU>::operator()( + const base_device::DEVICE_GPU *ctx, + const int nstart, + const int ldh, + const std::complex *_hcc, + const int m, + float *_eigenvalue, + std::complex *_vcc) { std::vector> hcc(ldh * ldh, {0, 0}); std::vector> vcc(ldh * ldh, {0, 0}); std::vector eigenvalue(ldh, 0); - hipErrcheck(hipMemcpy(hcc.data(), _hcc, sizeof(std::complex) * hcc.size(), hipMemcpyDeviceToHost)); - base_device::DEVICE_CPU* cpu_ctx = {}; - heevx_op, base_device::DEVICE_CPU>()(cpu_ctx, - nstart, - ldh, - hcc.data(), - m, - eigenvalue.data(), - vcc.data()); - hipErrcheck(hipMemcpy(_vcc, vcc.data(), sizeof(std::complex) * vcc.size(), hipMemcpyHostToDevice)); - hipErrcheck(hipMemcpy(_eigenvalue, eigenvalue.data(), sizeof(float) * eigenvalue.size(), hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + hcc.data(), + _hcc, + sizeof(std::complex) * hcc.size(), + hipMemcpyDeviceToHost)); + base_device::DEVICE_CPU *cpu_ctx = {}; + heevx_op, base_device::DEVICE_CPU>()( + cpu_ctx, + nstart, + ldh, + hcc.data(), + m, + eigenvalue.data(), + vcc.data()); + hipErrcheck(hipMemcpy( + _vcc, + vcc.data(), + sizeof(std::complex) * vcc.size(), + hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + _eigenvalue, + eigenvalue.data(), + sizeof(float) * eigenvalue.size(), + hipMemcpyHostToDevice)); } template <> -void heevx_op, base_device::DEVICE_GPU>::operator()(const base_device::DEVICE_GPU* ctx, - const int nstart, - const int ldh, - const std::complex* _hcc, - const int m, - double* _eigenvalue, - std::complex* _vcc) +void heevx_op, base_device::DEVICE_GPU>::operator()( + const base_device::DEVICE_GPU *ctx, + const int nstart, + const int ldh, + const std::complex *_hcc, + const int m, + double *_eigenvalue, + std::complex *_vcc) { std::vector> hcc(ldh * ldh, {0, 0}); std::vector> vcc(ldh * ldh, {0, 0}); std::vector eigenvalue(ldh, 0); - hipErrcheck(hipMemcpy(hcc.data(), _hcc, sizeof(std::complex) * hcc.size(), hipMemcpyDeviceToHost)); - base_device::DEVICE_CPU* cpu_ctx = {}; - heevx_op, base_device::DEVICE_CPU>()(cpu_ctx, - nstart, - ldh, - hcc.data(), - m, - eigenvalue.data(), - vcc.data()); - hipErrcheck(hipMemcpy(_vcc, vcc.data(), sizeof(std::complex) * vcc.size(), hipMemcpyHostToDevice)); - hipErrcheck(hipMemcpy(_eigenvalue, eigenvalue.data(), sizeof(double) * eigenvalue.size(), hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + hcc.data(), + _hcc, + sizeof(std::complex) * hcc.size(), + hipMemcpyDeviceToHost)); + base_device::DEVICE_CPU *cpu_ctx = {}; + heevx_op, base_device::DEVICE_CPU>()( + cpu_ctx, + nstart, + ldh, + hcc.data(), + m, + eigenvalue.data(), + vcc.data()); + hipErrcheck(hipMemcpy( + _vcc, + vcc.data(), + sizeof(std::complex) * vcc.size(), + hipMemcpyHostToDevice)); + hipErrcheck(hipMemcpy( + _eigenvalue, + eigenvalue.data(), + sizeof(double) * eigenvalue.size(), + hipMemcpyHostToDevice)); } template <> -void hegvx_op, base_device::DEVICE_GPU>::operator()(const base_device::DEVICE_GPU* d, - const int nbase, - const int ldh, - std::complex* hcc, - std::complex* scc, - const int m, - float* eigenvalue, - std::complex* vcc) +void hegvx_op, base_device::DEVICE_GPU>::operator()( + const base_device::DEVICE_GPU *d, + const int nbase, + const int ldh, + std::complex *hcc, + std::complex *scc, + const int m, + float *eigenvalue, + std::complex *vcc) { } template <> -void hegvx_op, base_device::DEVICE_GPU>::operator()(const base_device::DEVICE_GPU* d, - const int nbase, - const int ldh, - std::complex* hcc, - std::complex* scc, - const int m, - double* eigenvalue, - std::complex* vcc) +void hegvx_op, base_device::DEVICE_GPU>::operator()( + const base_device::DEVICE_GPU *d, + const int nbase, + const int ldh, + std::complex *hcc, + std::complex *scc, + const int m, + double *eigenvalue, + std::complex *vcc) { } #ifdef __LCAO template <> -void hegvx_op::operator()(const base_device::DEVICE_GPU* d, - const int nbase, - const int ldh, - double* hcc, - double* scc, - const int m, - double* eigenvalue, - double* vcc) +void hegvx_op::operator()( + const base_device::DEVICE_GPU *d, + const int nbase, + const int ldh, + double *hcc, + double *scc, + const int m, + double *eigenvalue, + double *vcc) { } #endif // __LCAO diff --git a/source/source_hsolver/kernels/test/math_hegvd_test.cpp b/source/source_hsolver/kernels/test/math_hegvd_test.cpp index adf2ab95eef..48714d33312 100644 --- a/source/source_hsolver/kernels/test/math_hegvd_test.cpp +++ b/source/source_hsolver/kernels/test/math_hegvd_test.cpp @@ -1,7 +1,7 @@ #include "source_base/complexmatrix.h" +#include "source_base/kernels/math_kernel_op.h" #include "source_base/module_device/memory_op.h" #include "source_hsolver/kernels/hegvd_op.h" -#include "source_base/kernels/math_kernel_op.h" #include #include @@ -18,35 +18,23 @@ class TestModuleHsolverMathDngvd : public ::testing::Test using resize_memory_op_D = base_device::memory::resize_memory_op; using delete_memory_op_D = base_device::memory::delete_memory_op; // from CPU to GPU - using synchronize_memory_op_C2G_Z = base_device::memory:: - synchronize_memory_op, base_device::DEVICE_GPU, base_device::DEVICE_CPU>; - using synchronize_memory_op_C2G_D - = base_device::memory::synchronize_memory_op; - using synchronize_memory_op_G2C_Z = base_device::memory:: - synchronize_memory_op, base_device::DEVICE_CPU, base_device::DEVICE_GPU>; - using synchronize_memory_op_G2C_D - = base_device::memory::synchronize_memory_op; - - const base_device::DEVICE_CPU* cpu_ctx = {}; - const base_device::DEVICE_GPU* gpu_ctx = {}; + using synchronize_memory_op_C2G_Z = base_device::memory::synchronize_memory_op< + std::complex, base_device::DEVICE_GPU, base_device::DEVICE_CPU>; + using synchronize_memory_op_C2G_D = + base_device::memory::synchronize_memory_op; + using synchronize_memory_op_G2C_Z = base_device::memory::synchronize_memory_op< + std::complex, base_device::DEVICE_CPU, base_device::DEVICE_GPU>; + using synchronize_memory_op_G2C_D = + base_device::memory::synchronize_memory_op; + + const base_device::DEVICE_CPU *cpu_ctx = {}; + const base_device::DEVICE_GPU *gpu_ctx = {}; // prepare A & B in CPU std::vector> matrix_A = { - {-0.351417, -1.73472}, - {-8.32667, 2.3744}, - {4.16334, 3.64292}, - {5.20417, -3.85976}, - {-8.32667, -2.3744}, - {0.551651, -2.60209}, - {2.08167, 1.9082}, - {-6.93889, 1.04083}, - {4.16334, -3.64292}, - {2.08167, -1.9082}, - {0.551651, -2.25514}, - {-1.31839, 5.20417}, - {5.20417, 3.85976}, - {-6.93889, -1.04083}, - {-1.31839, -5.20417}, + {-0.351417, -1.73472}, {-8.32667, 2.3744}, {4.16334, 3.64292}, {5.20417, -3.85976}, {-8.32667, -2.3744}, + {0.551651, -2.60209}, {2.08167, 1.9082}, {-6.93889, 1.04083}, {4.16334, -3.64292}, {2.08167, -1.9082}, + {0.551651, -2.25514}, {-1.31839, 5.20417}, {5.20417, 3.85976}, {-6.93889, -1.04083}, {-1.31839, -5.20417}, {0.551651, -3.64292} // {-4.280e-01,3.084e-17}, {-1.288e-16,9.021e-17}, {5.204e-18,-3.990e-17}, {-5.204e-17,-2.776e-17}, @@ -60,70 +48,31 @@ class TestModuleHsolverMathDngvd : public ::testing::Test // {-5.204e-17,2.776e-17}, {8.327e-17,1.041e-17}, {-6.939e-18,-4.163e-17}, {4.574e-01,-3.451e-17} }; std::vector> matrix_B = { - {1, 0}, - {0, 0}, - {0, 0}, - {0, 0}, - {0, 0}, - {1, 0}, + {1, 0}, {0, 0}, {0, 0}, {0, 0}, {0, 0}, {1, 0}, {0, 0}, {0, 0}, {0, 0}, {0, 0}, {1, 0}, {0, 0}, {0, 0}, {0, 0}, {0, 0}, - {0, 0}, - {0, 0}, - {0, 0}, - {1, 0}, - {0, 0}, - {0, 0}, - {0, 0}, - {0, 0}, - {1, 0} - // {1.000e+00,0.000e+00}, {-7.069e-17,-3.123e-17}, {-5.204e-17,0.000e+00}, {5.551e-17,-2.082e-17}, - // {-7.069e-17,3.123e-17}, {1.000e+00,0.000e+00}, {1.110e-16,-7.286e-17}, {8.327e-17,-1.110e-16}, - // {-5.204e-17,0.000e+00}, {1.110e-16,7.286e-17}, {1.000e+00,0.000e+00}, {-9.714e-17,2.776e-17}, - // {5.551e-17,2.082e-17}, {8.327e-17,1.110e-16}, {-9.714e-17,-2.776e-17}, {1.000e+00,0.000e+00} + {1, 0} // {1.000e+00,0.000e+00}, {-7.069e-17,-3.123e-17}, {-5.204e-17,0.000e+00}, {5.551e-17,-2.082e-17}, + // {-7.069e-17,3.123e-17}, {1.000e+00,0.000e+00}, {1.110e-16,-7.286e-17}, {8.327e-17,-1.110e-16}, + // {-5.204e-17,0.000e+00}, {1.110e-16,7.286e-17}, {1.000e+00,0.000e+00}, {-9.714e-17,2.776e-17}, + // {5.551e-17,2.082e-17}, {8.327e-17,1.110e-16}, {-9.714e-17,-2.776e-17}, {1.000e+00,0.000e+00} }; const int matrix_size = 16; // prepare W & V in CPU in dngv_op - std::vector W_dngv_op = {0.0, 0.0, 0.0, 0.0}; - std::vector> matrix_V_dngv_op = {{0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}}; + std::vector W_dngv_op = {0.0, + 0.0, + 0.0, + 0.0}; + std::vector> matrix_V_dngv_op = {{0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, + {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, + {0.0, 0.0}, {0.0, 0.0}}; // prepare W & V in CPU in dngvx_op std::vector W_DNGVX = {0.0, 0.0}; - std::vector> matrix_V_DNGVX = {{0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}, - {0.0, 0.0}}; + std::vector> matrix_V_DNGVX = {{0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, + {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, {0.0, 0.0}, + {0.0, 0.0}, {0.0, 0.0}}; }; - - #if __UT_USE_CUDA || __UT_USE_ROCM TEST_F(TestModuleHsolverMathDngvd, transpose_gpu) @@ -138,16 +87,17 @@ TEST_F(TestModuleHsolverMathDngvd, transpose_gpu) {4.16334, 3.64292}, // {-0.351417,-1.73472}, {-8.32667,2.3744}, {4.16334,3.64292} }; - std::complex* device_transpose = nullptr; + std::complex *device_transpose = nullptr; resize_memory_op_Z()(device_transpose, matrix_size); synchronize_memory_op_C2G_Z()(device_transpose, transpose.data(), transpose.size()); // run ModuleBase::createGpuBlasHandle(); - ModuleBase::matrixTranspose_op, base_device::DEVICE_GPU>()(2, - 3, - device_transpose, - device_transpose); + ModuleBase::matrixTranspose_op, base_device::DEVICE_GPU>()( + 2, + 3, + device_transpose, + device_transpose); ModuleBase::destoryBLAShandle(); // copy transpose data from GPU to CPU diff --git a/source/source_hsolver/kernels/test/perf_math_kernel.cpp b/source/source_hsolver/kernels/test/perf_math_kernel.cpp index c9258d30fd9..84b95dfeaea 100644 --- a/source/source_hsolver/kernels/test/perf_math_kernel.cpp +++ b/source/source_hsolver/kernels/test/perf_math_kernel.cpp @@ -1,20 +1,20 @@ #include "source_base/constants.h" -#include "source_base/module_device/memory_op.h" #include "source_base/kernels/math_kernel_op.h" +#include "source_base/module_device/memory_op.h" -#include #include -#include +#include #include +#include #include -#include +#include /************************************************ -* performace test of class math_kernel_op -***********************************************/ + * performace test of class math_kernel_op + ***********************************************/ /** - * Tested function: + * Tested function: * - zdot_real_cpu_op * - vector_mul_real_op_cpu * - vector_mul_vector_op_cpu @@ -32,54 +32,56 @@ * - scal_gpu */ -class PerfModuleHsolverMathKernel : public benchmark::Fixture { - public: - +class PerfModuleHsolverMathKernel : public benchmark::Fixture +{ + public: // DEVICE SYMBOL - const base_device::DEVICE_CPU* cpu_ctx = {}; + const base_device::DEVICE_CPU *cpu_ctx = {}; int dim_vector = 1; - std::complex* test_zvector_a = nullptr; - std::complex* test_zvector_b = nullptr; - std::complex* result_zvector = nullptr; + std::complex *test_zvector_a = nullptr; + std::complex *test_zvector_b = nullptr; + std::complex *result_zvector = nullptr; - double* test_dvector_a = nullptr; + double *test_dvector_a = nullptr; double dconstant_a = 1.0; double dconstant_b = 1.0; - std::complex zconstant_a = {1.0,1.0}; + std::complex zconstant_a = {1.0, 1.0}; - #if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM - const base_device::DEVICE_GPU * gpu_ctx = {}; +#if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM + const base_device::DEVICE_GPU *gpu_ctx = {}; // from CPU to GPU - using synchronize_memory_op - = base_device::memory::synchronize_memory_op, base_device::DEVICE_GPU, base_device::DEVICE_CPU>; + using synchronize_memory_op = base_device::memory::synchronize_memory_op< + std::complex, base_device::DEVICE_GPU, base_device::DEVICE_CPU>; // form GPU to CPU - using synchronize_memory_op_gpu - = base_device::memory::synchronize_memory_op, base_device::DEVICE_CPU, base_device::DEVICE_GPU>; + using synchronize_memory_op_gpu = base_device::memory::synchronize_memory_op< + std::complex, base_device::DEVICE_CPU, base_device::DEVICE_GPU>; using resize_memory_op = base_device::memory::resize_memory_op, base_device::DEVICE_GPU>; using delete_memory_op = base_device::memory::delete_memory_op, base_device::DEVICE_GPU>; using resize_memory_op_double = base_device::memory::resize_memory_op; using delete_memory_op_double = base_device::memory::delete_memory_op; - using synchronize_memory_op_double = base_device::memory::synchronize_memory_op; + using synchronize_memory_op_double = + base_device::memory::synchronize_memory_op; using set_memory_op = base_device::memory::set_memory_op, base_device::DEVICE_GPU>; using set_memory_op_double = base_device::memory::set_memory_op; - std::complex* test_zvector_a_gpu = nullptr; - std::complex* test_zvector_b_gpu = nullptr; - std::complex* result_zvector_gpu = nullptr; + std::complex *test_zvector_a_gpu = nullptr; + std::complex *test_zvector_b_gpu = nullptr; + std::complex *result_zvector_gpu = nullptr; - double* test_dvector_a_gpu = nullptr; + double *test_dvector_a_gpu = nullptr; - #endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM +#endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM - void SetUp(const benchmark::State& state){ + void SetUp(const benchmark::State &state) + { dim_vector = state.range(0); // Generate vectors with different diminsions (1,10,100,...,1e6) // This should be complex vectors @@ -90,18 +92,21 @@ class PerfModuleHsolverMathKernel : public benchmark::Fixture { // The following is double vectors test_dvector_a = new double[dim_vector + 1]; - - for (int i=0;i{(double)rand()+(double)rand()/(RAND_MAX+1.0),(double)rand()+(double)rand()/(RAND_MAX+1.0)}; - test_zvector_b[i] = std::complex{(double)rand()+(double)rand()/(RAND_MAX+1.0),(double)rand()+(double)rand()/(RAND_MAX+1.0)}; - test_dvector_a[i] = (double)rand()+(double)rand()/(RAND_MAX+1.0); + test_zvector_a[i] = std::complex{ + (double)rand() + (double)rand() / (RAND_MAX + 1.0), (double)rand() + (double)rand() / (RAND_MAX + 1.0)}; + test_zvector_b[i] = std::complex{ + (double)rand() + (double)rand() / (RAND_MAX + 1.0), (double)rand() + (double)rand() / (RAND_MAX + 1.0)}; + test_dvector_a[i] = (double)rand() + (double)rand() / (RAND_MAX + 1.0); } - dconstant_a = (double)rand()+(double)rand()/(RAND_MAX+1.0); - dconstant_b = (double)rand()+(double)rand()/(RAND_MAX+1.0); + dconstant_a = (double)rand() + (double)rand() / (RAND_MAX + 1.0); + dconstant_b = (double)rand() + (double)rand() / (RAND_MAX + 1.0); - zconstant_a = std::complex{(double)rand()+(double)rand()/(RAND_MAX+1.0),(double)rand()+(double)rand()/(RAND_MAX+1.0)}; + zconstant_a = std::complex{ + (double)rand() + (double)rand() / (RAND_MAX + 1.0), (double)rand() + (double)rand() / (RAND_MAX + 1.0)}; #if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM resize_memory_op()(test_zvector_a_gpu, dim_vector); @@ -115,10 +120,10 @@ class PerfModuleHsolverMathKernel : public benchmark::Fixture { ModuleBase::createGpuBlasHandle(); - #endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM } - void TearDown(const benchmark::State& state){ + void TearDown(const benchmark::State &state) + { delete[] test_zvector_a; delete[] test_zvector_b; delete[] result_zvector; @@ -128,18 +133,17 @@ class PerfModuleHsolverMathKernel : public benchmark::Fixture { #endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM } - // OPs need benchmark // CPU operator using zdot_real_cpu_op = ModuleBase::dot_real_op, base_device::DEVICE_CPU>; - + using vector_mul_real_op_cpu = ModuleBase::vector_mul_real_op, base_device::DEVICE_CPU>; - using vector_mul_vector_op_cpu - = ModuleBase::vector_mul_vector_op, base_device::DEVICE_CPU, double>; - using vector_div_vector_op_cpu - = ModuleBase::vector_div_vector_op, base_device::DEVICE_CPU, double>; - using vector_add_vector_op_cpu - = ModuleBase::vector_add_vector_op, base_device::DEVICE_CPU, double>; + using vector_mul_vector_op_cpu = + ModuleBase::vector_mul_vector_op, base_device::DEVICE_CPU, double>; + using vector_div_vector_op_cpu = + ModuleBase::vector_div_vector_op, base_device::DEVICE_CPU, double>; + using vector_add_vector_op_cpu = + ModuleBase::vector_add_vector_op, base_device::DEVICE_CPU, double>; using axpy_op_cpu = ModuleBase::axpy_op, base_device::DEVICE_CPU>; using scal_op_cpu = ModuleBase::scal_op; using gemv_op_cpu = ModuleBase::gemv_op, base_device::DEVICE_CPU>; @@ -150,71 +154,135 @@ class PerfModuleHsolverMathKernel : public benchmark::Fixture { using zdot_real_gpu_op = ModuleBase::dot_real_op, base_device::DEVICE_GPU>; using vector_mul_real_op_gpu = ModuleBase::vector_mul_real_op, base_device::DEVICE_GPU>; - using vector_mul_vector_op_gpu - = ModuleBase::vector_mul_vector_op, base_device::DEVICE_GPU, double>; - using vector_div_vector_op_gpu - = ModuleBase::vector_div_vector_op, base_device::DEVICE_GPU, double>; - using vector_add_vector_op_gpu - = ModuleBase::vector_add_vector_op, base_device::DEVICE_GPU, double>; + using vector_mul_vector_op_gpu = + ModuleBase::vector_mul_vector_op, base_device::DEVICE_GPU, double>; + using vector_div_vector_op_gpu = + ModuleBase::vector_div_vector_op, base_device::DEVICE_GPU, double>; + using vector_add_vector_op_gpu = + ModuleBase::vector_add_vector_op, base_device::DEVICE_GPU, double>; using axpy_op_gpu = ModuleBase::axpy_op, base_device::DEVICE_GPU>; using scal_op_gpu = ModuleBase::scal_op; #endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM }; - -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_zdot_real_cpu_op)(benchmark::State& state) { - for (auto _ : state) { - double result = zdot_real_cpu_op()(dim_vector, test_zvector_a, test_zvector_b, false); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_zdot_real_cpu_op)(benchmark::State &state) +{ + for (auto _ : state) + { + double result = zdot_real_cpu_op()( + dim_vector, + test_zvector_a, + test_zvector_b, + false); } } - -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_mul_real_op_cpu)(benchmark::State& state) { - for (auto _ : state) { - vector_mul_real_op_cpu()(dim_vector, result_zvector, test_zvector_a, dconstant_a); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_mul_real_op_cpu)(benchmark::State &state) +{ + for (auto _ : state) + { + vector_mul_real_op_cpu()( + dim_vector, + result_zvector, + test_zvector_a, + dconstant_a); } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_mul_vector_op_cpu)(benchmark::State& state) { - for (auto _ : state) { - vector_mul_vector_op_cpu()(dim_vector, result_zvector, test_zvector_a, test_dvector_a, false); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_mul_vector_op_cpu)(benchmark::State &state) +{ + for (auto _ : state) + { + vector_mul_vector_op_cpu()( + dim_vector, + result_zvector, + test_zvector_a, + test_dvector_a, + false); } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_cpu)(benchmark::State& state) { - for (auto _ : state) { - vector_div_vector_op_cpu()(dim_vector, result_zvector, test_zvector_a, test_dvector_a); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_cpu)(benchmark::State &state) +{ + for (auto _ : state) + { + vector_div_vector_op_cpu()( + dim_vector, + result_zvector, + test_zvector_a, + test_dvector_a); } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_cpu)(benchmark::State& state) { - for (auto _ : state) { - vector_add_vector_op_cpu()(dim_vector, result_zvector, test_zvector_a, dconstant_a ,test_zvector_b, dconstant_b); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_cpu)(benchmark::State &state) +{ + for (auto _ : state) + { + vector_add_vector_op_cpu()( + dim_vector, + result_zvector, + test_zvector_a, + dconstant_a, + test_zvector_b, + dconstant_b); } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_axpy_op_cpu)(benchmark::State& state) { - for (auto _ : state) { - axpy_op_cpu()(dim_vector, &zconstant_a, test_zvector_a, 1 ,test_zvector_b, 1); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_axpy_op_cpu)(benchmark::State &state) +{ + for (auto _ : state) + { + axpy_op_cpu()( + dim_vector, + &zconstant_a, + test_zvector_a, + 1, + test_zvector_b, + 1); } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_scal_op_cpu)(benchmark::State& state) { - for (auto _ : state) { - scal_op_cpu()(dim_vector, &zconstant_a, test_zvector_a, 1); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_scal_op_cpu)(benchmark::State &state) +{ + for (auto _ : state) + { + scal_op_cpu()( + dim_vector, + &zconstant_a, + test_zvector_a, + 1); } } - -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_zdot_real_cpu_op)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_real_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_vector_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_axpy_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_scal_op_cpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); - +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_zdot_real_cpu_op) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_real_op_cpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_vector_op_cpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_cpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_cpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_axpy_op_cpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_scal_op_cpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); #if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM @@ -233,60 +301,128 @@ BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_zdot_real_gpu_op)(benchmark:: } */ -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_zdot_real_gpu_op)(benchmark::State& state) { - for (auto _ : state) { - double result = zdot_real_gpu_op()(dim_vector, test_zvector_a_gpu, test_zvector_b_gpu, false); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_zdot_real_gpu_op)(benchmark::State &state) +{ + for (auto _ : state) + { + double result = zdot_real_gpu_op()( + dim_vector, + test_zvector_a_gpu, + test_zvector_b_gpu, + false); } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_mul_real_op_gpu)(benchmark::State& state) { - for (auto _ : state) { - vector_mul_real_op_gpu()(dim_vector, result_zvector_gpu, test_zvector_a_gpu, dconstant_a); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_mul_real_op_gpu)(benchmark::State &state) +{ + for (auto _ : state) + { + vector_mul_real_op_gpu()( + dim_vector, + result_zvector_gpu, + test_zvector_a_gpu, + dconstant_a); } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_mul_vector_op_gpu)(benchmark::State& state) { - for (auto _ : state) { - vector_mul_vector_op_gpu()(dim_vector, result_zvector_gpu, test_zvector_a_gpu, test_dvector_a_gpu, false); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_mul_vector_op_gpu)(benchmark::State &state) +{ + for (auto _ : state) + { + vector_mul_vector_op_gpu()( + dim_vector, + result_zvector_gpu, + test_zvector_a_gpu, + test_dvector_a_gpu, + false); } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_gpu)(benchmark::State& state) { - for (auto _ : state) { - vector_div_vector_op_gpu()(dim_vector, result_zvector_gpu, test_zvector_a_gpu, test_dvector_a_gpu); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_gpu)(benchmark::State &state) +{ + for (auto _ : state) + { + vector_div_vector_op_gpu()( + dim_vector, + result_zvector_gpu, + test_zvector_a_gpu, + test_dvector_a_gpu); } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_gpu)(benchmark::State& state) { - for (auto _ : state) { - vector_add_vector_op_gpu()(dim_vector, result_zvector_gpu, test_zvector_a_gpu, dconstant_a ,test_zvector_b_gpu, dconstant_b); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_gpu)(benchmark::State &state) +{ + for (auto _ : state) + { + vector_add_vector_op_gpu()( + dim_vector, + result_zvector_gpu, + test_zvector_a_gpu, + dconstant_a, + test_zvector_b_gpu, + dconstant_b); } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_axpy_op_gpu)(benchmark::State& state) { - for (auto _ : state) { - axpy_op_gpu()(dim_vector, &zconstant_a, test_zvector_a_gpu, 1 ,test_zvector_b_gpu, 1); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_axpy_op_gpu)(benchmark::State &state) +{ + for (auto _ : state) + { + axpy_op_gpu()( + dim_vector, + &zconstant_a, + test_zvector_a_gpu, + 1, + test_zvector_b_gpu, + 1); } } -BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_scal_op_gpu)(benchmark::State& state) { - for (auto _ : state) { - scal_op_gpu()(dim_vector, &zconstant_a, test_zvector_a_gpu, 1); +BENCHMARK_DEFINE_F(PerfModuleHsolverMathKernel, BM_scal_op_gpu)(benchmark::State &state) +{ + for (auto _ : state) + { + scal_op_gpu()( + dim_vector, + &zconstant_a, + test_zvector_a_gpu, + 1); } } // If you want to use manual timer, you can refer to this. -// BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_zdot_real_gpu_op)->RangeMultiplier(10)->Range(1,10e6)->UseManualTime()->Unit(benchmark::kMicrosecond); - -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_zdot_real_gpu_op)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_real_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_vector_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_axpy_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); -BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_scal_op_gpu)->RangeMultiplier(10)->Range(1,10e6)->Unit(benchmark::kMicrosecond); +// BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, +// BM_zdot_real_gpu_op)->RangeMultiplier(10)->Range(1,10e6)->UseManualTime()->Unit(benchmark::kMicrosecond); + +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_zdot_real_gpu_op) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_real_op_gpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_mul_vector_op_gpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_div_vector_op_gpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_vector_add_vector_op_gpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_axpy_op_gpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); +BENCHMARK_REGISTER_F(PerfModuleHsolverMathKernel, BM_scal_op_gpu) + ->RangeMultiplier(10) + ->Range(1, 10e6) + ->Unit(benchmark::kMicrosecond); #endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM - BENCHMARK_MAIN(); diff --git a/source/source_hsolver/my_elpa.h b/source/source_hsolver/my_elpa.h index 65c57df4e0e..658e5856077 100644 --- a/source/source_hsolver/my_elpa.h +++ b/source/source_hsolver/my_elpa.h @@ -12,11 +12,11 @@ extern "C" { #endif -struct elpa_struct; -typedef struct elpa_struct *elpa_t; + struct elpa_struct; + typedef struct elpa_struct *elpa_t; -struct elpa_autotune_struct; -typedef struct elpa_autotune_struct *elpa_autotune_t; + struct elpa_autotune_struct; + typedef struct elpa_autotune_struct *elpa_autotune_t; #include #include @@ -33,14 +33,14 @@ typedef struct elpa_autotune_struct *elpa_autotune_t; #endif #include -#define ELPA_2STAGE_REAL_GPU ELPA_2STAGE_REAL_NVIDIA_GPU +#define ELPA_2STAGE_REAL_GPU ELPA_2STAGE_REAL_NVIDIA_GPU #define ELPA_2STAGE_COMPLEX_GPU ELPA_2STAGE_COMPLEX_NVIDIA_GPU -const char *elpa_strerr(int elpa_error); + const char *elpa_strerr(int elpa_error); #if ELPA_API_VERSION < 20221101 -}//extern "C" +} // extern "C" #endif #endif //__ELPA -#endif //ELPA_H +#endif // ELPA_H diff --git a/source/source_hsolver/para_lin_tf.cpp b/source/source_hsolver/para_lin_tf.cpp index a7093fa49d2..f8086d56eb6 100644 --- a/source/source_hsolver/para_lin_tf.cpp +++ b/source/source_hsolver/para_lin_tf.cpp @@ -9,8 +9,7 @@ #include namespace hsolver { -template -PLinearTransform::~PLinearTransform() +template PLinearTransform::~PLinearTransform() { #ifdef __MPI delmem_dev_op()(U_tmp_); @@ -19,14 +18,15 @@ PLinearTransform::~PLinearTransform() #endif } template -void PLinearTransform::set_dimension(const int nrowA, - const int ncolA, - const int ncolB, - const int LDA, +void PLinearTransform::set_dimension( + const int nrowA, + const int ncolA, + const int ncolB, + const int LDA, #ifdef __MPI - MPI_Comm col_world, + MPI_Comm col_world, #endif - const bool localU) + const bool localU) { this->nrowA = nrowA; this->ncolA = ncolA; @@ -40,7 +40,14 @@ void PLinearTransform::set_dimension(const int nrowA, { this->localU = localU; colA_loc.resize(nproc_col); - MPI_Allgather(&ncolA, 1, MPI_INT, colA_loc.data(), 1, MPI_INT, col_world); + MPI_Allgather( + &ncolA, + 1, + MPI_INT, + colA_loc.data(), + 1, + MPI_INT, + col_world); start_colA.resize(nproc_col); start_colA[0] = 0; for (int ip = 1; ip < nproc_col; ++ip) @@ -51,7 +58,14 @@ void PLinearTransform::set_dimension(const int nrowA, this->max_colA = *std::max_element(colA_loc.begin(), colA_loc.end()); std::vector colB_loc(nproc_col); - MPI_Allgather(&ncolB, 1, MPI_INT, colB_loc.data(), 1, MPI_INT, col_world); + MPI_Allgather( + &ncolB, + 1, + MPI_INT, + colB_loc.data(), + 1, + MPI_INT, + col_world); start_colB.resize(nproc_col); start_colB[0] = 0; for (int ip = 1; ip < nproc_col; ++ip) @@ -78,7 +92,12 @@ void PLinearTransform::set_dimension(const int nrowA, #endif } template -void PLinearTransform::act(const T alpha, const T* A, const T* U, const T beta, T* B) +void PLinearTransform::act( + const T alpha, + const T *A, + const T *U, + const T beta, + T *B) { ModuleBase::timer::start("PLinearTransform", "act"); #ifdef __MPI @@ -92,18 +111,44 @@ void PLinearTransform::act(const T alpha, const T* A, const T* U, con if (rank_col != ip) { int size = LDA * ncolA; - Parallel_Common::isend_dev(A, size, ip, 0, col_world, &requests[ip], isend_tmp_.data()); + Parallel_Common::isend_dev( + A, + size, + ip, + 0, + col_world, + &requests[ip], + isend_tmp_.data()); } } // local part const int start = this->localU ? 0 : start_colB[rank_col]; - const T* U_part = U + start_colA[rank_col] + start * ncolA_glo; - ModuleBase::matrixCopy()(ncolB, ncolA, U_part, ncolA_glo, U_tmp_, ncolA); - ModuleBase::gemm_op()('N', 'N', nrowA, ncolB, ncolA, &alpha, A, LDA, U_tmp_, ncolA, &beta, B, LDA); + const T *U_part = U + start_colA[rank_col] + start * ncolA_glo; + ModuleBase::matrixCopy()( + ncolB, + ncolA, + U_part, + ncolA_glo, + U_tmp_, + ncolA); + ModuleBase::gemm_op()( + 'N', + 'N', + nrowA, + ncolB, + ncolA, + &alpha, + A, + LDA, + U_tmp_, + ncolA, + &beta, + B, + LDA); // Receive - T* Atmp_device = nullptr; + T *Atmp_device = nullptr; if (std::is_same::value) { Atmp_device = A_tmp_device_; @@ -118,28 +163,48 @@ void PLinearTransform::act(const T alpha, const T* A, const T* U, con { T zero = 0.0; const int ncolA_ip = colA_loc[ip]; - const T* U_part = U + start_colA[ip] + start * ncolA_glo; - ModuleBase::matrixCopy()(ncolB, ncolA_ip, U_part, ncolA_glo, U_tmp_, ncolA_ip); + const T *U_part = U + start_colA[ip] + start * ncolA_glo; + ModuleBase::matrixCopy()( + ncolB, + ncolA_ip, + U_part, + ncolA_glo, + U_tmp_, + ncolA_ip); int size = LDA * ncolA_ip; MPI_Status status; - Parallel_Common::recv_dev(Atmp_device, size, ip, 0, col_world, &status, A_tmp_.data()); - ModuleBase::gemm_op()('N', - 'N', - nrowA, - ncolB, - ncolA_ip, - &alpha, - Atmp_device, - LDA, - U_tmp_, - ncolA_ip, - &zero, - B_tmp_, - LDA); + Parallel_Common::recv_dev( + Atmp_device, + size, + ip, + 0, + col_world, + &status, + A_tmp_.data()); + ModuleBase::gemm_op()( + 'N', + 'N', + nrowA, + ncolB, + ncolA_ip, + &alpha, + Atmp_device, + LDA, + U_tmp_, + ncolA_ip, + &zero, + B_tmp_, + LDA); // sum all the results T one = 1.0; - ModuleBase::axpy_op()(ncolB * LDA, &one, B_tmp_, 1, B, 1); + ModuleBase::axpy_op()( + ncolB * LDA, + &one, + B_tmp_, + 1, + B, + 1); } } @@ -155,19 +220,20 @@ void PLinearTransform::act(const T alpha, const T* A, const T* U, con else #endif { - ModuleBase::gemm_op()('N', - 'N', - nrowA, - ncolB, - ncolA, - &alpha, - A, - LDA, - U, - ncolA, - &beta, - B, - LDA); + ModuleBase::gemm_op()( + 'N', + 'N', + nrowA, + ncolB, + ncolA, + &alpha, + A, + LDA, + U, + ncolA, + &beta, + B, + LDA); } ModuleBase::timer::end("PLinearTransform", "act"); }; diff --git a/source/source_hsolver/para_lin_tf.h b/source/source_hsolver/para_lin_tf.h index 3e60d0eba4d..ad3f10918e2 100644 --- a/source/source_hsolver/para_lin_tf.h +++ b/source/source_hsolver/para_lin_tf.h @@ -15,8 +15,7 @@ namespace hsolver * A and B are local matrice * U can be a local matrix or a global matrix */ -template -class PLinearTransform +template class PLinearTransform { public: using syncmem_dev_op = base_device::memory::synchronize_memory_op; @@ -46,14 +45,15 @@ class PLinearTransform * A: LDA * nrow, U_global: ncolA_global * ncolB_global, U_local: ncolA_global * ncolB * B: LDA * ncolB */ - void set_dimension(const int nrowA, - const int ncolA, - const int ncolB, - const int LDA, + void set_dimension( + const int nrowA, + const int ncolA, + const int ncolB, + const int LDA, #ifdef __MPI - MPI_Comm col_world, + MPI_Comm col_world, #endif - const bool localU); + const bool localU); /** * @brief B = alpha * A * U + beta * B @@ -70,15 +70,20 @@ class PLinearTransform * @param B : input/output matrix * */ - void act(const T alpha, const T* A, const T* U_global, const T beta, T* B); + void act( + const T alpha, + const T *A, + const T *U_global, + const T beta, + T *B); #ifdef __MPI private: std::vector A_tmp_; // temperory memory for A std::vector isend_tmp_; // temperory memory for isend - T* U_tmp_ = nullptr; // temperory memory for U - T* B_tmp_ = nullptr; // temperory memory for B - T* A_tmp_device_ = nullptr; // temperory pointer + T *U_tmp_ = nullptr; // temperory memory for U + T *B_tmp_ = nullptr; // temperory memory for B + T *A_tmp_device_ = nullptr; // temperory pointer #endif }; } // namespace hsolver diff --git a/source/source_hsolver/parallel_k2d.cpp b/source/source_hsolver/parallel_k2d.cpp index deb63db058d..b27eae3a1d0 100644 --- a/source/source_hsolver/parallel_k2d.cpp +++ b/source/source_hsolver/parallel_k2d.cpp @@ -1,81 +1,101 @@ #include "parallel_k2d.h" #include "source_base/matrix_block.h" -#include "source_base/parallel_global.h" +#include "source_base/memory_recorder.h" #include "source_base/module_external/scalapack_connector.h" +#include "source_base/parallel_global.h" #include "source_base/timer.h" -#include "source_base/memory_recorder.h" template -void Parallel_K2D::set_para_env(int nks, - const int& nw, - const int& nb2d, - const int& nproc, - const int& my_rank, - const int& nspin) { +void Parallel_K2D::set_para_env( + int nks, + const int &nw, + const int &nb2d, + const int &nproc, + const int &my_rank, + const int &nspin) +{ const int kpar = this->get_kpar(); - Parallel_Global::divide_mpi_groups(nproc, - kpar, - my_rank, - this->NPROC_IN_POOL, - this->MY_POOL, - this->RANK_IN_POOL); + Parallel_Global::divide_mpi_groups( + nproc, + kpar, + my_rank, + this->NPROC_IN_POOL, + this->MY_POOL, + this->RANK_IN_POOL); #ifdef __MPI - MPI_Comm_split(MPI_COMM_WORLD, - this->MY_POOL, - this->RANK_IN_POOL, - &this->POOL_WORLD_K2D); + MPI_Comm_split( + MPI_COMM_WORLD, + this->MY_POOL, + this->RANK_IN_POOL, + &this->POOL_WORLD_K2D); #endif this->Pkpoints.reset(new Parallel_Kpoints()); this->P2D_global.reset(new Parallel_2D()); this->P2D_pool.reset(new Parallel_2D()); - this->Pkpoints - ->kinfo(nks, kpar, this->MY_POOL, this->RANK_IN_POOL, nproc, nspin); - this->P2D_global->init(nw, nw, nb2d, MPI_COMM_WORLD); - this->P2D_pool->init(nw, nw, nb2d, this->POOL_WORLD_K2D); + this->Pkpoints->kinfo( + nks, + kpar, + this->MY_POOL, + this->RANK_IN_POOL, + nproc, + nspin); + this->P2D_global->init( + nw, + nw, + nb2d, + MPI_COMM_WORLD); + this->P2D_pool->init( + nw, + nw, + nb2d, + this->POOL_WORLD_K2D); } template -void Parallel_K2D::distribute_hsk(hsolver::HSMatrix& hs, - const std::vector& ik_kpar, - const int& nw) { +void Parallel_K2D::distribute_hsk(hsolver::HSMatrix &hs, const std::vector &ik_kpar, const int &nw) +{ #ifdef __MPI ModuleBase::timer::start("Parallel_K2D", "distribute_hsk"); for (int ipool = 0; ipool < ik_kpar.size(); ++ipool) { ModuleBase::MatrixBlock HK_global, SK_global; hs.hs_at_k(ik_kpar[ipool], HK_global, SK_global); - if (this->MY_POOL == this->Pkpoints->whichpool[ik_kpar[ipool]]) { + if (this->MY_POOL == this->Pkpoints->whichpool[ik_kpar[ipool]]) + { this->hk_pool.resize(this->P2D_pool->get_local_size(), 0.0); this->sk_pool.resize(this->P2D_pool->get_local_size(), 0.0); } int desc_pool[9]; std::copy(this->P2D_pool->desc, this->P2D_pool->desc + 9, desc_pool); - if (this->MY_POOL != this->Pkpoints->whichpool[ik_kpar[ipool]]) { + if (this->MY_POOL != this->Pkpoints->whichpool[ik_kpar[ipool]]) + { desc_pool[1] = -1; } - Cpxgemr2d(nw, - nw, - HK_global.p, - 1, - 1, - this->P2D_global->desc, - hk_pool.data(), - 1, - 1, - desc_pool, - this->P2D_global->blacs_ctxt); - Cpxgemr2d(nw, - nw, - SK_global.p, - 1, - 1, - this->P2D_global->desc, - sk_pool.data(), - 1, - 1, - desc_pool, - this->P2D_global->blacs_ctxt); + Cpxgemr2d( + nw, + nw, + HK_global.p, + 1, + 1, + this->P2D_global->desc, + hk_pool.data(), + 1, + 1, + desc_pool, + this->P2D_global->blacs_ctxt); + Cpxgemr2d( + nw, + nw, + SK_global.p, + 1, + 1, + this->P2D_global->desc, + sk_pool.data(), + 1, + 1, + desc_pool, + this->P2D_global->blacs_ctxt); } ModuleBase::Memory::record("Parallel_K2D::hsk_pool", this->P2D_pool->get_local_size() * 2 * sizeof(TK)); ModuleBase::timer::end("Parallel_K2D", "distribute_hsk"); @@ -83,19 +103,19 @@ void Parallel_K2D::distribute_hsk(hsolver::HSMatrix& hs, #endif } -template -void Parallel_K2D::unset_para_env() { +template void Parallel_K2D::unset_para_env() +{ this->Pkpoints.reset(); this->P2D_global.reset(); this->P2D_pool.reset(); MPI_Comm_free(&this->POOL_WORLD_K2D); } -template -void Parallel_K2D::set_kpar(int kpar) { - if (kpar < 1) { - ModuleBase::WARNING_QUIT("Parallel_K2D::set_kpar", - "kpar must be greater than 0."); +template void Parallel_K2D::set_kpar(int kpar) +{ + if (kpar < 1) + { + ModuleBase::WARNING_QUIT("Parallel_K2D::set_kpar", "kpar must be greater than 0."); } this->kpar_ = kpar; } diff --git a/source/source_hsolver/parallel_k2d.h b/source/source_hsolver/parallel_k2d.h index 8b205b31003..7675b2a3244 100644 --- a/source/source_hsolver/parallel_k2d.h +++ b/source/source_hsolver/parallel_k2d.h @@ -17,10 +17,10 @@ * It is now designed only to work with 2D Eigenvalue solver parallelism. */ -template -class Parallel_K2D { +template class Parallel_K2D +{ public: - /// private constructor + /// private constructor Parallel_K2D() {} /// private destructor ~Parallel_K2D() {} @@ -29,18 +29,17 @@ class Parallel_K2D { */ /// this function sets the parallel environment for k-points parallelism /// including the glabal and pool 2D parallel distribution - void set_para_env(int nks, - const int& nw, - const int& nb2d, - const int& nproc, - const int& my_rank, - const int& nspin); + void set_para_env( + int nks, + const int &nw, + const int &nb2d, + const int &nproc, + const int &my_rank, + const int &nspin); /// this function distributes the Hk and Sk matrices to hk_pool and sk_pool; /// hs supplies H(k) and S(k) for one k point, this class only redistributes them - void distribute_hsk(hsolver::HSMatrix& hs, - const std::vector& ik_kpar, - const int& nw); + void distribute_hsk(hsolver::HSMatrix &hs, const std::vector &ik_kpar, const int &nw); /// this function unsets the parallel environment for k-points parallelism /// including the glabal and pool 2D parallel distribution @@ -52,11 +51,11 @@ class Parallel_K2D { /// get my pool int get_my_pool() const { return this->MY_POOL; } /// get pKpoints - Parallel_Kpoints* get_pKpoints() const { return this->Pkpoints.get(); } + Parallel_Kpoints *get_pKpoints() const { return this->Pkpoints.get(); } /// get p2D_global - Parallel_2D* get_p2D_global() const { return this->P2D_global.get(); } + Parallel_2D *get_p2D_global() const { return this->P2D_global.get(); } /// get p2D_pool - Parallel_2D* get_p2D_pool() const { return this->P2D_pool.get(); } + Parallel_2D *get_p2D_pool() const { return this->P2D_pool.get(); } /** * the local Hk, Sk matrices in POOL_WORLD_K2D