diff --git a/source/source_esolver/esolver_ks_pw.cpp b/source/source_esolver/esolver_ks_pw.cpp index 6714821d02f..190bbcb4f0f 100644 --- a/source/source_esolver/esolver_ks_pw.cpp +++ b/source/source_esolver/esolver_ks_pw.cpp @@ -224,6 +224,8 @@ void ESolver_KS_PW::hamilt2rho_single(UnitCell& ucell, const int iste hsolver::DiagoIterAssist::need_subspace, PARAM.inp.use_k_continuity); + hsolver_pw_obj.set_diago_precision_mode(parse_precision_mode(PARAM.inp.diago_precision_mode)); + hsolver_pw_obj.solve(static_cast*>(this->p_hamilt), *this->stp.template get_psi_t(), this->pelec, this->pelec->ekb.c, GlobalV::RANK_IN_POOL, GlobalV::NPROC_IN_POOL, skip_charge, ucell.tpiba, ucell.nat); } diff --git a/source/source_esolver/esolver_sdft_pw.cpp b/source/source_esolver/esolver_sdft_pw.cpp index 654f45a19cf..f2d0ad6a861 100644 --- a/source/source_esolver/esolver_sdft_pw.cpp +++ b/source/source_esolver/esolver_sdft_pw.cpp @@ -165,6 +165,8 @@ void ESolver_SDFT_PW::hamilt2rho_single(UnitCell& ucell, int istep, i hsolver::DiagoIterAssist::PW_DIAG_THR, hsolver::DiagoIterAssist::need_subspace); + hsolver_pw_sdft_obj.set_diago_precision_mode(parse_precision_mode(PARAM.inp.diago_precision_mode)); + hsolver_pw_sdft_obj.solve(ucell, static_cast*>(this->p_hamilt), *this->stp.template get_psi_t(), diff --git a/source/source_hsolver/diago_cg.cpp b/source/source_hsolver/diago_cg.cpp index bc2d45e1269..47a73565e46 100644 --- a/source/source_hsolver/diago_cg.cpp +++ b/source/source_hsolver/diago_cg.cpp @@ -40,6 +40,7 @@ DiagoCG::DiagoCG(const std::string& basis_type, pw_diag_thr_ = pw_diag_thr; pw_diag_nmax_ = pw_diag_nmax; nproc_in_pool_ = nproc_in_pool; + precision_mode_ = PrecisionMode::kDouble; this->one_ = new T(static_cast(1.0)); this->zero_ = new T(static_cast(0.0)); this->neg_one_ = new T(static_cast(-1.0)); @@ -578,6 +579,166 @@ bool DiagoCG::test_exit_cond(const int& ntry, const int& notconv) con return f1 && (f2 || f3); } +template +double DiagoCG::diag_mixed_precision(const HPsiFunc& hpsi_func, + const SPsiFunc& spsi_func, + const int ld_psi, + const int nband, + const int dim, + T* psi_in, + Real* eigenvalue_in, + const std::vector& ethr_band, + const Real* prec) +{ +#ifdef ENABLE_MIXED_PRECISION + using MixedT = typename std::conditional::value, + float, + std::complex>::type; + using MixedReal = typename GetTypeReal::type; + + auto psi = ct::TensorMap(psi_in, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nband, ld_psi})); + auto psi_temp = psi.slice({0, 0}, {nband, dim}); + auto psi_mixed = psi_temp.cast(); + + ct::Tensor prec_mixed; + if (prec != nullptr) + { + auto prec_map = ct::TensorMap(const_cast(prec), + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({dim})); + prec_mixed = prec_map.template cast().template to_device(); + } + + std::vector eigen_mixed(nband, static_cast(0.0)); + + auto hpsi_func_mixed = [hpsi_func](MixedT* psi_in_mixed, + MixedT* hpsi_out_mixed, + const int ld_psi_mixed, + const int nvec) { + auto psi_in_map = ct::TensorMap(psi_in_mixed, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + auto psi_in_double = psi_in_map.cast(); + auto hpsi_double = ct::Tensor(ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + hpsi_func(psi_in_double.template data(), hpsi_double.template data(), ld_psi_mixed, nvec); + auto hpsi_mixed_out = hpsi_double.cast(); + ct::TensorMap hpsi_out_tensor(hpsi_out_mixed, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + hpsi_out_tensor.CopyFrom(hpsi_mixed_out); + }; + + auto spsi_func_mixed = [spsi_func](MixedT* psi_in_mixed, + MixedT* spsi_out_mixed, + const int ld_psi_mixed, + const int nvec) { + auto psi_in_map = ct::TensorMap(psi_in_mixed, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + auto psi_in_double = psi_in_map.cast(); + auto spsi_double = ct::Tensor(ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + spsi_func(psi_in_double.template data(), spsi_double.template data(), ld_psi_mixed, nvec); + auto spsi_mixed_out = spsi_double.cast(); + ct::TensorMap spsi_out_tensor(spsi_out_mixed, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + spsi_out_tensor.CopyFrom(spsi_mixed_out); + }; + + auto double_subspace = subspace_func_; + auto subspace_func_mixed = [double_subspace](MixedT* psi_in_mixed, + MixedT* psi_out_mixed, + const int ld_psi_mixed, + const int nband_mixed, + const bool S_orth) { + if (!double_subspace) + { + return; + } + auto psi_in_map = ct::TensorMap(psi_in_mixed, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nband_mixed, ld_psi_mixed})); + auto psi_in_double = psi_in_map.cast(); + auto psi_out_double = ct::Tensor(ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nband_mixed, ld_psi_mixed})); + double_subspace(psi_in_double.template data(), psi_out_double.template data(), ld_psi_mixed, nband_mixed, S_orth); + auto psi_out_mixed_tensor = psi_out_double.cast(); + ct::TensorMap psi_out_tensor(psi_out_mixed, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nband_mixed, ld_psi_mixed})); + psi_out_tensor.CopyFrom(psi_out_mixed_tensor); + }; + + hsolver::DiagoCG mixed_solver( + basis_type_, + calculation_, + need_subspace_, + subspace_func_mixed, + pw_diag_thr_, + pw_diag_nmax_, + nproc_in_pool_); + mixed_solver.set_precision_mode(hsolver::PrecisionMode::kFloat); + + double float_avg_iter = mixed_solver.diag(hpsi_func_mixed, + spsi_func_mixed, + ld_psi, + nband, + dim, + psi_mixed.template data(), + eigen_mixed.data(), + ethr_band, + prec != nullptr ? prec_mixed.template data() : nullptr); + + auto psi_refined = psi_mixed.template cast(); + psi_temp.CopyFrom(psi_refined); + + ct::Tensor eigen = ct::TensorMap(eigenvalue_in, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nband})); + + ct::Tensor prec_tensor; + if (prec != nullptr) + { + prec_tensor = ct::TensorMap(const_cast(prec), + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({dim})) + .template to_device(); + } + + avg_iter_ += float_avg_iter; + this->diag_once(prec_tensor, psi_temp, eigen, ethr_band); + + if (this->notconv_ > std::max(5, this->n_band_ / 4)) + { + std::cout << "\n notconv = " << this->notconv_; + std::cout << "\n DiagoCG::diag_mixed_precision', too many bands are not converged! \n"; + } + + psi.zero(); + psi.sync(psi_temp); + return avg_iter_; +#else + return 0.0; +#endif +} + template double DiagoCG::diag(const HPsiFunc& hpsi_func, const SPsiFunc& spsi_func, @@ -593,6 +754,21 @@ double DiagoCG::diag(const HPsiFunc& hpsi_func, REQUIRES_OK(static_cast(ethr_band.size()) >= nband, "DiagoCG::diag: ethr_band size must be >= nband"); + if (precision_mode_ == PrecisionMode::kMixed) + { +#ifdef ENABLE_MIXED_PRECISION + return diag_mixed_precision(hpsi_func, + spsi_func, + ld_psi, + nband, + dim, + psi_in, + eigenvalue_in, + ethr_band, + prec); +#endif + } + auto psi = ct::TensorMap(psi_in, ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, diff --git a/source/source_hsolver/diago_cg.h b/source/source_hsolver/diago_cg.h index 99d9369a0a3..121654b0b5a 100644 --- a/source/source_hsolver/diago_cg.h +++ b/source/source_hsolver/diago_cg.h @@ -10,6 +10,42 @@ #include #include +#include + +namespace hsolver { + +/** + * @brief Precision mode for diagonalization solvers. + */ +enum class PrecisionMode +{ + kDouble = 0, ///< Pure double precision (default) + kFloat = 1, ///< Pure single precision + kMixed = 2 ///< Mixed precision (float iteration + double refinement) +}; + +} // namespace hsolver + +inline hsolver::PrecisionMode parse_precision_mode(const std::string& mode_str) +{ + if (mode_str == "float" || mode_str == "single") + return hsolver::PrecisionMode::kFloat; + if (mode_str == "mixed" || mode_str == "auto") + return hsolver::PrecisionMode::kMixed; + return hsolver::PrecisionMode::kDouble; +} + +inline std::string precision_mode_to_string(hsolver::PrecisionMode mode) +{ + switch (mode) + { + case hsolver::PrecisionMode::kFloat: return "float"; + case hsolver::PrecisionMode::kMixed: return "mixed"; + case hsolver::PrecisionMode::kDouble: + default: return "double"; + } +} + namespace hsolver { template @@ -25,6 +61,7 @@ class DiagoCG final using HPsiFunc = std::function; using SPsiFunc = std::function; using SubspaceFunc = std::function; + // Constructor need: // 1. temporary mock of Hamiltonian "Hamilt_PW" // 2. precondition pointer should point to place of precondition array. @@ -38,6 +75,8 @@ class DiagoCG final const int& pw_diag_nmax, const int& nproc_in_pool); + void set_precision_mode(const PrecisionMode mode) { precision_mode_ = mode; } + ~DiagoCG(); // virtual void init(){}; @@ -80,6 +119,7 @@ class DiagoCG final std::string calculation_ = {}; bool need_subspace_ = false; + PrecisionMode precision_mode_ = PrecisionMode::kDouble; /// A function object that performs the hPsi calculation. HPsiFunc hpsi_func_ = nullptr; /// A function object that performs the sPsi calculation. @@ -133,6 +173,16 @@ class DiagoCG final ct::Tensor& eigen, const std::vector& ethr_band); + double diag_mixed_precision(const HPsiFunc& hpsi_func, + const SPsiFunc& spsi_func, + const int ld_psi, + const int nband, + const int dim, + T* psi_in, + Real* eigenvalue_in, + const std::vector& ethr_band, + const Real* prec); + bool test_exit_cond(const int& ntry, const int& notconv) const; using dot_real_op = ModuleBase::dot_real_op; diff --git a/source/source_hsolver/diago_david.cpp b/source/source_hsolver/diago_david.cpp index 04e50e76c68..21d0f11a99d 100644 --- a/source/source_hsolver/diago_david.cpp +++ b/source/source_hsolver/diago_david.cpp @@ -8,6 +8,10 @@ #include "source_base/kernels/math_kernel_op.h" #include "source_base/parallel_comm.h" +#include +#include +#include + using namespace hsolver; @@ -1003,6 +1007,144 @@ void DiagoDavid::planSchmidtOrth(const int nband, std::vector& p } +template +int DiagoDavid::diag_mixed_precision(const HPsiFunc& hpsi_func, + const SPsiFunc& spsi_func, + 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) +{ +#ifdef ENABLE_MIXED_PRECISION + // Mixed precision: convert to float, run Davidson, then refine in double + using MixedT = typename std::conditional::value, + float, + std::complex>::type; + using MixedReal = typename GetTypeReal::type; + + // Mixed precision currently only supported on CPU; fallback to double on GPU + if (this->device == base_device::GpuDevice) + { + // Fallback: run standard double-precision diag + return this->diag(hpsi_func, spsi_func, ld_psi, psi_in, eigenvalue_in, + ethr_band, david_maxiter, ntry_max, notconv_max); + } + + // Convert psi to mixed precision + auto psi_tensor = ct::TensorMap(psi_in, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nband, ld_psi})); + auto psi_slice = psi_tensor.slice({0, 0}, {nband, dim}); + auto psi_mixed = psi_slice.cast(); + + // Convert precondition to mixed precision + ct::Tensor prec_mixed; + if (this->precondition != nullptr) + { + auto prec_map = ct::TensorMap(const_cast(this->precondition), + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({dim})); + prec_mixed = prec_map.template cast(); + } + + // Wrap H*psi and S*psi to operate in double but return mixed precision results + auto hpsi_func_mixed = [hpsi_func](MixedT* psi_in_mixed, + MixedT* hpsi_out_mixed, + const int ld_psi_mixed, + const int nvec) { + auto psi_in_map = ct::TensorMap(psi_in_mixed, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + auto psi_in_double = psi_in_map.cast(); + auto hpsi_double = ct::Tensor(ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + hpsi_func(psi_in_double.template data(), hpsi_double.template data(), ld_psi_mixed, nvec); + auto hpsi_mixed_out = hpsi_double.cast(); + ct::TensorMap hpsi_out_tensor(hpsi_out_mixed, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + hpsi_out_tensor.CopyFrom(hpsi_mixed_out); + }; + + auto spsi_func_mixed = [spsi_func](MixedT* psi_in_mixed, + MixedT* spsi_out_mixed, + const int ld_psi_mixed, + const int nvec) { + auto psi_in_map = ct::TensorMap(psi_in_mixed, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + auto psi_in_double = psi_in_map.cast(); + auto spsi_double = ct::Tensor(ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + spsi_func(psi_in_double.template data(), spsi_double.template data(), ld_psi_mixed, nvec); + auto spsi_mixed_out = spsi_double.cast(); + ct::TensorMap spsi_out_tensor(spsi_out_mixed, + ct::DataTypeToEnum::value, + ct::DeviceTypeToEnum::value, + ct::TensorShape({nvec, ld_psi_mixed})); + spsi_out_tensor.CopyFrom(spsi_mixed_out); + }; + + // Allocate mixed precision eigenvalue storage + std::vector eigen_mixed(nband, static_cast(0.0)); + + // Run Davidson in mixed (float) precision + diag_comm_info comm_info_mixed = this->diag_comm; + DiagoDavid david_mixed( + prec_mixed.NumElements() > 0 ? prec_mixed.template data() : nullptr, + nband, dim, david_ndim, comm_info_mixed); + david_mixed.set_precision_mode(PrecisionMode::kFloat); + + int mixed_iter = david_mixed.diag( + hpsi_func_mixed, + spsi_func_mixed, + ld_psi, + psi_mixed.template data(), + eigen_mixed.data(), + ethr_band, + david_maxiter, + ntry_max, + notconv_max); + + // Convert back to double precision + auto psi_refined = psi_mixed.template cast(); + psi_slice.CopyFrom(psi_refined); + + // Copy eigenvalues to output + for (int i = 0; i < nband; ++i) + { + eigenvalue_in[i] = static_cast(eigen_mixed[i]); + } + + // Refinement: run one double-precision Davidson iteration + int refine_iter = this->diag_once(hpsi_func, spsi_func, + dim, nband, ld_psi, + psi_in, eigenvalue_in, + ethr_band, david_maxiter); + + if (this->notconv > std::max(5, nband / 4)) + { + std::cout << "\n notconv = " << this->notconv; + std::cout << "\n DiagoDavid::diag_mixed_precision', too many bands are not converged! \n"; + } + + return mixed_iter + refine_iter; +#else + return 0; +#endif +} + + template int DiagoDavid::diag(const HPsiFunc& hpsi_func, const SPsiFunc& spsi_func, @@ -1014,6 +1156,17 @@ int DiagoDavid::diag(const HPsiFunc& hpsi_func, const int ntry_max, const int notconv_max) { + // Dispatch to mixed precision if requested + if (precision_mode_ == PrecisionMode::kMixed) + { +#ifdef ENABLE_MIXED_PRECISION + return diag_mixed_precision(hpsi_func, spsi_func, + ld_psi, psi_in, eigenvalue_in, + ethr_band, david_maxiter, + ntry_max, notconv_max); +#endif + } + /// record the times of trying iterative diagonalization int ntry = 0; this->notconv = 0; diff --git a/source/source_hsolver/diago_david.h b/source/source_hsolver/diago_david.h index e9ee3a50fde..8522b5b1c99 100644 --- a/source/source_hsolver/diago_david.h +++ b/source/source_hsolver/diago_david.h @@ -9,12 +9,15 @@ #include "source_hsolver/diag_comm_info.h" #include "source_hsolver/kernels/hegvd_op.h" +#include "source_hsolver/diago_cg.h" #include #include +#include namespace hsolver { + /** * @class DiagoDavid * @brief A class that implements the block-Davidson algorithm for solving generalized eigenvalue problems. @@ -58,6 +61,8 @@ class DiagoDavid const int david_ndim_in, const diag_comm_info& diag_comm_in); + void set_precision_mode(const PrecisionMode mode) { precision_mode_ = mode; } + /** * @brief Destructor for the DiagoDavid class. * @@ -139,11 +144,36 @@ class DiagoDavid 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 + /** + * @brief Mixed precision diagonalization using float iteration + double refinement. + * + * Converts wavefunctions to float/complex, performs Davidson iteration + * in single precision, then refines the result in double precision. + * The refinement may perform multiple Davidson iterations up to the + * configured iteration limit, depending on convergence behavior. + * + * @return Total number of iterations (single-precision iterations + + * double-precision refinement iterations). + */ + int diag_mixed_precision( + const HPsiFunc& hpsi_func, + const SPsiFunc& spsi_func, + 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); + private: int test_david = 0; diag_comm_info diag_comm; + /// Precision mode: kDouble (default), kFloat, or kMixed + PrecisionMode precision_mode_ = PrecisionMode::kDouble; + /// number of required eigenpairs const int nband; /// dimension of the input matrix to be diagonalized diff --git a/source/source_hsolver/hsolver_pw.cpp b/source/source_hsolver/hsolver_pw.cpp index b88bc3b90dd..3a363a99546 100644 --- a/source/source_hsolver/hsolver_pw.cpp +++ b/source/source_hsolver/hsolver_pw.cpp @@ -299,6 +299,7 @@ void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, this->diag_thr, this->diag_iter_max, this->nproc_in_pool); + cg.set_precision_mode(this->diago_precision_mode_); DiagoIterAssist::avg_iter += static_cast( cg.diag(hpsi_func, @@ -367,6 +368,7 @@ void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, const int ld_psi = psi.get_nbasis(); /// leading dimension of psi DiagoDavid david(pre_condition.data(), nband, dim, PARAM.inp.pw_diag_ndim, comm_info); + david.set_precision_mode(this->diago_precision_mode_); // do diag and add davidson iteration counts up to avg_iter DiagoIterAssist::avg_iter += static_cast( david.diag(hpsi_func, diff --git a/source/source_hsolver/hsolver_pw.h b/source/source_hsolver/hsolver_pw.h index 7300972cc5c..2594a3c015f 100644 --- a/source/source_hsolver/hsolver_pw.h +++ b/source/source_hsolver/hsolver_pw.h @@ -5,6 +5,7 @@ #include "source_hamilt/hamilt.h" #include "source_base/macros.h" #include "source_basis/module_pw/pw_basis_k.h" +#include "source_hsolver/diago_cg.h" #include namespace hsolver @@ -39,6 +40,13 @@ class HSolverPW diag_iter_max(diag_iter_max_in), diag_thr(diag_thr_in), need_subspace(need_subspace_in), use_k_continuity(use_k_continuity_in) {}; + /// @brief Set the precision mode for diagonalization solvers + /// @param mode PrecisionMode enum value (kDouble, kFloat, or kMixed) + void set_diago_precision_mode(const PrecisionMode mode) { diago_precision_mode_ = mode; } + + /// @brief Get the current precision mode + PrecisionMode get_diago_precision_mode() const { return diago_precision_mode_; } + /// @brief solve function for pw /// @param pHamilt interface to hamilt /// @param psi reference to psi @@ -85,6 +93,9 @@ class HSolverPW const bool use_k_continuity; + /// Precision mode for diagonalization: kDouble (default), kFloat, or kMixed + PrecisionMode diago_precision_mode_ = PrecisionMode::kDouble; + protected: Device* ctx = {}; diff --git a/source/source_hsolver/test/diago_cg_mixed_test.cpp b/source/source_hsolver/test/diago_cg_mixed_test.cpp new file mode 100644 index 00000000000..8a4d9b8251c --- /dev/null +++ b/source/source_hsolver/test/diago_cg_mixed_test.cpp @@ -0,0 +1,129 @@ +#include "gtest/gtest.h" +#include "source_hsolver/diago_cg.h" +#include +#include +#include + +using Complex = std::complex; + +static void make_hermitian(int n, std::vector& H) +{ + H.resize(n * n); + std::mt19937_64 rng(12345); + std::uniform_real_distribution dist(-1.0, 1.0); + for (int i = 0; i < n; ++i) { + for (int j = 0; j <= i; ++j) { + const double real = dist(rng); + const double imag = (i == j ? 0.0 : dist(rng)); + H[i * n + j] = Complex(real, imag); + H[j * n + i] = std::conj(H[i * n + j]); + } + } +} + +static void make_random_psi(int nband, int dim, std::vector& psi) +{ + psi.resize(static_cast(nband) * dim); + std::mt19937_64 rng(54321); + std::uniform_real_distribution dist(-0.5, 0.5); + for (int i = 0; i < nband * dim; ++i) { + psi[i] = Complex(dist(rng), dist(rng)); + } +} + +static void apply_hamiltonian(const std::vector& H, + int n, + const Complex* psi_in, + Complex* hpsi_out, + int ld, int nvec) +{ + for (int v = 0; v < nvec; ++v) { + const Complex* psi_vec = psi_in + static_cast(v) * ld; + Complex* out_vec = hpsi_out + static_cast(v) * ld; + for (int i = 0; i < n; ++i) { + Complex sum = 0.0; + for (int j = 0; j < n; ++j) { + sum += H[static_cast(i) * n + j] * psi_vec[j]; + } + out_vec[i] = sum; + } + } +} + +static void apply_overlap(const Complex* psi_in, + Complex* spsi_out, + int ld, + int nvec) +{ + for (int i = 0; i < nvec * ld; ++i) { + spsi_out[i] = psi_in[i]; + } +} + +TEST(DiagoCGMixedTest, MixedPrecisionMatchesDouble) +{ + const int dim = 8; + const int nband = 3; + const int ld_psi = dim; + + std::vector H; + make_hermitian(dim, H); + + std::vector psi_initial; + make_random_psi(nband, dim, psi_initial); + + std::vector psi_double = psi_initial; + std::vector psi_mixed = psi_initial; + std::vector eigen_double(nband, 0.0); + std::vector eigen_mixed(nband, 0.0); + + auto hpsi_func = [&H, dim](Complex* psi_in, Complex* hpsi_out, const int ld, const int nvec) { + apply_hamiltonian(H, dim, psi_in, hpsi_out, ld, nvec); + }; + auto spsi_func = [](Complex* psi_in, Complex* spsi_out, const int ld, const int nvec) { + apply_overlap(psi_in, spsi_out, ld, nvec); + }; + + std::vector ethr_band(nband, 1e-6); + + hsolver::DiagoCG cg_double( + "pw", + "nscf", + false, + hsolver::DiagoCG::SubspaceFunc(), + 1e-6, + 200, + 1); + cg_double.diag(hpsi_func, + spsi_func, + ld_psi, + nband, + dim, + psi_double.data(), + eigen_double.data(), + ethr_band, + nullptr); + + hsolver::DiagoCG cg_mixed( + "pw", + "nscf", + false, + hsolver::DiagoCG::SubspaceFunc(), + 1e-6, + 200, + 1, + hsolver::PrecisionMode::kMixed); + cg_mixed.diag(hpsi_func, + spsi_func, + ld_psi, + nband, + dim, + psi_mixed.data(), + eigen_mixed.data(), + ethr_band, + nullptr); + + for (int i = 0; i < nband; ++i) { + EXPECT_NEAR(eigen_double[i], eigen_mixed[i], 1e-6) << "Index=" << i; + } +} diff --git a/source/source_hsolver/test/hsolver_pw_sup.h b/source/source_hsolver/test/hsolver_pw_sup.h index b41196c396d..8446b065dfb 100644 --- a/source/source_hsolver/test/hsolver_pw_sup.h +++ b/source/source_hsolver/test/hsolver_pw_sup.h @@ -79,6 +79,7 @@ DiagoCG::DiagoCG(const std::string& basis_type, pw_diag_thr_ = pw_diag_thr; pw_diag_nmax_ = pw_diag_nmax; nproc_in_pool_ = nproc_in_pool; + precision_mode_ = PrecisionMode::kDouble; this->one_ = new T(static_cast(1.0)); this->zero_ = new T(static_cast(0.0)); this->neg_one_ = new T(static_cast(-1.0)); diff --git a/source/source_io/module_parameter/input_parameter.h b/source/source_io/module_parameter/input_parameter.h index 029ad364eb5..e1e56acea0c 100644 --- a/source/source_io/module_parameter/input_parameter.h +++ b/source/source_io/module_parameter/input_parameter.h @@ -91,6 +91,7 @@ struct Input_para bool diago_smooth_ethr = false; ///< smooth ethr for iter methods int pw_diag_ndim = 4; ///< dimension of workspace for Davidson diagonalization int diago_cg_prec = 1; ///< mohan add 2012-03-31 + std::string diago_precision_mode = "double"; ///< precision mode for diagonalization: double, float, or mixed int diag_subspace = 0; // 0: Lapack, 1: elpa, 2: scalapack bool use_k_continuity = false; ///< whether to use k-point continuity for initializing wave functions