Fix GPU single-precision DAV convergence, stress accuracy and meta-GGA XC contributions

python/pyabacus/src/hsolver/diago_adapter.hpp

@@ -85,7 +85,7 @@ class PyDiagoDavidAdapter
     {
         const PyHSOperator<T> op(mm_op);
 
-        solver_ = std::make_unique<SolverType>(precond_vec.data(), nband_, nbasis_, dav_ndim, comm_info);
+        solver_ = std::make_unique<SolverType>(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);
     }

source/source_base/kernels/cuda/math_kernel_op.cu

@@ -5,6 +5,7 @@
 
 #include <base/macros/macros.h>
 #include <cuda_runtime.h>
+#include <algorithm>
 #include <thrust/complex.h>
 #include <thrust/execution_policy.h>
 #include <thrust/inner_product.h>
@@ -293,7 +294,22 @@ void gemm_op<std::complex<float>, 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<float> 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 <>

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<float>;
+    using Zd = std::complex<double>;
+
+    struct Case
+    {
+        int m;
+        int n;
+        int k;
+        int lda;
+        int ldb;
+        int ldc;
+        Cf alpha;
+        Cf beta;
+    };
+
+    const std::vector<Case> 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<Cf> A(static_cast<size_t>(c.lda) * c.m, Cf(0.0f, 0.0f));
+        std::vector<Cf> B(static_cast<size_t>(c.ldb) * c.n, Cf(0.0f, 0.0f));
+        std::vector<Cf> C(static_cast<size_t>(c.ldc) * c.n, Cf(0.0f, 0.0f));
+        std::vector<Cf> C0(static_cast<size_t>(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<size_t>(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<size_t>(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<size_t>(j) * c.ldc] = Cf(0.1f * i - 0.2f, 0.05f * j + 0.1f);
+                C[i + static_cast<size_t>(j) * c.ldc] = C0[i + static_cast<size_t>(j) * c.ldc];
+            }
+        }
+
+        std::vector<Zd> ref(static_cast<size_t>(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<size_t>(i) * c.lda]))
+                           * Zd(B[l + static_cast<size_t>(j) * c.ldb]);
+                }
+                ref[i + static_cast<size_t>(j) * c.ldc]
+                    = Zd(c.alpha) * sum + Zd(c.beta) * Zd(C0[i + static_cast<size_t>(j) * c.ldc]);
+            }
+        }
+
+        Cf* dA = nullptr;
+        Cf* dB = nullptr;
+        Cf* dC = nullptr;
+        base_device::memory::resize_memory_op<Cf, base_device::DEVICE_GPU>()(dA, A.size());
+        base_device::memory::resize_memory_op<Cf, base_device::DEVICE_GPU>()(dB, B.size());
+        base_device::memory::resize_memory_op<Cf, base_device::DEVICE_GPU>()(dC, C.size());
+        base_device::memory::synchronize_memory_op<Cf, base_device::DEVICE_GPU, base_device::DEVICE_CPU>()(
+            dA, A.data(), A.size());
+        base_device::memory::synchronize_memory_op<Cf, base_device::DEVICE_GPU, base_device::DEVICE_CPU>()(
+            dB, B.data(), B.size());
+        base_device::memory::synchronize_memory_op<Cf, base_device::DEVICE_GPU, base_device::DEVICE_CPU>()(
+            dC, C.data(), C.size());
+
+        ModuleBase::gemm_op<Cf, base_device::DEVICE_GPU>()(
+            '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<Cf, base_device::DEVICE_CPU, base_device::DEVICE_GPU>()(
+            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<size_t>(j) * c.ldc])
+                                - ref[i + static_cast<size_t>(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<Cf, base_device::DEVICE_GPU>()(dA);
+        base_device::memory::delete_memory_op<Cf, base_device::DEVICE_GPU>()(dB);
+        base_device::memory::delete_memory_op<Cf, base_device::DEVICE_GPU>()(dC);
+    }
+    ModuleBase::destoryBLAShandle();
+}
+#endif // __UT_USE_CUDA
+
 #endif // __UT_USE_CUDA || __UT_USE_ROCM
 
 /*

source/source_base/module_container/ATen/kernels/cuda/lapack.cu

@@ -7,6 +7,7 @@
 #include <thrust/complex.h>
 
 #include <cassert>
+#include <cstddef>
 
 #include "source_base/module_device/device_check.h"
 
@@ -129,7 +130,8 @@ struct lapack_geqrf_inplace<T, DEVICE_GPU> {
         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<T, DEVICE_GPU> {
 
         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<std::size_t>(j) * lda + j,
+                                      sizeof(T),
+                                      cudaMemcpyDeviceToHost));
+            }
+        }
+
         cuSolverConnector::orgqr(cusolver_handle, m, n, k, d_A, lda, d_tau);
 
         CHECK_CUDA(cudaFree(d_tau));

source/source_base/module_container/ATen/kernels/lapack.cpp

@@ -116,7 +116,8 @@ struct lapack_geqrf_inplace<T, DEVICE_CPU> {
         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<T, DEVICE_CPU> {
             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;

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

source/source_base/module_container/ATen/kernels/test/lapack_test.cpp

@@ -1,5 +1,7 @@
 #include <gtest/gtest.h>
 
+#include <vector>
+
 #include <ATen/core/tensor.h>
 #include <ATen/kernels/lapack.h>
 #include <base/utils/gtest.h>
@@ -134,9 +136,10 @@ TYPED_TEST(LapackTest, GeqrfInPlace) {
     }).to_device<Device>());
 
     Tensor A = A_input; // will be overwritten as Q
+    std::vector<Type> rdiag(n);
 
-    // do geqrf -> get orthogonal Q
-    geqrfCalculator(m, n, A.data<Type>(), lda);
+    // do geqrf -> get orthogonal Q and the diagonal of R
+    geqrfCalculator(m, n, A.data<Type>(), lda, rdiag.data());
 
     // check on CPU
     Tensor Q = A.to_device<DEVICE_CPU>();
@@ -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<Type>(0.0), static_cast<Type>(0.0), static_cast<Type>(0.0),
+        static_cast<Type>(0.0), static_cast<Type>(0.0), static_cast<Type>(0.0),
+        static_cast<Type>(0.0), static_cast<Type>(0.0), static_cast<Type>(0.0)
+    }).to_device<DEVICE_CPU>());
+    Tensor A_cpu = A_input.to_device<DEVICE_CPU>();
+    gemm('C', 'N',
+         n, n, m,
+         &alpha,
+         Q_data, lda,
+         A_cpu.data<Type>(), lda,
+         &beta,
+         R.data<Type>(), n);
+    for (int j = 0; j < n; ++j) {
+        EXPECT_NEAR(std::abs(R.data<Type>()[j + j * n]), std::abs(rdiag[j]), 1e-5)
+            << "R diagonal mismatch at " << j;
+    }
 }
 
 // Test for lapack_heevd and lapack_heevx:

source/source_estate/elecstate_pw.cpp

@@ -1,5 +1,8 @@
 #include "elecstate_pw.h"
 
+#include <type_traits>
+#include <vector>
+
 #include "source_base/constants.h"
 #include "source_base/libm/libm.h"
 #include "source_base/module_device/device.h"
@@ -139,13 +142,35 @@ void ElecStatePW<T, Device>::psiToRho(const psi::Psi<T, Device>& psi)
         }
     }
 
+    // Keep the FFTs in single precision, but sum k-point densities in double.
+    const bool sum_k_in_double = std::is_same<Real, float>::value
+                                 && std::is_same<Device, base_device::DEVICE_GPU>::value;
+    std::vector<double> 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++)
         {

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 <typename T, typename Device,
+              typename Real = typename GetTypeReal<T>::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<double>* rhog,
         ModuleBase::Vector3<double>* gdr,