diff --git a/python/pyabacus/src/hsolver/diago_adapter.hpp b/python/pyabacus/src/hsolver/diago_adapter.hpp index 2d6e652ba46..3ce2c89d561 100644 --- a/python/pyabacus/src/hsolver/diago_adapter.hpp +++ b/python/pyabacus/src/hsolver/diago_adapter.hpp @@ -85,7 +85,7 @@ class PyDiagoDavidAdapter { const PyHSOperator op(mm_op); - solver_ = std::make_unique(precond_vec.data(), nband_, nbasis_, dav_ndim, comm_info); + solver_ = std::make_unique(precond_vec.data(), nband_, nbasis_, dav_ndim, false, comm_info); return solver_->diag(op, nbasis_, storage_.psi_ptr(), storage_.eigenvalue_ptr(), diag_ethr, max_iter); } diff --git a/source/source_base/kernels/cuda/math_kernel_op.cu b/source/source_base/kernels/cuda/math_kernel_op.cu index 0f2d51f4394..fa6419fea26 100644 --- a/source/source_base/kernels/cuda/math_kernel_op.cu +++ b/source/source_base/kernels/cuda/math_kernel_op.cu @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -293,7 +294,22 @@ void gemm_op, base_device::DEVICE_GPU>::operator()(const cha { cublasOperation_t cutransA = judge_trans_op(true, transa, "gemm_op"); cublasOperation_t cutransB = judge_trans_op(true, transb, "gemm_op"); - CHECK_CUBLAS(cublasCgemm(cublas_handle, cutransA, cutransB, m, n ,k, (float2*)alpha, (float2*)a , lda, (float2*)b, ldb, (float2*)beta, (float2*)c, ldc)); + // Bound rounding error in long single-precision conjugate inner products. + constexpr int block = 2048; + if (cutransA == CUBLAS_OP_C && cutransB == CUBLAS_OP_N && k > block) + { + const std::complex one(1.0f, 0.0f); + for (int offset = 0; offset < k; offset += block) + { + CHECK_CUBLAS(cublasCgemm(cublas_handle, cutransA, cutransB, m, n, std::min(block, k - offset), + (float2*)alpha, (float2*)(a + offset), lda, (float2*)(b + offset), ldb, + (float2*)(offset == 0 ? beta : &one), (float2*)c, ldc)); + } + } + else + { + CHECK_CUBLAS(cublasCgemm(cublas_handle, cutransA, cutransB, m, n ,k, (float2*)alpha, (float2*)a , lda, (float2*)b, ldb, (float2*)beta, (float2*)c, ldc)); + } } template <> diff --git a/source/source_base/kernels/test/math_kernel_test.cpp b/source/source_base/kernels/test/math_kernel_test.cpp index 709f9b35cc7..0d8abc1182c 100644 --- a/source/source_base/kernels/test/math_kernel_test.cpp +++ b/source/source_base/kernels/test/math_kernel_test.cpp @@ -907,6 +907,120 @@ TEST_F(TestModuleHsolverMathKernel, matrixCopy_op_gpu) delete_memory_op()(device_B); } + +#if __UT_USE_CUDA +TEST_F(TestModuleHsolverMathKernel, gemm_op_cgemm_blocked_gpu) +{ + using Cf = std::complex; + using Zd = std::complex; + + struct Case + { + int m; + int n; + int k; + int lda; + int ldb; + int ldc; + Cf alpha; + Cf beta; + }; + + const std::vector cases = { + // k > 2048, one full block plus a short tail; padded leading dimensions. + {3, 2, 2050, 2057, 2053, 8, Cf(1.3f, -0.2f), Cf(0.5f, 0.25f)}, + // Two full 2048 blocks plus a 5-column tail. + {4, 3, 4101, 4106, 4103, 7, Cf(-0.7f, 0.4f), Cf(-0.3f, 0.7f)}, + }; + + ModuleBase::createGpuBlasHandle(); + for (const Case& c : cases) + { + std::vector A(static_cast(c.lda) * c.m, Cf(0.0f, 0.0f)); + std::vector B(static_cast(c.ldb) * c.n, Cf(0.0f, 0.0f)); + std::vector C(static_cast(c.ldc) * c.n, Cf(0.0f, 0.0f)); + std::vector C0(static_cast(c.ldc) * c.n, Cf(0.0f, 0.0f)); + + for (int i = 0; i < c.m; ++i) + { + for (int l = 0; l < c.k; ++l) + { + A[l + static_cast(i) * c.lda] + = Cf(0.001f * (l % 37) - 0.02f, 0.0005f * (l % 23) - 0.01f); + } + } + for (int j = 0; j < c.n; ++j) + { + for (int l = 0; l < c.k; ++l) + { + B[l + static_cast(j) * c.ldb] + = Cf(0.002f * (l % 29) - 0.03f, 0.0007f * (l % 19) + 0.01f); + } + } + for (int j = 0; j < c.n; ++j) + { + for (int i = 0; i < c.m; ++i) + { + C0[i + static_cast(j) * c.ldc] = Cf(0.1f * i - 0.2f, 0.05f * j + 0.1f); + C[i + static_cast(j) * c.ldc] = C0[i + static_cast(j) * c.ldc]; + } + } + + std::vector ref(static_cast(c.ldc) * c.n, Zd(0.0, 0.0)); + for (int j = 0; j < c.n; ++j) + { + for (int i = 0; i < c.m; ++i) + { + Zd sum(0.0, 0.0); + for (int l = 0; l < c.k; ++l) + { + sum += std::conj(Zd(A[l + static_cast(i) * c.lda])) + * Zd(B[l + static_cast(j) * c.ldb]); + } + ref[i + static_cast(j) * c.ldc] + = Zd(c.alpha) * sum + Zd(c.beta) * Zd(C0[i + static_cast(j) * c.ldc]); + } + } + + Cf* dA = nullptr; + Cf* dB = nullptr; + Cf* dC = nullptr; + base_device::memory::resize_memory_op()(dA, A.size()); + base_device::memory::resize_memory_op()(dB, B.size()); + base_device::memory::resize_memory_op()(dC, C.size()); + base_device::memory::synchronize_memory_op()( + dA, A.data(), A.size()); + base_device::memory::synchronize_memory_op()( + dB, B.data(), B.size()); + base_device::memory::synchronize_memory_op()( + dC, C.data(), C.size()); + + ModuleBase::gemm_op()( + 'C', 'N', c.m, c.n, c.k, &c.alpha, dA, c.lda, dB, c.ldb, &c.beta, dC, c.ldc); + + base_device::memory::synchronize_memory_op()( + C.data(), dC, C.size()); + + for (int j = 0; j < c.n; ++j) + { + for (int i = 0; i < c.m; ++i) + { + const Zd diff = Zd(C[i + static_cast(j) * c.ldc]) + - ref[i + static_cast(j) * c.ldc]; + SCOPED_TRACE(::testing::Message() << "m=" << c.m << " n=" << c.n << " k=" << c.k + << " i=" << i << " j=" << j); + EXPECT_LT(std::abs(diff), 1e-3); + } + } + + base_device::memory::delete_memory_op()(dA); + base_device::memory::delete_memory_op()(dB); + base_device::memory::delete_memory_op()(dC); + } + ModuleBase::destoryBLAShandle(); +} +#endif // __UT_USE_CUDA + #endif // __UT_USE_CUDA || __UT_USE_ROCM /* diff --git a/source/source_base/module_container/ATen/kernels/cuda/lapack.cu b/source/source_base/module_container/ATen/kernels/cuda/lapack.cu index 4c69973b6be..2412bfd851e 100644 --- a/source/source_base/module_container/ATen/kernels/cuda/lapack.cu +++ b/source/source_base/module_container/ATen/kernels/cuda/lapack.cu @@ -7,6 +7,7 @@ #include #include +#include #include "source_base/module_device/device_check.h" @@ -129,7 +130,8 @@ struct lapack_geqrf_inplace { const int m, const int n, T *d_A, - const int lda) + const int lda, + T* rdiag) { const int k = std::min(m, n); @@ -139,6 +141,18 @@ struct lapack_geqrf_inplace { cuSolverConnector::geqrf(cusolver_handle, m, n, d_A, lda, d_tau); + // Capture the diagonal of R before orgqr overwrites d_A with Q. + if (rdiag != nullptr) + { + for (int j = 0; j < k; ++j) + { + CHECK_CUDA(cudaMemcpy(rdiag + j, + d_A + static_cast(j) * lda + j, + sizeof(T), + cudaMemcpyDeviceToHost)); + } + } + cuSolverConnector::orgqr(cusolver_handle, m, n, k, d_A, lda, d_tau); CHECK_CUDA(cudaFree(d_tau)); diff --git a/source/source_base/module_container/ATen/kernels/lapack.cpp b/source/source_base/module_container/ATen/kernels/lapack.cpp index 2ab02f35c81..984bb75e8ec 100644 --- a/source/source_base/module_container/ATen/kernels/lapack.cpp +++ b/source/source_base/module_container/ATen/kernels/lapack.cpp @@ -116,7 +116,8 @@ struct lapack_geqrf_inplace { const int m, const int n, T *A, - const int lda) + const int lda, + T* rdiag) { // Tensor or vector? // 1. tau for storing the Householder reflectors @@ -147,6 +148,15 @@ struct lapack_geqrf_inplace { throw std::runtime_error("geqrf failed with info = " + std::to_string(info)); } + // Capture the diagonal of R before orgqr overwrites A with Q. + if (rdiag != nullptr) + { + for (int j = 0; j < k; ++j) + { + rdiag[j] = A[j * lda + j]; + } + } + // 4. use orgqr to compute Q // workspace query lwork = -1; diff --git a/source/source_base/module_container/ATen/kernels/lapack.h b/source/source_base/module_container/ATen/kernels/lapack.h index 117f8ef24ba..0fe8a049430 100644 --- a/source/source_base/module_container/ATen/kernels/lapack.h +++ b/source/source_base/module_container/ATen/kernels/lapack.h @@ -80,12 +80,16 @@ struct lapack_geqrf_inplace { * @param n The number of columns in the matrix A. n >= 0 * @param A Pointer to the matrix A to be factorized. On exit, contains the QR factorization * @param lda The leading dimension of the matrix A. lda >= max(1, m) + * @param rdiag Optional host buffer receiving the diagonal of R before A is + * overwritten with Q. Must hold at least min(m, n) entries, or + * be nullptr when the rank is not needed. */ void operator()( const int m, const int n, T *A, - const int lda); + const int lda, + T* rdiag); }; // This is QR factorization diff --git a/source/source_base/module_container/ATen/kernels/test/lapack_test.cpp b/source/source_base/module_container/ATen/kernels/test/lapack_test.cpp index a318bd8c6b8..92a6054c8fa 100644 --- a/source/source_base/module_container/ATen/kernels/test/lapack_test.cpp +++ b/source/source_base/module_container/ATen/kernels/test/lapack_test.cpp @@ -1,5 +1,7 @@ #include +#include + #include #include #include @@ -134,9 +136,10 @@ TYPED_TEST(LapackTest, GeqrfInPlace) { }).to_device()); Tensor A = A_input; // will be overwritten as Q + std::vector rdiag(n); - // do geqrf -> get orthogonal Q - geqrfCalculator(m, n, A.data(), lda); + // do geqrf -> get orthogonal Q and the diagonal of R + geqrfCalculator(m, n, A.data(), lda, rdiag.data()); // check on CPU Tensor Q = A.to_device(); @@ -190,6 +193,26 @@ TYPED_TEST(LapackTest, GeqrfInPlace) { << "Q^T * Q not identity at (" << i << "," << j << ")"; } } + + // R = Q^H * A_original; its diagonal must match the values captured + // before geqrf overwrote A with Q. + Tensor R = std::move(Tensor({ + static_cast(0.0), static_cast(0.0), static_cast(0.0), + static_cast(0.0), static_cast(0.0), static_cast(0.0), + static_cast(0.0), static_cast(0.0), static_cast(0.0) + }).to_device()); + Tensor A_cpu = A_input.to_device(); + gemm('C', 'N', + n, n, m, + &alpha, + Q_data, lda, + A_cpu.data(), lda, + &beta, + R.data(), n); + for (int j = 0; j < n; ++j) { + EXPECT_NEAR(std::abs(R.data()[j + j * n]), std::abs(rdiag[j]), 1e-5) + << "R diagonal mismatch at " << j; + } } // Test for lapack_heevd and lapack_heevx: diff --git a/source/source_estate/elecstate_pw.cpp b/source/source_estate/elecstate_pw.cpp index 3037b88b286..5b58dc89823 100644 --- a/source/source_estate/elecstate_pw.cpp +++ b/source/source_estate/elecstate_pw.cpp @@ -1,5 +1,8 @@ #include "elecstate_pw.h" +#include +#include + #include "source_base/constants.h" #include "source_base/libm/libm.h" #include "source_base/module_device/device.h" @@ -139,13 +142,35 @@ void ElecStatePW::psiToRho(const psi::Psi& psi) } } + // Keep the FFTs in single precision, but sum k-point densities in double. + const bool sum_k_in_double = std::is_same::value + && std::is_same::value; + std::vector buffer(sum_k_in_double ? this->charge->nrxx : 0); + const auto accumulate_k = [&](Real** source, double** target, const bool first) + { + for (int is = 0; is < this->charge->nspin; ++is) + { + castmem_var_d2h_op()(buffer.data(), source[is], this->charge->nrxx); + for (int ir = 0; ir < this->charge->nrxx; ++ir) + target[is][ir] = first ? buffer[ir] : target[is][ir] + buffer[ir]; + setmem_var_op()(source[is], 0, this->charge->nrxx); + } + }; + for (int ik = 0; ik < psi.get_nk(); ++ik) { psi.fix_k(ik); this->updateRhoK(psi); + if (sum_k_in_double) + { + accumulate_k(this->rho, this->charge->rho, ik == 0); + if (XC_Functional::get_func_type() == 3) + accumulate_k(this->kin_r, this->charge->kin_r, ik == 0); + } } - if (PARAM.inp.device == "gpu" || PARAM.inp.precision == "single") + if ((!sum_k_in_double || psi.get_nk() == 0) + && (PARAM.inp.device == "gpu" || PARAM.inp.precision == "single")) { for (int ii = 0; ii < PARAM.inp.nspin; ii++) { diff --git a/source/source_hamilt/module_xc/xc_functional.h b/source/source_hamilt/module_xc/xc_functional.h index a8b57f00465..882edab71ed 100644 --- a/source/source_hamilt/module_xc/xc_functional.h +++ b/source/source_hamilt/module_xc/xc_functional.h @@ -324,6 +324,18 @@ class XC_Functional const T* rhog, T* grad); + // Overload with precomputed precision-matched gcar/kvec_c device buffers. + template ::type> + static void grad_wfc( + const int ik, + const Real tpiba, + const ModulePW::PW_Basis_K* wfc_basis, + const Real* gcar, + const Real* kvec_c, + const T* rhog, + T* grad); + static void grad_rho( const std::complex* rhog, ModuleBase::Vector3* gdr, diff --git a/source/source_hamilt/module_xc/xc_grad_wfc.cpp b/source/source_hamilt/module_xc/xc_grad_wfc.cpp index 85dbbcdad92..e978eb82124 100644 --- a/source/source_hamilt/module_xc/xc_grad_wfc.cpp +++ b/source/source_hamilt/module_xc/xc_grad_wfc.cpp @@ -15,16 +15,56 @@ void XC_Functional::grad_wfc( const ModulePW::PW_Basis_K* wfc_basis, const T* rhog, T* grad) +{ + using ct_Device = typename ct::PsiToContainer::type; + // gcar/kvec_c are stored as Vector3 on the host. Build + // precision-matched host buffers and upload them, so single precision + // never reinterprets a double tensor and does not depend on the + // device-only PW_Basis_K caches being populated in the stress path. + const int gcar_size = wfc_basis->nks * wfc_basis->npwk_max; + std::vector gcar_host(3 * gcar_size); + for (int i = 0; i < gcar_size; ++i) + { + gcar_host[3 * i] = static_cast(wfc_basis->gcar[i].x); + gcar_host[3 * i + 1] = static_cast(wfc_basis->gcar[i].y); + gcar_host[3 * i + 2] = static_cast(wfc_basis->gcar[i].z); + } + std::vector kvec_host(3 * wfc_basis->nks); + for (int i = 0; i < wfc_basis->nks; ++i) + { + kvec_host[3 * i] = static_cast(wfc_basis->kvec_c[i].x); + kvec_host[3 * i + 1] = static_cast(wfc_basis->kvec_c[i].y); + kvec_host[3 * i + 2] = static_cast(wfc_basis->kvec_c[i].z); + } + auto gcar = ct::TensorMap(gcar_host.data(), + ct::DataTypeToEnum::value, + ct::DeviceType::CpuDevice, + {gcar_size, 3}) + .to_device(); + auto kvec_c = ct::TensorMap(kvec_host.data(), + ct::DataTypeToEnum::value, + ct::DeviceType::CpuDevice, + {wfc_basis->nks, 3}) + .to_device(); + + grad_wfc(ik, tpiba, wfc_basis, gcar.template data(), kvec_c.template data(), rhog, grad); +} + +template +void XC_Functional::grad_wfc( + const int ik, + const Real tpiba, + const ModulePW::PW_Basis_K* wfc_basis, + const Real* gcar, + const Real* kvec_c, + const T* rhog, + T* grad) { using ct_Device = typename ct::PsiToContainer::type; const int npw_k = wfc_basis->npwk[ik]; auto porter = std::move(ct::Tensor( ct::DataTypeToEnum::value, ct::DeviceTypeToEnum::value, {wfc_basis->nmaxgr})); - auto gcar = ct::TensorMap( - &wfc_basis->gcar[0][0], ct::DataType::DT_DOUBLE, ct::DeviceType::CpuDevice, {wfc_basis->nks * wfc_basis->npwk_max, 3}).to_device(); - auto kvec_c = ct::TensorMap( - &wfc_basis->kvec_c[0][0],ct::DataType::DT_DOUBLE, ct::DeviceType::CpuDevice, {wfc_basis->nks, 3}).to_device(); auto xc_functional_grad_wfc_solver = hamilt::xc_functional_grad_wfc_op(); @@ -34,8 +74,8 @@ void XC_Functional::grad_wfc( xc_functional_grad_wfc_solver( ik, ipol, npw_k, wfc_basis->npwk_max, tpiba, - gcar.template data(), - kvec_c.template data(), + gcar, + kvec_c, rhog, porter.data()); // bring the gdr from G --> R @@ -54,6 +94,14 @@ template void XC_Functional::grad_wfc, base_device::DEVICE_ const ModulePW::PW_Basis_K* wfc_basis, const std::complex* rhog, std::complex* grad); +template void XC_Functional::grad_wfc, base_device::DEVICE_CPU, double>( + const int ik, + const double tpiba, + const ModulePW::PW_Basis_K* wfc_basis, + const double* gcar, + const double* kvec_c, + const std::complex* rhog, + std::complex* grad); #if __CUDA || __ROCM template void XC_Functional::grad_wfc, base_device::DEVICE_GPU, double>( const int ik, @@ -61,4 +109,26 @@ template void XC_Functional::grad_wfc, base_device::DEVICE_ const ModulePW::PW_Basis_K* wfc_basis, const std::complex* rhog, std::complex* grad); +template void XC_Functional::grad_wfc, base_device::DEVICE_GPU, double>( + const int ik, + const double tpiba, + const ModulePW::PW_Basis_K* wfc_basis, + const double* gcar, + const double* kvec_c, + const std::complex* rhog, + std::complex* grad); +template void XC_Functional::grad_wfc, base_device::DEVICE_GPU, float>( + const int ik, + const float tpiba, + const ModulePW::PW_Basis_K* wfc_basis, + const std::complex* rhog, + std::complex* grad); +template void XC_Functional::grad_wfc, base_device::DEVICE_GPU, float>( + const int ik, + const float tpiba, + const ModulePW::PW_Basis_K* wfc_basis, + const float* gcar, + const float* kvec_c, + const std::complex* rhog, + std::complex* grad); #endif diff --git a/source/source_hsolver/diago_dav_subspace.cpp b/source/source_hsolver/diago_dav_subspace.cpp index 9b4764fc9d5..ce08b3b6524 100644 --- a/source/source_hsolver/diago_dav_subspace.cpp +++ b/source/source_hsolver/diago_dav_subspace.cpp @@ -1,4 +1,5 @@ #include +#include #include "diago_dav_subspace.h" #include "diago_iter_assist.h" @@ -442,18 +443,36 @@ void Diago_DavSubspace::cal_grad(const HSOperator& op, Real* psi_norm_host = nullptr; resmem_real_h_op()(psi_norm_host, notconv); syncmem_var_d2h_op()(psi_norm_host, psi_norm, notconv); + std::vector valid_indices; for (int i = 0; i < notconv; i++) { - if (psi_norm_host[i] <= 1.0e-12) + const bool finite = std::isfinite(psi_norm_host[i]); + if (!finite || psi_norm_host[i] <= 1.0e-12) { - std::cout << "Diago_DavSubspace::cal_grad: psi_norm <= 0 for band " << i << std::endl; - std::cout << "This may be due to npwx < nbands: the number of plane waves is less than" << std::endl; - std::cout << "the number of bands, leading to a rank-deficient problem." << std::endl; - std::cout << "Please increase ecutwfc or reduce nbands." << std::endl; - delmem_real_h_op()(psi_norm_host); - delmem_real_op()(psi_norm); - ModuleBase::WARNING_QUIT("cal_grad", "psi_norm <= 0"); + std::cout << "Diago_DavSubspace::cal_grad: norm of new direction for band " << i + << (finite ? " is too small" : " is not finite") + << ", treating as converged." << std::endl; + continue; } + valid_indices.push_back(i); + } + if (valid_indices.empty()) + { + this->notconv = 0; + } + else if (static_cast(valid_indices.size()) < notconv) + { + for (int k = 0; k < static_cast(valid_indices.size()); ++k) + { + const int src = valid_indices[k]; + if (src != k) + { + syncmem_complex_op()(psi_iter + (nbase + k) * this->dim, + psi_iter + (nbase + src) * this->dim, + this->dim); + } + } + this->notconv = static_cast(valid_indices.size()); } delmem_real_h_op()(psi_norm_host); delmem_real_op()(psi_norm); @@ -472,25 +491,59 @@ void Diago_DavSubspace::cal_grad(const HSOperator& op, psi_norm); // Check for zero norms (CPU path) + std::vector valid_indices; for (int i = 0; i < notconv; i++) { + const bool finite = std::isfinite(psi_norm[i]); + if (!finite) + { + ModuleBase::WARNING("Diago_DavSubspace::cal_grad", + "Non-finite norm for new direction of band " + + std::to_string(i) + + "; skipping the direction instead of " + "reporting a converged eigenpair."); + continue; + } if (psi_norm[i] <= 1.0e-12) { - std::cout << "Diago_DavSubspace::cal_grad: psi_norm <= 0 for band " << i << std::endl; - std::cout << "This may be due to npwx < nbands: the number of plane waves is less than" << std::endl; - std::cout << "the number of bands, leading to a rank-deficient problem." << std::endl; - std::cout << "Please increase ecutwfc or reduce nbands." << std::endl; - delmem_real_h_op()(psi_norm); - ModuleBase::WARNING_QUIT("cal_grad", "psi_norm <= 0"); + ModuleBase::WARNING("Diago_DavSubspace::cal_grad", + "Near-zero norm for new direction of band " + + std::to_string(i) + + "; treating the direction as converged."); + continue; } + valid_indices.push_back(i); + } + if (valid_indices.empty()) + { + this->notconv = 0; + } + else if (static_cast(valid_indices.size()) < notconv) + { + for (int k = 0; k < static_cast(valid_indices.size()); ++k) + { + const int src = valid_indices[k]; + if (src != k) + { + syncmem_complex_op()(psi_iter + (nbase + k) * this->dim, + psi_iter + (nbase + src) * this->dim, + this->dim); + } + } + this->notconv = static_cast(valid_indices.size()); } delmem_real_h_op()(psi_norm); } // 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); + // An all-rejected block leaves no new direction; asking the operator for + // zero vectors would build an invalid [0, -1] range. + if (this->notconv > 0) + { + op.hpsi(psi_iter + nbase * dim, hpsi + nbase * this->dim, this->dim, this->notconv); + op.spsi(psi_iter + nbase * dim, spsi + nbase * this->dim, this->dim, this->notconv); + } ModuleBase::timer::end("Diago_DavSubspace", "cal_grad"); return; diff --git a/source/source_hsolver/diago_david.cpp b/source/source_hsolver/diago_david.cpp index 31e6fcfb0c2..b8c3e6c6d63 100644 --- a/source/source_hsolver/diago_david.cpp +++ b/source/source_hsolver/diago_david.cpp @@ -1,9 +1,16 @@ #include #include "diago_david.h" +#include +#include +#include +#include +#include #include "source_base/parallel_cell.h" +#include "source_base/parallel_device.h" #include "source_base/timer.h" // ModuleBase::timer::tick -#include "source_base/tool_title.h" // ModuleBase::TITLE +#include "source_base/tool_title.h" +#include "source_base/tool_quit.h" // ModuleBase::TITLE #include "source_base/module_device/device.h" #include "source_hsolver/kernels/hegvd_op.h" @@ -13,14 +20,32 @@ using namespace hsolver; +namespace +{ +void warn_skip_degenerate_direction(const int band) +{ + static std::atomic warned(false); + bool expected = false; + if (warned.compare_exchange_strong(expected, true)) + { + ModuleBase::WARNING("DiagoDavid::SchmidtOrth", + "Skipping a degenerate residual direction (band " + + std::to_string(band) + + ") with non-positive norm; treating it as converged."); + } +} +} // namespace + template DiagoDavid::DiagoDavid(const Real* precondition_in, const int nband_in, const int dim_in, const int david_ndim_in, + const bool use_uspp_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) + : nband(nband_in), dim(dim_in), nbase_x(david_ndim_in * nband_in), david_ndim(david_ndim_in), + use_uspp(use_uspp_in), diag_comm(diag_comm_in) { this->device = base_device::get_device_type(this->ctx); this->precondition = precondition_in; @@ -159,6 +184,51 @@ int DiagoDavid::diag_once(const HSOperator& op, std::vector pre_matrix_mv_m(nband, 1); this->planSchmidtOrth(nband, pre_matrix_mm_m, pre_matrix_mv_m); +#if defined(__CUDA) + const bool block_qr = this->use_block_qr(); + if (block_qr) + { + for (int m = 0; m < nband; ++m) + syncmem_complex_op()(basis + dim*m, psi_in + m*ld_psi, dim); + // Form the basis without subtracting nearly equal norm estimates. + std::vector rdiag(nband, T(0)); + ct::kernels::lapack_geqrf_inplace()(dim, nband, basis, dim, rdiag.data()); + + const int rdiag_size = std::min(dim, nband); + bool full_rank = rdiag_size == nband; + double max_diag = 0.0; + for (int m = 0; m < rdiag_size; ++m) + { + const double abs_diag = std::abs(rdiag[m]); + if (!std::isfinite(abs_diag)) + { + full_rank = false; + } + if (abs_diag > max_diag) + { + max_diag = abs_diag; + } + } + const double rank_tol = 1e-6 * max_diag; + for (int m = 0; m < rdiag_size && full_rank; ++m) + { + if (std::abs(rdiag[m]) <= rank_tol) + { + full_rank = false; + } + } + if (!full_rank) + { + ModuleBase::WARNING_QUIT("DiagoDavid", + "initial trial wavefunction is linearly dependent; " + "cannot construct an orthonormal Davidson basis."); + } + + op.spsi(basis, spsi, dim, nband); + } + else +#endif + { for (int m = 0; m < nband; m++) { { @@ -171,19 +241,26 @@ int DiagoDavid::diag_once(const HSOperator& op, { 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]); + if (!this->SchmidtOrth(dim, + nband, + m, + this->spsi, + &this->lagrange_matrix[m * nband], + pre_matrix_mm_m[m], + pre_matrix_mv_m[m])) + { + ModuleBase::WARNING_QUIT("DiagoDavid", + "initial trial wavefunction is linearly dependent; " + "cannot construct an orthonormal Davidson basis."); + } { // phm_in->sPsi(basis + dim*m, &this->spsi[m * dim], dim, dim, 1); op.spsi(basis + dim*m, &this->spsi[m * dim], dim, 1); } } + } + // end of SchmidtOrth and calculate H|psi> // hpsi_info dav_hpsi_in(&basis, psi::Range(true, 0, 0, nband - 1), this->hpsi); // phm_in->ops->hPsi(dav_hpsi_in); @@ -552,27 +629,114 @@ void DiagoDavid::cal_grad(const HSOperator& op, ); } - for (int m = 0; m < notconv; m++) +#if defined(__CUDA) + bool qr_used = false; + const bool block_qr = this->use_block_qr() && notconv > 0; + if (block_qr) { - this->SchmidtOrth(dim, - nbase + notconv, - nbase + m, - spsi, - &lagrange[m * (nbase + notconv)], - pre_matrix_mm_m[m], - pre_matrix_mv_m[m]); + // Keep a copy of the unprojected block so a rank-deficient QR can + // restore it and fall back to the regular Schmidt path. + T* qr_backup = nullptr; + resmem_complex_op()(qr_backup, dim * notconv); + syncmem_complex_op()(qr_backup, basis + dim * nbase, dim * notconv); + + // Project the new block once and orthogonalize its columns together. + ModuleBase::gemm_op()('N', 'N', dim, notconv, nbase, + this->neg_one, basis, dim, lagrange, nbase + notconv, + this->one, basis + dim*nbase, dim); + + std::vector rdiag(notconv, T(0)); + ct::kernels::lapack_geqrf_inplace()( + dim, notconv, basis + dim * nbase, dim, rdiag.data()); + + const int rdiag_size = std::min(dim, notconv); + bool full_rank = rdiag_size == notconv; + double max_diag = 0.0; + for (int j = 0; j < rdiag_size; ++j) { - // 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); + const double abs_diag = std::abs(rdiag[j]); + if (!std::isfinite(abs_diag)) + { + full_rank = false; + } + if (abs_diag > max_diag) + { + max_diag = abs_diag; + } + } + const double rank_tol = 1e-6 * max_diag; + for (int j = 0; j < rdiag_size && full_rank; ++j) + { + if (std::abs(rdiag[j]) <= rank_tol) + { + full_rank = false; + } + } + + if (full_rank) + { + op.spsi(basis + dim * nbase, spsi + dim * nbase, dim, notconv); + qr_used = true; + } + else + { + syncmem_complex_op()(basis + dim * nbase, qr_backup, dim * notconv); + ModuleBase::WARNING("DiagoDavid::cal_grad", + "CUDA float QR block is rank deficient; falling back to " + "Schmidt orthogonalization."); + } + delmem_complex_op()(qr_backup); + } + if (!qr_used) +#endif + { + std::vector valid_indices; + valid_indices.reserve(notconv); + for (int m = 0; m < notconv; m++) + { + if (this->SchmidtOrth(dim, + nbase + notconv, + nbase + m, + spsi, + &lagrange[m * (nbase + notconv)], + pre_matrix_mm_m[m], + pre_matrix_mv_m[m])) + { + valid_indices.push_back(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); + } + } + + // Drop rejected directions and compact the remaining columns so that + // cal_elem() does not see zero columns or enlarge the basis with them. + if (static_cast(valid_indices.size()) < notconv) + { + for (int k = 0; k < static_cast(valid_indices.size()); ++k) + { + const int src = valid_indices[k]; + if (src != k) + { + syncmem_complex_op()(basis + dim * (nbase + k), basis + dim * (nbase + src), dim); + syncmem_complex_op()(spsi + (nbase + k) * dim, spsi + (nbase + src) * dim, dim); + } + } + this->notconv = static_cast(valid_indices.size()); } } + // calculate H|psi> for not convergence bands // hpsi_info dav_hpsi_in(&basis, // psi::Range(true, 0, nbase, nbase + notconv - 1), // &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); + // An all-rejected block leaves no new direction; asking the operator for + // zero vectors would build an invalid [0, -1] range. + if (this->notconv > 0) + { + op.hpsi(basis + nbase * dim, hpsi + nbase * dim, dim, this->notconv); + } delmem_complex_op()(lagrange); delmem_complex_op()(vc_ev_vector); @@ -666,7 +830,23 @@ void DiagoDavid::diag_zhegvx(const int& nbase, 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); + if (std::is_same::value && !this->use_uspp) + { + // Use double only for this single reduced eigensolve. + auto h64 = ct::TensorMap(const_cast(hcc), ct::DataTypeToEnum::value, + ct::DeviceType::GpuDevice, {nbase, nbase_x}).cast>(); + ct::Tensor v64(ct::DataType::DT_COMPLEX_DOUBLE, ct::DeviceType::GpuDevice, {nbase, nbase_x}); + ct::Tensor e64(ct::DataType::DT_DOUBLE, ct::DeviceType::GpuDevice, {nbase_x}); + ct::kernels::lapack_heevx, ct_Device>()( + nbase, nbase_x, h64.data>(), nband, e64.data(), v64.data>()); + base_device::memory::cast_memory_op, Device, Device>()( + vcc, v64.data>(), static_cast(nband) * nbase_x); + base_device::memory::cast_memory_op()(eigenvalue_gpu, e64.data(), nband); + } + else + { + 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); @@ -685,9 +865,9 @@ 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); + Parallel_Common::bcast_data(&vcc[i * nbase_x], nbase, diag_comm.comm, 0); } - MPI_Bcast(this->eigenvalue, nband, MPI_DOUBLE, 0, diag_comm.comm); + Parallel_Common::bcast_data(this->eigenvalue, nband, diag_comm.comm, 0); } #endif @@ -820,7 +1000,7 @@ void DiagoDavid::refresh(const int& dim, template -void DiagoDavid::SchmidtOrth(const int& dim, +bool DiagoDavid::SchmidtOrth(const int& dim, const int nband, const int m, const T* spsi, @@ -845,6 +1025,8 @@ void DiagoDavid::SchmidtOrth(const int& dim, if (dim == 0) { // GEMV quick-returns on zero rows without overwriting the projection. + // Do not return early: every rank, including empty local blocks, must + // reach the pool reduction below with the same collective sequence. const int projection_size = m + 1; setmem_complex_op()(lagrange_m, 0, projection_size); } @@ -896,7 +1078,16 @@ void DiagoDavid::SchmidtOrth(const int& dim, syncmem_d2h_op()(&var, lagrange_m + m, 1); double psi_norm = get_real(var); - assert(psi_norm > 0.0); + if (psi_norm <= 0.0) + { + // The residual direction has a zero (or negative) S-norm: in single + // precision the residual H*c - lambda*S*c suffers catastrophic + // cancellation and collapses to numerical noise. Treat it as converged + // and leave the (near-zero) direction untouched instead of aborting. + warn_skip_degenerate_direction(m); + ModuleBase::timer::end("DiagoDavid", "SchmidtOrth"); + return false; + } // / psi_m = psi_m - \sum_{i < m} \langle psi(i)|S|psi(m) \rangle psi(i) // psi_m = psi_m - basis * lagrange_m @@ -915,44 +1106,30 @@ void DiagoDavid::SchmidtOrth(const int& dim, // psi_norm = psi_norm - lagrange_m ยท lagrange_m psi_norm -= ModuleBase::dot_real_op()(m, lagrange_m, lagrange_m, false); - // for (int j = 0; j < m; j++) - // { - // const std::complex alpha = std::complex(-1, 0) * lagrange_m[j]; - // zaxpy_(&npw, &alpha, &psi(j,0), &inc, psi_m, &inc); - // /*for (int ig = 0; ig < npw; ig++) - // { - // psi_m[ig] -= lagrange[j] * psi(j, ig); - // }*/ - // psi_norm -= (conj(lagrange_m[j]) * lagrange_m[j]).real(); - // } - - assert(psi_norm > 0.0); + if (psi_norm <= 0.0) + { + // Nearly rank-deficient direction after orthogonalization: skip it. + warn_skip_degenerate_direction(m); + ModuleBase::timer::end("DiagoDavid", "SchmidtOrth"); + return false; + } psi_norm = sqrt(psi_norm); if (psi_norm < 1.0e-12) { - std::cout << "DiagoDavid::SchmidtOrth:aborted for psi_norm <1.0e-12" << std::endl; - std::cout << "This may be due to npwx < nbands: the number of plane waves is less than" << std::endl; - std::cout << "the number of bands, leading to a rank-deficient problem." << std::endl; - std::cout << "Please increase ecutwfc or reduce nbands." << std::endl; - std::cout << "nband = " << nband << std::endl; - std::cout << "m = " << m << std::endl; - exit(0); - } - else - { - // psi_m = psi_m / 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; - // } + // Nearly rank-deficient direction: skip instead of aborting. + warn_skip_degenerate_direction(m); + ModuleBase::timer::end("DiagoDavid", "SchmidtOrth"); + return false; } + // psi_m = psi_m / psi_norm + ModuleBase::vector_mul_real_op()(dim, psi_m, psi_m, Real(1.0 / psi_norm)); + // delete[] lagrange; ModuleBase::timer::end("DiagoDavid", "SchmidtOrth"); - return; + return true; } @@ -1050,6 +1227,15 @@ int DiagoDavid::diag(const HSOperator& op, } +template +inline bool DiagoDavid::use_block_qr() const +{ + return this->device == base_device::GpuDevice + && std::is_same::value + && !this->use_uspp + && this->diag_comm.nproc == 1; +} + template inline bool DiagoDavid::check_block_conv(const int& ntry, const int& notconv, diff --git a/source/source_hsolver/diago_david.h b/source/source_hsolver/diago_david.h index 5dd29805378..f9c28f5e462 100644 --- a/source/source_hsolver/diago_david.h +++ b/source/source_hsolver/diago_david.h @@ -56,6 +56,7 @@ class DiagoDavid const int nband_in, const int dim_in, const int david_ndim_in, + const bool use_uspp_in, const diag_comm_info& diag_comm_in); /** @@ -109,6 +110,8 @@ class DiagoDavid const int nbase_x; /// dimension of the subspace allowed in Davidson const int david_ndim = 4; + /// whether the SCF uses ultrasoft pseudopotentials + const bool use_uspp = false; /// number of unconverged eigenvalues int notconv = 0; @@ -230,7 +233,7 @@ 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, + bool SchmidtOrth(const int& dim, const int nband, const int m, const T* spsi, @@ -278,6 +281,9 @@ class DiagoDavid */ bool check_block_conv(const int &ntry, const int ¬conv, const int &ntry_max, const int ¬conv_max); + /// @brief Whether to form CUDA float Davidson bases with a single-pass QR. + bool use_block_qr() const; + using resmem_complex_op = base_device::memory::resize_memory_op; using delmem_complex_op = base_device::memory::delete_memory_op; using setmem_complex_op = base_device::memory::set_memory_op; diff --git a/source/source_hsolver/hsolver_pw.cpp b/source/source_hsolver/hsolver_pw.cpp index 1b316b567e0..0228443e7ec 100644 --- a/source/source_hsolver/hsolver_pw.cpp +++ b/source/source_hsolver/hsolver_pw.cpp @@ -335,7 +335,7 @@ void HSolverPW::hamiltSolvePsiK(const HSOperator& op, 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); + DiagoDavid david(pre_condition.data(), nband, dim, this->pw_diag_ndim, this->use_uspp, comm_info); // do diag and add davidson iteration counts up to avg_iter DiagoIterAssist::avg_iter += static_cast( david.diag(op, diff --git a/source/source_hsolver/kernels/bpcg_kernel_op.cpp b/source/source_hsolver/kernels/bpcg_kernel_op.cpp index 3b04f1390e2..5c9523fd17f 100644 --- a/source/source_hsolver/kernels/bpcg_kernel_op.cpp +++ b/source/source_hsolver/kernels/bpcg_kernel_op.cpp @@ -194,7 +194,16 @@ struct normalize_op { psi_iter + (nbase + m) * dim, psi_iter + (nbase + m) * dim, true); - assert(psi_m_norm > 0.0); + // A zero (or NaN) norm means this direction is degenerate; in single + // precision tiny residual norms can underflow to exactly zero. Skip the + // normalization and let the caller treat the zero-norm band as converged + // (see Diago_DavSubspace::cal_grad). Also skip denormal/tiny norms to + // avoid amplifying them into huge directions that corrupt the basis. + if (psi_m_norm <= 1.0e-24) + { + if (psi_norm) psi_norm[m] = 0.0; + continue; + } psi_m_norm = sqrt(psi_m_norm); // Normalize using vector_div_constant_op diff --git a/source/source_hsolver/kernels/cuda/bpcg_kernel_op.cu b/source/source_hsolver/kernels/cuda/bpcg_kernel_op.cu index 48113fe1d43..dca51c8306c 100644 --- a/source/source_hsolver/kernels/cuda/bpcg_kernel_op.cu +++ b/source/source_hsolver/kernels/cuda/bpcg_kernel_op.cu @@ -356,9 +356,14 @@ __global__ void normalize_kernel( Real l2_sq = blockReduceSum(local_sum, shared); Real norm = sqrt(l2_sq); - // Normalize the vector - for (int i = tid; i < dim; i += thread_per_block) { - psi_iter[(nbase + m) * dim + i] /= norm; + // Guard against zero/denormal norm (degenerate direction, e.g. single-precision + // underflow): skip normalization and leave psi_norm[m] = 0 so the caller treats + // it as converged. + if (norm > 1.0e-12) { + // Normalize the vector + for (int i = tid; i < dim; i += thread_per_block) { + psi_iter[(nbase + m) * dim + i] /= norm; + } } // Store the norm if needed 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..ad22ebf812c 100644 --- a/source/source_hsolver/kernels/rocm/bpcg_kernel_op.hip.cu +++ b/source/source_hsolver/kernels/rocm/bpcg_kernel_op.hip.cu @@ -256,9 +256,14 @@ __global__ void normalize_kernel( Real norm = sqrt(sum[0]); - // Normalize the vector - for (int i = tid; i < dim; i += THREAD_PER_BLOCK) { - psi_iter[(nbase + m) * dim + i] /= norm; + // Guard against zero/denormal norm (degenerate direction, e.g. single-precision + // underflow): skip normalization and leave psi_norm[m] = 0 so the caller treats + // it as converged. + if (norm > 1.0e-12) { + // Normalize the vector + for (int i = tid; i < dim; i += THREAD_PER_BLOCK) { + psi_iter[(nbase + m) * dim + i] /= norm; + } } // Store the norm if needed diff --git a/source/source_hsolver/test/diago_david_float_test.cpp b/source/source_hsolver/test/diago_david_float_test.cpp index 930b4920cb6..5aeb6b52743 100644 --- a/source/source_hsolver/test/diago_david_float_test.cpp +++ b/source/source_hsolver/test/diago_david_float_test.cpp @@ -93,7 +93,7 @@ class DiagoDavPrepare const int dim = phi.get_current_ngk() ; const int nband = phi.get_nbands(); const int ld_psi =phi.get_nbasis(); - hsolver::DiagoDavid> dav(precondition, nband, dim, order, comm_info); + hsolver::DiagoDavid> dav(precondition, nband, dim, order, false, comm_info); hsolver::DiagoIterAssist>::PW_DIAG_NMAX = maxiter; hsolver::DiagoIterAssist>::PW_DIAG_THR = eps; diff --git a/source/source_hsolver/test/diago_david_real_test.cpp b/source/source_hsolver/test/diago_david_real_test.cpp index 3ff4e33bc7b..630006ea83c 100644 --- a/source/source_hsolver/test/diago_david_real_test.cpp +++ b/source/source_hsolver/test/diago_david_real_test.cpp @@ -92,7 +92,7 @@ class DiagoDavPrepare const int dim = phi.get_current_ngk(); const int nband = phi.get_nbands(); const int ld_psi = phi.get_nbasis(); - hsolver::DiagoDavid dav(precondition, nband, dim, order, comm_info); + hsolver::DiagoDavid dav(precondition, nband, dim, order, false, comm_info); hsolver::DiagoIterAssist::PW_DIAG_NMAX = maxiter; hsolver::DiagoIterAssist::PW_DIAG_THR = eps; diff --git a/source/source_hsolver/test/diago_david_test.cpp b/source/source_hsolver/test/diago_david_test.cpp index 7d4b8cbfea0..7deefc68855 100644 --- a/source/source_hsolver/test/diago_david_test.cpp +++ b/source/source_hsolver/test/diago_david_test.cpp @@ -97,7 +97,7 @@ class DiagoDavPrepare const int dim = phi.get_current_ngk(); const int nband = phi.get_nbands(); const int ld_psi = phi.get_nbasis(); - hsolver::DiagoDavid> dav(precondition, nband, dim, order, comm_info); + hsolver::DiagoDavid> dav(precondition, nband, dim, order, false, comm_info); hsolver::DiagoIterAssist>::PW_DIAG_NMAX = maxiter; hsolver::DiagoIterAssist>::PW_DIAG_THR = eps; diff --git a/source/source_hsolver/test/hsolver_pw_sup.h b/source/source_hsolver/test/hsolver_pw_sup.h index 75ec3a90500..5c8dfc6d94f 100644 --- a/source/source_hsolver/test/hsolver_pw_sup.h +++ b/source/source_hsolver/test/hsolver_pw_sup.h @@ -127,8 +127,10 @@ DiagoDavid::DiagoDavid(const Real* precondition_in, const int nband_in, const int dim_in, const int david_ndim_in, + const bool use_uspp_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) { + (void)use_uspp_in; this->device = base_device::get_device_type(this->ctx); this->precondition = precondition_in; diff --git a/source/source_hsolver/test/test_diago_david.cpp b/source/source_hsolver/test/test_diago_david.cpp index b2556513944..4d155774a79 100644 --- a/source/source_hsolver/test/test_diago_david.cpp +++ b/source/source_hsolver/test/test_diago_david.cpp @@ -25,6 +25,7 @@ class DiagonalOperator final : public hsolver::HSOperator void hpsi(const T* x, T* hx, const int ld, const int nvec) const override { EXPECT_EQ(ld, size_); + EXPECT_GT(nvec, 0); for (int band = 0; band < nvec; ++band) { @@ -37,6 +38,7 @@ class DiagonalOperator final : public hsolver::HSOperator void spsi(const T* x, T* sx, const int ld, const int nvec) const override { EXPECT_EQ(ld, size_); + EXPECT_GT(nvec, 0); const int count = size_ * nvec; std::copy_n(x, count, sx); } @@ -124,7 +126,7 @@ TYPED_TEST(DavidsonOwnershipTest, PhysicalEmptyAndUnevenBlocks) DiagonalOperator op(offset, local_size); if (method == 0) { - hsolver::DiagoDavid solver(precondition.data(), bands, local_size, 3, comm); + hsolver::DiagoDavid solver(precondition.data(), bands, local_size, 3, false, comm); solver.diag(op, local_size, vectors.data(), energies, thresholds, 300, 5, 0); } else if (method == 1) @@ -167,6 +169,95 @@ TYPED_TEST(DavidsonOwnershipTest, PhysicalEmptyAndUnevenBlocks) } } } + +TYPED_TEST(DavidsonOwnershipTest, ExactAndApproximateWarmStart) +{ + using T = TypeParam; + const auto domain = ModuleBase::world_comm_domain(); + const int rank = domain.rank(); + const int processes = domain.size(); + const int bands = 2; + // david_ndim=3 gives nbase_x=6; keep the asserted global dimension above it + // even when the test runs serially with a single process. + const int dimension = std::max(8, 2 * processes); +#ifdef __MPI + const hsolver::diag_comm_info comm = {POOL_WORLD, rank, processes}; +#else + const hsolver::diag_comm_info comm = {rank, processes}; +#endif + + std::vector sizes(processes, 0); + for (int process = 0; process < processes; ++process) + { + sizes[process] = dimension / processes; + if (process < dimension % processes) + { + ++sizes[process]; + } + } + int offset = 0; + for (int process = 0; process < rank; ++process) + { + offset += sizes[process]; + } + const int local_size = sizes[rank]; + const int storage_size = std::max(local_size * bands, 1); + + std::vector precondition(std::max(local_size, 1), 1.0); + for (int method = 0; method < 2; ++method) + { + SCOPED_TRACE(::testing::Message() << "method=" << method << " rank=" << rank); + std::vector vectors(storage_size, T(0)); + for (int band = 0; band < bands; ++band) + { + for (int row = 0; row < local_size; ++row) + { + const int global_row = offset + row; + if (band == 0 && global_row == 0) + { + vectors[band * local_size + row] = T(1); + } + if (band == 1 && (global_row == 1 || global_row == 2)) + { + vectors[band * local_size + row] = T(1.0 / std::sqrt(2.0)); + } + } + } + + double energies[bands] = {}; + const std::vector thresholds(bands, 1e-8); + DiagonalOperator op(offset, local_size); + if (method == 0) + { + hsolver::DiagoDavid solver(precondition.data(), bands, local_size, 3, false, comm); + solver.diag(op, local_size, vectors.data(), energies, thresholds, 100, 5, 0); + } + else + { + hsolver::Diago_DavSubspace solver(precondition, bands, local_size, 3, 1e-8, 100, comm, 0, 1); + solver.diag(op, vectors.data(), local_size, energies, thresholds, false); + } + + for (int band = 0; band < bands; ++band) + { + EXPECT_NEAR(energies[band], band + 1.0, 1e-5); + double residual = 0.0; + double norm = 0.0; + for (int row = 0; row < local_size; ++row) + { + const double amplitude = std::abs(vectors[band * local_size + row]); + const double error = (offset + row + 1.0 - energies[band]) * amplitude; + residual += error * error; + norm += amplitude * amplitude; + } + Parallel_Reduce::reduce_pool(residual); + Parallel_Reduce::reduce_pool(norm); + EXPECT_LT(std::sqrt(residual), 1e-4); + EXPECT_NEAR(norm, 1.0, 1e-6); + } + } +} + } // namespace int main(int argc, char** argv) diff --git a/source/source_lcao/module_lr/hsolver_lrtd.cpp b/source/source_lcao/module_lr/hsolver_lrtd.cpp index 94d2526bdc7..1bdfc37a644 100644 --- a/source/source_lcao/module_lr/hsolver_lrtd.cpp +++ b/source/source_lcao/module_lr/hsolver_lrtd.cpp @@ -114,6 +114,7 @@ namespace LR nband, dim, davidson_dimension, + false, comm_info); std::vector ethr_band(nband, diag_ethr); hsolver::DiagoIterAssist::avg_iter += static_cast(david.diag(op, diff --git a/source/source_pw/module_pwdft/fs_kin_tools.h b/source/source_pw/module_pwdft/fs_kin_tools.h index 78afc835c5c..5e6566edcdb 100644 --- a/source/source_pw/module_pwdft/fs_kin_tools.h +++ b/source/source_pw/module_pwdft/fs_kin_tools.h @@ -48,7 +48,7 @@ class FS_Kin_tools std::vector gk3_; std::vector gk; std::vector kfac; - std::vector s_kin; + std::vector s_kin; FPTYPE* d_gk = nullptr; FPTYPE* d_kfac = nullptr; const FPTYPE* wg = nullptr; diff --git a/source/source_pw/module_pwdft/kernels/cuda/force_op.cu b/source/source_pw/module_pwdft/kernels/cuda/force_op.cu index 806394a74f5..d34bba73f2a 100644 --- a/source/source_pw/module_pwdft/kernels/cuda/force_op.cu +++ b/source/source_pw/module_pwdft/kernels/cuda/force_op.cu @@ -933,8 +933,10 @@ void cal_force_ew_op::operator()( // for revertVkbValues functions instantiation template void revertVkbValues(const int *gcar_zero_ptrs, std::complex *vkb_ptr, const std::complex *vkb_save_ptr, int nkb, int gcar_zero_count, int npw, int ipol, int npwx, const std::complex coeff); +template void revertVkbValues(const int *gcar_zero_ptrs, std::complex *vkb_ptr, const std::complex *vkb_save_ptr, int nkb, int gcar_zero_count, int npw, int ipol, int npwx, const std::complex coeff); // for saveVkbValues functions instantiation template void saveVkbValues(const int *gcar_zero_ptrs, const std::complex *vkb_ptr, std::complex *vkb_save_ptr, int nkb, int gcar_zero_count, int npw, int ipol, int npwx); +template void saveVkbValues(const int *gcar_zero_ptrs, const std::complex *vkb_ptr, std::complex *vkb_save_ptr, int nkb, int gcar_zero_count, int npw, int ipol, int npwx); template struct cal_vkb1_nl_op; template struct cal_force_nl_op; diff --git a/source/source_pw/module_pwdft/kernels/cuda/stress_op.cu b/source/source_pw/module_pwdft/kernels/cuda/stress_op.cu index da5c9f1c1b7..e40796f1fa0 100644 --- a/source/source_pw/module_pwdft/kernels/cuda/stress_op.cu +++ b/source/source_pw/module_pwdft/kernels/cuda/stress_op.cu @@ -1006,7 +1006,7 @@ __global__ void cal_stress_onsite( const int *atom_nh, const int *atom_na, const FPTYPE *d_wg, - const double* lambda, + const FPTYPE* lambda, const thrust::complex *becp, const thrust::complex *dbecp, FPTYPE *stress) @@ -1127,7 +1127,7 @@ void cal_stress_nl_op::operator()(const base_de const int* atom_nh, const int* atom_na, const FPTYPE* d_wg, - const double* lambda, + const FPTYPE* lambda, const int* isk, const std::complex* becp, const std::complex* dbecp, diff --git a/source/source_pw/module_pwdft/kernels/rocm/force_op.hip.cu b/source/source_pw/module_pwdft/kernels/rocm/force_op.hip.cu index 68aa6a9ec59..a38c4b55dfd 100644 --- a/source/source_pw/module_pwdft/kernels/rocm/force_op.hip.cu +++ b/source/source_pw/module_pwdft/kernels/rocm/force_op.hip.cu @@ -700,8 +700,10 @@ void revertVkbValues( // for revertVkbValues functions instantiation template void revertVkbValues(const int *gcar_zero_ptrs, std::complex *vkb_ptr, const std::complex *vkb_save_ptr, int nkb, int gcar_zero_count, int npw, int ipol, int npwx, const std::complex coeff); +template void revertVkbValues(const int *gcar_zero_ptrs, std::complex *vkb_ptr, const std::complex *vkb_save_ptr, int nkb, int gcar_zero_count, int npw, int ipol, int npwx, const std::complex coeff); // for saveVkbValues functions instantiation template void saveVkbValues(const int *gcar_zero_ptrs, const std::complex *vkb_ptr, std::complex *vkb_save_ptr, int nkb, int gcar_zero_count, int npw, int ipol, int npwx); +template void saveVkbValues(const int *gcar_zero_ptrs, const std::complex *vkb_ptr, std::complex *vkb_save_ptr, int nkb, int gcar_zero_count, int npw, int ipol, int npwx); template struct cal_vkb1_nl_op; template struct cal_force_nl_op; diff --git a/source/source_pw/module_pwdft/kernels/rocm/stress_op.hip.cu b/source/source_pw/module_pwdft/kernels/rocm/stress_op.hip.cu index 966badd83e4..73c8441ea6b 100644 --- a/source/source_pw/module_pwdft/kernels/rocm/stress_op.hip.cu +++ b/source/source_pw/module_pwdft/kernels/rocm/stress_op.hip.cu @@ -995,7 +995,7 @@ __global__ void cal_stress_onsite( const int *atom_nh, const int *atom_na, const FPTYPE *d_wg, - const double* lambda, + const FPTYPE* lambda, const thrust::complex *becp, const thrust::complex *dbecp, FPTYPE *stress) @@ -1116,7 +1116,7 @@ void cal_stress_nl_op::operator()(const base_de const int* atom_nh, const int* atom_na, const FPTYPE* d_wg, - const double* lambda, + const FPTYPE* lambda, const int* isk, const std::complex* becp, const std::complex* dbecp, diff --git a/source/source_pw/module_pwdft/kernels/stress_op.h b/source/source_pw/module_pwdft/kernels/stress_op.h index fd61c483550..28f3943d451 100644 --- a/source/source_pw/module_pwdft/kernels/stress_op.h +++ b/source/source_pw/module_pwdft/kernels/stress_op.h @@ -148,7 +148,7 @@ struct cal_stress_nl_op const int* atom_nh, const int* atom_na, const FPTYPE* d_wg, - const double* lambda, + const FPTYPE* lambda, const int* isk, const std::complex* becp, const std::complex* dbecp, @@ -356,7 +356,7 @@ struct cal_stress_nl_op const int* atom_nh, const int* atom_na, const FPTYPE* d_wg, - const double* lambda, + const FPTYPE* lambda, const int* isk, const std::complex* becp, const std::complex* dbecp, diff --git a/source/source_pw/module_pwdft/nonlocal_maths.cpp b/source/source_pw/module_pwdft/nonlocal_maths.cpp index fb1ff95a912..3307aeb03ad 100644 --- a/source/source_pw/module_pwdft/nonlocal_maths.cpp +++ b/source/source_pw/module_pwdft/nonlocal_maths.cpp @@ -287,6 +287,19 @@ void Nonlocal_maths::dylmr2(const int nylm, FPTYPE* dylm, const int ipol) { + // A 1e-6 finite difference loses significant digits when evaluated in float. + // Keep the difference in double, then store the derivative in the requested type. + if (sizeof(FPTYPE) < sizeof(double)) + { + const std::vector q(gk, gk + 3 * ngy); + std::vector derivative(nylm * ngy); + Nonlocal_maths::dylmr2(nylm, ngy, q.data(), derivative.data(), ipol); + for (int i = 0; i < nylm * ngy; ++i) + { + dylm[i] = static_cast(derivative[i]); + } + return; + } //----------------------------------------------------------------------- // // compute \partial Y_lm(G) \over \partial (G)_ipol diff --git a/source/source_pw/module_pwdft/stress_cc.cpp b/source/source_pw/module_pwdft/stress_cc.cpp index a0b6cb3acb3..10692e5fa51 100644 --- a/source/source_pw/module_pwdft/stress_cc.cpp +++ b/source/source_pw/module_pwdft/stress_cc.cpp @@ -35,7 +35,7 @@ void Stress_Func::stress_cc(ModuleBase::matrix& sigma, fact = 2.0; //is_pw:PW basis, gamma_only need to FPTYPE. } - FPTYPE sigmadiag; + double sigmadiag; FPTYPE* rhocg = nullptr; int judge=0; diff --git a/source/source_pw/module_pwdft/stress_ewa.cpp b/source/source_pw/module_pwdft/stress_ewa.cpp index 8134da9d4c2..3a3702e702e 100644 --- a/source/source_pw/module_pwdft/stress_ewa.cpp +++ b/source/source_pw/module_pwdft/stress_ewa.cpp @@ -43,9 +43,10 @@ void Stress_Func::stress_ewa(const UnitCell& ucell, } while(upperbound>1e-7); - //G-space sum here - //Determine if this processor contains G=0 and set the constant term - FPTYPE sdewald=0.0; + // G-space sum here + // Determine if this processor contains G=0 and set the constant term + // Long reciprocal-space sums need double accumulation in both precision modes. + double sdewald = 0.0; const int ig0 = rho_basis->ig_gge0; if( ig0 >= 0) { @@ -66,9 +67,9 @@ void Stress_Func::stress_ewa(const UnitCell& ucell, // else fact=1.0; #pragma omp parallel -{ - ModuleBase::matrix local_sigma(3, 3); - FPTYPE local_sdewald = 0; + { + ModuleBase::matrix local_sigma(3, 3); + double local_sdewald = 0.0; FPTYPE g2,g2a; FPTYPE arg; diff --git a/source/source_pw/module_pwdft/stress_loc.cpp b/source/source_pw/module_pwdft/stress_loc.cpp index 94f1746dbe0..33b9f9a7f0d 100644 --- a/source/source_pw/module_pwdft/stress_loc.cpp +++ b/source/source_pw/module_pwdft/stress_loc.cpp @@ -20,9 +20,10 @@ void Stress_Func::stress_loc(const UnitCell& ucell, ModuleBase::TITLE("Stress","stress_loc"); ModuleBase::timer::start("Stress","stress_loc"); - std::vector dvloc(rho_basis->npw); - FPTYPE evloc=0.0; - FPTYPE fact=1.0; + std::vector dvloc(rho_basis->npw); + // The diagonal potential term sums all plane waves on each MPI rank. + double evloc = 0.0; + FPTYPE fact = 1.0; const int nspin_rho = (PARAM.inp.nspin == 2) ? 2 : 1; diff --git a/source/source_pw/module_pwdft/stress_mgga.cpp b/source/source_pw/module_pwdft/stress_mgga.cpp index 7ecbe1cb914..1df76fc5779 100644 --- a/source/source_pw/module_pwdft/stress_mgga.cpp +++ b/source/source_pw/module_pwdft/stress_mgga.cpp @@ -7,6 +7,8 @@ #include #include +#include +#include "source_basis/module_pw/pw_basis_k.h" // calculate the Pulay term of mGGA stress correction in PW template @@ -31,10 +33,37 @@ void Stress_Func::stress_mgga(const UnitCell& ucell, std::complex* psi; int ipol2xy[3][3] = {{0, 1, 3}, {1, 2, 4}, {3, 4, 5}}; - FPTYPE sigma_mgga[3][3]; + double sigma_mgga[3][3]; using ct_Device = typename ct::PsiToContainer::type; + // gcar/kvec_c are constant over bands: prepare device buffers once. + const int gcar_size = wfc_basis->nks * wfc_basis->npwk_max; + std::vector gcar_host(3 * gcar_size); + std::vector kvec_host(3 * wfc_basis->nks); + for (int i = 0; i < gcar_size; ++i) + { + gcar_host[3 * i] = static_cast(wfc_basis->gcar[i].x); + gcar_host[3 * i + 1] = static_cast(wfc_basis->gcar[i].y); + gcar_host[3 * i + 2] = static_cast(wfc_basis->gcar[i].z); + } + for (int i = 0; i < wfc_basis->nks; ++i) + { + kvec_host[3 * i] = static_cast(wfc_basis->kvec_c[i].x); + kvec_host[3 * i + 1] = static_cast(wfc_basis->kvec_c[i].y); + kvec_host[3 * i + 2] = static_cast(wfc_basis->kvec_c[i].z); + } + auto gcar_device = ct::TensorMap(gcar_host.data(), + ct::DataTypeToEnum::value, + ct::DeviceType::CpuDevice, + {gcar_size, 3}) + .to_device(); + auto kvec_device = ct::TensorMap(kvec_host.data(), + ct::DataTypeToEnum::value, + ct::DeviceType::CpuDevice, + {wfc_basis->nks, 3}) + .to_device(); + auto gradwfc = ct::Tensor( ct::DataTypeToEnum>::value, ct::DeviceTypeToEnum::value, @@ -59,10 +88,13 @@ void Stress_Func::stress_mgga(const UnitCell& ucell, { const FPTYPE w1 = wg(ik, ibnd) / ucell.omega; const std::complex* psi = &psi_in[0](ik, ibnd, 0); - XC_Functional::grad_wfc, Device>(ik, - ucell.tpiba, - wfc_basis, - psi, + XC_Functional::grad_wfc, Device>( + ik, + ucell.tpiba, + wfc_basis, + gcar_device.template data(), + kvec_device.template data(), + psi, gradwfc.data>()); cal_stress_mgga_solver( diff --git a/source/source_pw/module_pwdft/stress_nl.cpp b/source/source_pw/module_pwdft/stress_nl.cpp index 21503be82b6..327f6d471f2 100644 --- a/source/source_pw/module_pwdft/stress_nl.cpp +++ b/source/source_pw/module_pwdft/stress_nl.cpp @@ -1,3 +1,5 @@ +#include + #include "source_base/math_polyint.h" #include "source_base/parallel_reduce.h" #include "source_io/module_parameter/parameter.h" @@ -28,14 +30,19 @@ void Stress_Func::stress_nl(ModuleBase::matrix& sigma, } ModuleBase::timer::start("Stress", "stress_nl"); + const int nks = p_kv->get_nks(); + constexpr bool split_k = std::is_same::value + && std::is_same::value; + const int stress_count = 9 * (split_k && nks > 0 ? nks : 1); FPTYPE* stress_device = nullptr; - resmem_var_op()(stress_device, 9); - setmem_var_op()(stress_device, 0, 9); - std::vector sigmanlc(9, 0.0); + resmem_var_op()(stress_device, stress_count); + setmem_var_op()(stress_device, 0, stress_count); + std::vector stress_host(stress_count, 0.0); + using SumType = typename std::conditional::type; + std::vector sigmanlc(9, 0.0); hamilt::FS_Nonlocal_tools nl_tools(&nlpp_in, &ucell_in, p_kv, wfc_basis, p_sf, wg, &ekb); - const int nks = p_kv->get_nks(); const int max_nbands = wg.nc; for (int ik = 0; ik < nks; ik++) // loop k points { @@ -62,13 +69,20 @@ void Stress_Func::stress_nl(ModuleBase::matrix& sigma, for (int jpol = 0; jpol <= ipol; jpol++) { nl_tools.cal_vkb_deri_s(ik, max_nbands, ipol, jpol); - nl_tools.cal_dbecp_s(ik, npm, &psi_in[0](ik,0,0)); - nl_tools.cal_stress(ik, npm, true, ipol, jpol, stress_device); + nl_tools.cal_dbecp_s(ik, npm, &psi_in[0](ik, 0, 0)); + nl_tools.cal_stress(ik, npm, true, ipol, jpol, stress_device + (split_k ? 9 * ik : 0)); } } } // transfer stress from device to host - syncmem_var_d2h_op()(sigmanlc.data(), stress_device, 9); + syncmem_var_d2h_op()(stress_host.data(), stress_device, stress_count); + for (int ik = 0; ik < (split_k ? nks : 1); ++ik) + { + for (int i = 0; i < 9; ++i) + { + sigmanlc[i] += static_cast(stress_host[9 * ik + i]); + } + } delmem_var_op()(stress_device); // sum up forcenl from all processors for (int l = 0; l < 3; l++) diff --git a/source/source_pw/module_pwdft/vnl_pw.cpp b/source/source_pw/module_pwdft/vnl_pw.cpp index 233e0590bca..b510a338e70 100644 --- a/source/source_pw/module_pwdft/vnl_pw.cpp +++ b/source/source_pw/module_pwdft/vnl_pw.cpp @@ -169,11 +169,10 @@ void pseudopot_cell_vnl::init(const UnitCell& ucell, resmem_cd_op()(c_deeq_nc, PARAM.inp.nspin * ucell.nat * this->nhm * this->nhm); resmem_cd_op()(c_qq_so, ntype * 4 * this->nhm * this->nhm); } - if (PARAM.globalv.has_double_data) - { - resmem_zd_op()(z_deeq_nc, PARAM.inp.nspin * ucell.nat * this->nhm * this->nhm); - resmem_zd_op()(z_qq_so, ntype * 4 * this->nhm * this->nhm); - } + // The double-precision force/stress path always reads these + // buffers, even when the SCF runs in single precision. + resmem_zd_op()(z_deeq_nc, PARAM.inp.nspin * ucell.nat * this->nhm * this->nhm); + resmem_zd_op()(z_qq_so, ntype * 4 * this->nhm * this->nhm); resmem_dd_op()(d_deeq, PARAM.inp.nspin * ucell.nat * this->nhm * this->nhm); resmem_dd_op()(d_indv, ntype * this->nhm); resmem_dd_op()(d_nhtol, ntype * this->nhm); @@ -192,11 +191,11 @@ void pseudopot_cell_vnl::init(const UnitCell& ucell, resmem_ch_op()(c_deeq_nc, PARAM.inp.nspin * ucell.nat * this->nhm * this->nhm, "VNL::c_deeq_nc"); resmem_ch_op()(c_qq_so, ntype * 4 * this->nhm * this->nhm, "VNL::c_qq_so"); } - if (PARAM.globalv.has_double_data) - { - this->z_deeq_nc = this->deeq_nc.ptr; - this->z_qq_so = this->qq_so.ptr; - } + // The double-precision deeq_nc / qq_so buffers are always required by + // the double-precision force/stress path (Forces / Stress_PW), + // even when the main SCF runs in single precision (has_float_data only). + this->z_deeq_nc = this->deeq_nc.ptr; + this->z_qq_so = this->qq_so.ptr; this->d_deeq = this->deeq.ptr; this->d_indv = this->indv.c; this->d_nhtol = this->nhtol.c; diff --git a/source/source_pw/module_pwdft/vnl_pw_deeq.cpp b/source/source_pw/module_pwdft/vnl_pw_deeq.cpp index 54ad733b274..beca6a4dd19 100644 --- a/source/source_pw/module_pwdft/vnl_pw_deeq.cpp +++ b/source/source_pw/module_pwdft/vnl_pw_deeq.cpp @@ -151,12 +151,10 @@ void pseudopot_cell_vnl::cal_effective_D(const ModuleBase::matrix& veff, this->deeq_nc.ptr, PARAM.inp.nspin * cell.nat * this->nhm * this->nhm); } - if (PARAM.globalv.has_double_data) - { - syncmem_z2z_h2d_op()(this->z_deeq_nc, - this->deeq_nc.ptr, - PARAM.inp.nspin * cell.nat * this->nhm * this->nhm); - } + // Always keep the double buffer in sync for force/stress. + syncmem_z2z_h2d_op()(this->z_deeq_nc, + this->deeq_nc.ptr, + PARAM.inp.nspin * cell.nat * this->nhm * this->nhm); syncmem_d2d_h2d_op()(this->d_deeq, this->deeq.ptr, PARAM.inp.nspin * cell.nat * this->nhm * this->nhm); diff --git a/source/source_pw/module_pwdft/vnl_pw_init_vnl.cpp b/source/source_pw/module_pwdft/vnl_pw_init_vnl.cpp index f44f2187d24..43cb80dcab7 100644 --- a/source/source_pw/module_pwdft/vnl_pw_init_vnl.cpp +++ b/source/source_pw/module_pwdft/vnl_pw_init_vnl.cpp @@ -329,7 +329,8 @@ void pseudopot_cell_vnl::init_vnl(UnitCell& cell, const ModulePW::PW_Basis* rho_ castmem_z2c_h2d_op()(this->c_qq_so, this->qq_so.ptr, this->qq_so.getSize()); } } - if (PARAM.globalv.has_double_data && this->nhm > 0) + // Always keep the double buffer in sync for force/stress. + if (this->nhm > 0) { syncmem_z2z_h2d_op()(this->z_qq_so, this->qq_so.ptr, this->qq_so.getSize()); }