From 791bf2abef25ba7112d6003f37b0a8b56190fcc6 Mon Sep 17 00:00:00 2001 From: dzzz2001 Date: Wed, 20 May 2026 22:27:22 +0800 Subject: [PATCH] fix(gint): use fp64 accumulators in mixed-precision paths to avoid precision loss Across CPU and GPU gint paths, accumulator buffers (hr_gint, phi_dm, rho, and the vbatched GEMM C output) are now always allocated as double, even when the input phi/dm/vr_eff are fp32. Multiplies stay in fp32 (cheap), but per-block and global reductions are widened to fp64 so that summing many atom-pair contributions into the same element does not drift. Co-Authored-By: Claude Opus 4.7 (1M context) --- source/source_lcao/module_gint/gint_rho.cpp | 72 ++---------- source/source_lcao/module_gint/gint_rho.h | 10 +- .../source_lcao/module_gint/gint_rho_gpu.cpp | 26 ++--- source/source_lcao/module_gint/gint_vl.cpp | 30 +---- source/source_lcao/module_gint/gint_vl.h | 3 - .../source_lcao/module_gint/gint_vl_gpu.cpp | 39 +++---- .../module_gint/kernel/dgemm_vbatch.cu | 8 +- .../module_gint/kernel/dgemm_vbatch.h | 12 +- .../module_gint/kernel/gemm_nn_vbatch.cuh | 13 ++- .../module_gint/kernel/gemm_tn_vbatch.cuh | 13 ++- .../module_gint/kernel/phi_operator_gpu.cu | 16 +-- .../module_gint/kernel/phi_operator_gpu.h | 22 ++-- .../module_gint/kernel/phi_operator_kernel.cu | 30 ++--- .../kernel/phi_operator_kernel.cuh | 13 ++- source/source_lcao/module_gint/phi_operator.h | 34 ++++-- .../source_lcao/module_gint/phi_operator.hpp | 103 ++++++++++++++---- .../module_hcontainer/base_matrix.h | 14 +++ 17 files changed, 231 insertions(+), 227 deletions(-) diff --git a/source/source_lcao/module_gint/gint_rho.cpp b/source/source_lcao/module_gint/gint_rho.cpp index 634659bc88b..c28e69a8b81 100644 --- a/source/source_lcao/module_gint/gint_rho.cpp +++ b/source/source_lcao/module_gint/gint_rho.cpp @@ -1,6 +1,3 @@ -#include -#include - #include "source_base/global_function.h" #include "gint_rho.h" #include "gint_common.h" @@ -30,15 +27,14 @@ template void Gint_rho::cal_gint_impl_() { std::vector> dm_gint_vec = init_dm_gint_(); - std::vector> rho_cache(nspin_); - std::vector rho_data(nspin_); + // rho_[is] is always double; phi_dot_phi accumulates into it directly. + std::vector rho_data(nspin_); for (int is = 0; is < nspin_; ++is) { - rho_data[is] = get_rho_data_(is, rho_cache); + rho_data[is] = rho_[is]; } dm_2d_to_gint(*gint_info_, dm_vec_, dm_gint_vec); cal_rho_(dm_gint_vec, rho_data); - transfer_rho_cache_(rho_cache); } template @@ -52,64 +48,19 @@ std::vector> Gint_rho::init_dm_gint_() const return dm_gint_vec; } -// Overloaded helpers (C++11-compatible alternative to if constexpr). -// The double overload is preferred by overload resolution when Real=double. - -inline double* get_rho_data(double* const* rho, int is, int /*local_mgrid_num*/, - std::vector>& /*rho_cache*/) -{ - return rho[is]; -} - -template -Real* get_rho_data(double* const* rho, int is, int local_mgrid_num, - std::vector>& rho_cache) -{ - rho_cache[is].resize(local_mgrid_num); - std::transform(rho[is], rho[is] + local_mgrid_num, rho_cache[is].begin(), [](const double value) { - return static_cast(value); - }); - return rho_cache[is].data(); -} - -inline void transfer_rho_back(double* const* /*rho*/, int /*nspin*/, int /*local_mgrid_num*/, - const std::vector>& /*rho_cache*/) -{ - // Nothing to do: double rho was written directly. -} - -template -void transfer_rho_back(double* const* rho, int nspin, int local_mgrid_num, - const std::vector>& rho_cache) -{ - for (int is = 0; is < nspin; ++is) - { - for (int ir = 0; ir < local_mgrid_num; ++ir) - { - rho[is][ir] = static_cast(rho_cache[is][ir]); - } - } -} - - - -template -Real* Gint_rho::get_rho_data_(int is, std::vector>& rho_cache) const -{ - return get_rho_data( - rho_, is, gint_info_->get_local_mgrid_num(), rho_cache); -} - template void Gint_rho::cal_rho_( const std::vector>& dm_gint_vec, - const std::vector& rho_data) const + const std::vector& rho_data) const { #pragma omp parallel { PhiOperator phi_op; std::vector phi; - std::vector phi_dm; + // phi_dm is always double: phi_mul_dm writes the cast-to-double result + // into it, and phi_dot_phi reads it as fp64 (so the rho reduction's + // right-hand side is uniformly fp64 even on the fp32 path). + std::vector phi_dm; #pragma omp for schedule(dynamic) for (int i = 0; i < gint_info_->get_bgrids_num(); i++) { @@ -132,11 +83,4 @@ void Gint_rho::cal_rho_( } } -template -void Gint_rho::transfer_rho_cache_(const std::vector>& rho_cache) const -{ - transfer_rho_back( - rho_, nspin_, gint_info_->get_local_mgrid_num(), rho_cache); -} - } diff --git a/source/source_lcao/module_gint/gint_rho.h b/source/source_lcao/module_gint/gint_rho.h index b75a1220958..5bcede5f19f 100644 --- a/source/source_lcao/module_gint/gint_rho.h +++ b/source/source_lcao/module_gint/gint_rho.h @@ -28,16 +28,12 @@ class Gint_rho : public Gint template std::vector> init_dm_gint_() const; - template - Real* get_rho_data_(int is, std::vector>& rho_cache) const; - + // rho is always accumulated in double (see phi_dot_phi). When Real=float, + // only phi and phi_dm are fp32; the per-meshgrid reduction is fp64. template void cal_rho_( const std::vector>& dm_gint_vec, - const std::vector& rho_data) const; - - template - void transfer_rho_cache_(const std::vector>& rho_cache) const; + const std::vector& rho_data) const; // input const std::vector*> dm_vec_; diff --git a/source/source_lcao/module_gint/gint_rho_gpu.cpp b/source/source_lcao/module_gint/gint_rho_gpu.cpp index dfa6fcc8606..885c2453879 100644 --- a/source/source_lcao/module_gint/gint_rho_gpu.cpp +++ b/source/source_lcao/module_gint/gint_rho_gpu.cpp @@ -5,8 +5,6 @@ #include "kernel/phi_operator_gpu.h" #include "source_base/module_device/device_check.h" -#include - namespace ModuleGint { @@ -40,13 +38,14 @@ void Gint_rho_gpu::cal_gint_impl_() // 2. Transfer dm from 2D parallel distribution to gint serial distribution dm_2d_to_gint(*gint_info_, dm_vec_, dm_gint_vec); - // 3. Transfer dm to GPU + // 3. Transfer dm to GPU. rho_d is always double — the kernel accumulates + // in fp64 regardless of the input precision. std::vector> dm_gint_d_vec(nspin_); - std::vector> rho_d_vec(nspin_); + std::vector> rho_d_vec(nspin_); for (int is = 0; is < nspin_; is++) { dm_gint_d_vec[is] = CudaMemWrapper(dm_gint_vec[is].get_nnr(), 0, false); - rho_d_vec[is] = CudaMemWrapper(gint_info_->get_local_mgrid_num(), 0, false); + rho_d_vec[is] = CudaMemWrapper(gint_info_->get_local_mgrid_num(), 0, false); CHECK_CUDA(cudaMemcpy(dm_gint_d_vec[is].get_device_ptr(), dm_gint_vec[is].get_wrapper(), dm_gint_vec[is].get_nnr() * sizeof(Real), cudaMemcpyHostToDevice)); } @@ -61,7 +60,9 @@ void Gint_rho_gpu::cal_gint_impl_() CHECK_CUDA(cudaStreamCreate(&stream)); PhiOperatorGpu phi_op(gint_info_->get_gpu_vars(), stream); CudaMemWrapper phi(BatchBigGrid::get_max_phi_len(), stream, false); - CudaMemWrapper phi_dm(BatchBigGrid::get_max_phi_len(), stream, false); + // phi_dm is always double: the gemm_nn_vbatch kernel accumulates fp32 + // multiplies into a fp64 register, then atomicAdd's into phi_dm. + CudaMemWrapper phi_dm(BatchBigGrid::get_max_phi_len(), stream, false); #pragma omp for schedule(dynamic) for (int i = 0; i < gint_info_->get_bgrid_batches_num(); ++i) { @@ -83,18 +84,13 @@ void Gint_rho_gpu::cal_gint_impl_() CHECK_CUDA(cudaStreamDestroy(stream)); } - // 5. Transfer rho back to CPU and convert to double if needed + // 5. Transfer rho back to CPU (already double — copy straight into rho_[is]) const int local_mgrid_num = gint_info_->get_local_mgrid_num(); for (int is = 0; is < nspin_; is++) { - std::vector rho_tmp(local_mgrid_num); - CHECK_CUDA(cudaMemcpy(rho_tmp.data(), rho_d_vec[is].get_device_ptr(), - local_mgrid_num * sizeof(Real), cudaMemcpyDeviceToHost)); - for (int ir = 0; ir < local_mgrid_num; ++ir) - { - rho_[is][ir] = static_cast(rho_tmp[ir]); - } + CHECK_CUDA(cudaMemcpy(rho_[is], rho_d_vec[is].get_device_ptr(), + local_mgrid_num * sizeof(double), cudaMemcpyDeviceToHost)); } } -} // namespace ModuleGint \ No newline at end of file +} // namespace ModuleGint diff --git a/source/source_lcao/module_gint/gint_vl.cpp b/source/source_lcao/module_gint/gint_vl.cpp index 410eeb65885..73619353805 100644 --- a/source/source_lcao/module_gint/gint_vl.cpp +++ b/source/source_lcao/module_gint/gint_vl.cpp @@ -1,5 +1,4 @@ #include -#include #include "gint_common.h" #include "gint_vl.h" @@ -30,12 +29,6 @@ void Gint_vl::cal_gint() // Private functions //======================== -template -HContainer Gint_vl::init_hr_gint_() const -{ - return gint_info_->get_hr(); -} - // Overloaded helpers (C++11-compatible alternative to if constexpr). // The double overload is preferred by overload resolution when Real=double; // the template overload handles all other types (e.g. float). @@ -57,24 +50,13 @@ const Real* get_vr_eff_data_(const double* vr_eff, int local_mgrid_num, return vr_eff_buffer.data(); } -inline void finalize_hr_gint_(HContainer& hr_gint, HContainer* hR) -{ - compose_hr_gint(hr_gint); - hr_gint_to_hR(hr_gint, *hR); -} - -template -void finalize_hr_gint_(HContainer& hr_gint, HContainer* hR) -{ - HContainer hr_gint_dp = make_cast_hcontainer(hr_gint); - compose_hr_gint(hr_gint_dp); - hr_gint_to_hR(hr_gint_dp, *hR); -} - template void Gint_vl::cal_gint_impl_() { - HContainer hr_gint = init_hr_gint_(); + // hr_gint is always allocated as HContainer: when Real=float, the + // fp32 multiplies feed into a fp64 accumulator inside phi_mul_phi to avoid + // catastrophic precision loss in the global reduction. + HContainer hr_gint = gint_info_->get_hr(); std::vector vr_eff_buffer; const Real* vr_eff = get_vr_eff_data_( vr_eff_, gint_info_->get_local_mgrid_num(), vr_eff_buffer); @@ -102,8 +84,8 @@ void Gint_vl::cal_gint_impl_() } } - finalize_hr_gint_(hr_gint, hR_); + compose_hr_gint(hr_gint); + hr_gint_to_hR(hr_gint, *hR_); } } // namespace ModuleGint - diff --git a/source/source_lcao/module_gint/gint_vl.h b/source/source_lcao/module_gint/gint_vl.h index eede31aa8b5..cd21f178c01 100644 --- a/source/source_lcao/module_gint/gint_vl.h +++ b/source/source_lcao/module_gint/gint_vl.h @@ -24,9 +24,6 @@ class Gint_vl : public Gint template void cal_gint_impl_(); - template - HContainer init_hr_gint_() const; - // input const double* vr_eff_ = nullptr; diff --git a/source/source_lcao/module_gint/gint_vl_gpu.cpp b/source/source_lcao/module_gint/gint_vl_gpu.cpp index ef2643f4a28..66822b77e25 100644 --- a/source/source_lcao/module_gint/gint_vl_gpu.cpp +++ b/source/source_lcao/module_gint/gint_vl_gpu.cpp @@ -28,32 +28,18 @@ void Gint_vl_gpu::cal_gint() ModuleBase::timer::end("Gint", "cal_gint_vl"); } -// Helper: finalize hr_gint (double path — no cast needed) -inline void finalize_hr_gint_gpu_(HContainer& hr_gint, HContainer* hR) -{ - compose_hr_gint(hr_gint); - hr_gint_to_hR(hr_gint, *hR); -} - -// Helper: finalize hr_gint (non-double path — cast to double first) -template -void finalize_hr_gint_gpu_(HContainer& hr_gint, HContainer* hR) -{ - HContainer hr_gint_dp = make_cast_hcontainer(hr_gint); - compose_hr_gint(hr_gint_dp); - hr_gint_to_hR(hr_gint_dp, *hR); -} - template void Gint_vl_gpu::cal_gint_impl_() { - // 1. Initialize hr_gint as HContainer - HContainer hr_gint = gint_info_->get_hr(); + // hr_gint is always allocated as HContainer: the per-atom-pair GEMM + // accumulates fp32 multiplies into a fp64 register/atomicAdd inside the + // kernel so that the global reduction across many biggrids stays accurate. + HContainer hr_gint = gint_info_->get_hr(); - // 2. Convert vr_eff to Real and transfer to GPU + // 1. Convert vr_eff to Real and transfer to GPU const int local_mgrid_num = gint_info_->get_local_mgrid_num(); CudaMemWrapper vr_eff_d(local_mgrid_num, 0, false); - CudaMemWrapper hr_gint_d(hr_gint.get_nnr(), 0, false); + CudaMemWrapper hr_gint_d(hr_gint.get_nnr(), 0, false); if (std::is_same::value) { @@ -71,7 +57,7 @@ void Gint_vl_gpu::cal_gint_impl_() local_mgrid_num * sizeof(Real), cudaMemcpyHostToDevice)); } - // 3. Calculate hr_gint on GPU + // 2. Calculate hr_gint on GPU #pragma omp parallel num_threads(gint_info_->get_streams_num()) { // 20240620 Note that it must be set again here because @@ -101,12 +87,13 @@ void Gint_vl_gpu::cal_gint_impl_() CHECK_CUDA(cudaStreamDestroy(stream)); } - // 4. Transfer hr_gint back to CPU + // 3. Transfer hr_gint back to CPU CHECK_CUDA(cudaMemcpy(hr_gint.get_wrapper(), hr_gint_d.get_device_ptr(), - hr_gint.get_nnr() * sizeof(Real), cudaMemcpyDeviceToHost)); + hr_gint.get_nnr() * sizeof(double), cudaMemcpyDeviceToHost)); - // 5. Compose and transfer to hR (with cast if needed) - finalize_hr_gint_gpu_(hr_gint, hR_); + // 4. Compose and transfer to hR (already double, no cast needed) + compose_hr_gint(hr_gint); + hr_gint_to_hR(hr_gint, *hR_); } -} \ No newline at end of file +} diff --git a/source/source_lcao/module_gint/kernel/dgemm_vbatch.cu b/source/source_lcao/module_gint/kernel/dgemm_vbatch.cu index b448f8d5a01..38946d51492 100644 --- a/source/source_lcao/module_gint/kernel/dgemm_vbatch.cu +++ b/source/source_lcao/module_gint/kernel/dgemm_vbatch.cu @@ -9,7 +9,7 @@ void gemm_nn_vbatch( const int* m_d, const int* n_d, const int* k_d, const T* const* A_array_d, const int* lda_d, const T* const* B_array_d, const int* ldb_d, - T** C_array_d, const int* ldc_d, + double** C_array_d, const int* ldc_d, int batchCount, cudaStream_t stream, const T* alpha) { @@ -28,7 +28,7 @@ void gemm_tn_vbatch( const int* m_d, const int* n_d, const int* k_d, const T* const* A_array_d, const int* lda_d, const T* const* B_array_d, const int* ldb_d, - T** C_array_d, const int* ldc_d, + double** C_array_d, const int* ldc_d, int batchCount, cudaStream_t stream, const T* alpha) { @@ -49,7 +49,7 @@ template void gemm_nn_vbatch( template void gemm_nn_vbatch( int, int, int, const int*, const int*, const int*, const float* const*, const int*, const float* const*, const int*, - float**, const int*, int, cudaStream_t, const float*); + double**, const int*, int, cudaStream_t, const float*); template void gemm_tn_vbatch( int, int, int, const int*, const int*, const int*, @@ -59,4 +59,4 @@ template void gemm_tn_vbatch( template void gemm_tn_vbatch( int, int, int, const int*, const int*, const int*, const float* const*, const int*, const float* const*, const int*, - float**, const int*, int, cudaStream_t, const float*); + double**, const int*, int, cudaStream_t, const float*); diff --git a/source/source_lcao/module_gint/kernel/dgemm_vbatch.h b/source/source_lcao/module_gint/kernel/dgemm_vbatch.h index 76835743561..68505c1f4e4 100644 --- a/source/source_lcao/module_gint/kernel/dgemm_vbatch.h +++ b/source/source_lcao/module_gint/kernel/dgemm_vbatch.h @@ -3,24 +3,30 @@ #include // Template version: C(batch_id) = alpha * A(batch_id) * B(batch_id) + C(batch_id) +// As with gemm_tn_vbatch, the C accumulator is always double regardless of the +// input type T so the per-block reduction and device-side atomicAdd run in fp64. template void gemm_nn_vbatch( int max_m, int max_n, int max_k, const int* m_d, const int* n_d, const int* k_d, const T* const* A_array_d, const int* lda_d, const T* const* B_array_d, const int* ldb_d, - T** C_array_d, const int* ldc_d, + double** C_array_d, const int* ldc_d, int batchCount, cudaStream_t stream, const T* alpha = nullptr); // Template version: C(batch_id) = alpha * A(batch_id)^T * B(batch_id) + C(batch_id) +// The C accumulator is always double regardless of input type T: a fp32 GEMM +// path (T=float) feeds fp32 multiplies into fp64 accumulators (registers and +// device-side atomicAdds) to avoid catastrophic precision loss across many +// atom-pair contributions to the same hr_gint element. template void gemm_tn_vbatch( int max_m, int max_n, int max_k, const int* m_d, const int* n_d, const int* k_d, const T* const* A_array_d, const int* lda_d, const T* const* B_array_d, const int* ldb_d, - T** C_array_d, const int* ldc_d, + double** C_array_d, const int* ldc_d, int batchCount, cudaStream_t stream, const T* alpha = nullptr); @@ -48,6 +54,8 @@ inline void dgemm_tn_vbatch( int batchCount, cudaStream_t stream, const double* alpha = nullptr) { + // T=double path: A, B, and C are all double — the C-channel double-fix + // matches the legacy signature here. gemm_tn_vbatch(max_m, max_n, max_k, m_d, n_d, k_d, A_array_d, lda_d, B_array_d, ldb_d, C_array_d, ldc_d, batchCount, stream, alpha); diff --git a/source/source_lcao/module_gint/kernel/gemm_nn_vbatch.cuh b/source/source_lcao/module_gint/kernel/gemm_nn_vbatch.cuh index cac7abfd346..39e10e4252d 100644 --- a/source/source_lcao/module_gint/kernel/gemm_nn_vbatch.cuh +++ b/source/source_lcao/module_gint/kernel/gemm_nn_vbatch.cuh @@ -35,7 +35,7 @@ static __device__ void vbatched_gemm_nn_device(int M, int LDA, const T* __restrict__ B, int LDB, - T* __restrict__ C, + double* __restrict__ C, int LDC, T* sA, int slda, @@ -57,7 +57,8 @@ static __device__ void vbatched_gemm_nn_device(int M, int blx = blockIdx.x; // block's m dimension int bly = blockIdx.y; // block's n dimension - // Registers for the innermost loop + // Registers for the innermost loop. rC accumulates in T; the widening to + // double happens only at the final atomicAdd into C. T rC[THR_N][THR_M]; T rA[THR_M]; T rB[THR_N]; @@ -85,7 +86,7 @@ static __device__ void vbatched_gemm_nn_device(int M, #pragma unroll for (m = 0; m < THR_M; m++) { - rC[n][m] = 0.0; + rC[n][m] = T(0); } } @@ -245,7 +246,7 @@ static __device__ void vbatched_gemm_nn_device(int M, { int offsC = coord_dCn * LDC + coord_dCm; - atomicAdd(C + offsC, rC[n][m] * alpha); + atomicAdd(C + offsC, static_cast(rC[n][m] * alpha)); } } } @@ -269,7 +270,7 @@ static __global__ void vbatched_gemm_nn_kernel(const int* M, const int* global_lda, const T* const* global_B_array, const int* global_ldb, - T** global_C_array, + double** global_C_array, const int* global_ldc, const T* alpha) { @@ -378,7 +379,7 @@ void vbatched_gemm_nn_impl(int max_m, const int* global_lda, const T* const* global_B_array, const int* global_ldb, - T** global_C_array, + double** global_C_array, const int* global_ldc, int batchCount, cudaStream_t stream, diff --git a/source/source_lcao/module_gint/kernel/gemm_tn_vbatch.cuh b/source/source_lcao/module_gint/kernel/gemm_tn_vbatch.cuh index 380eb9efa84..ae6a4a00551 100644 --- a/source/source_lcao/module_gint/kernel/gemm_tn_vbatch.cuh +++ b/source/source_lcao/module_gint/kernel/gemm_tn_vbatch.cuh @@ -35,7 +35,7 @@ static __device__ void vbatched_gemm_nt_device(int M, int LDA, const T* __restrict__ B, int LDB, - T* __restrict__ C, + double* __restrict__ C, int LDC, T* sA, int slda, @@ -57,7 +57,8 @@ static __device__ void vbatched_gemm_nt_device(int M, int blx = blockIdx.x; // block's m dimension int bly = blockIdx.y; // block's n dimension - // Registers for the innermost loop + // Registers for the innermost loop. rC accumulates in T; the widening to + // double happens only at the final atomicAdd into C. T rC[THR_N][THR_M]; T rA[THR_M]; T rB[THR_N]; @@ -85,7 +86,7 @@ static __device__ void vbatched_gemm_nt_device(int M, #pragma unroll for (m = 0; m < THR_M; m++) { - rC[n][m] = 0.0; + rC[n][m] = T(0); } } @@ -245,7 +246,7 @@ static __device__ void vbatched_gemm_nt_device(int M, { int offsC = coord_dCn * LDC + coord_dCm; - atomicAdd(C + offsC, rC[n][m] * alpha); + atomicAdd(C + offsC, static_cast(rC[n][m] * alpha)); } } } @@ -269,7 +270,7 @@ static __global__ void vbatched_gemm_nt_kernel(const int* M, const int* global_lda, const T* const* global_B_array, const int* global_ldb, - T** global_C_array, + double** global_C_array, const int* global_ldc, const T* alpha) { @@ -403,7 +404,7 @@ void vbatched_gemm_tn_impl(int max_m, const int* global_lda, const T* const* global_B_array, const int* global_ldb, - T** global_C_array, + double** global_C_array, const int* global_ldc, int batchCount, cudaStream_t stream, diff --git a/source/source_lcao/module_gint/kernel/phi_operator_gpu.cu b/source/source_lcao/module_gint/kernel/phi_operator_gpu.cu index 29e27a852a7..24cea8c8565 100644 --- a/source/source_lcao/module_gint/kernel/phi_operator_gpu.cu +++ b/source/source_lcao/module_gint/kernel/phi_operator_gpu.cu @@ -219,8 +219,8 @@ template void PhiOperatorGpu::phi_mul_phi( const Real* phi_d, const Real* phi_vldr3_d, - HContainer& hRGint, - Real* hr_d) const + HContainer& hRGint, + double* hr_d) const { // ap_num means number of atom pairs int ap_num = 0; @@ -285,7 +285,7 @@ void PhiOperatorGpu::phi_mul_phi( gemm_n_.copy_host_to_device_async(ap_num); gemm_k_.copy_host_to_device_async(ap_num); CHECK_CUDA(cudaEventRecord(event_, stream_)); - + gemm_tn_vbatch(max_m, max_n, max_k, @@ -309,9 +309,9 @@ void PhiOperatorGpu::phi_mul_dm( const Real* dm_d, const HContainer& dm, const bool is_symm, - Real* phi_dm_d) + double* phi_dm_d) { - CHECK_CUDA(cudaMemsetAsync(phi_dm_d, 0, phi_len_ * sizeof(Real), stream_)); + CHECK_CUDA(cudaMemsetAsync(phi_dm_d, 0, phi_len_ * sizeof(double), stream_)); // ap_num means number of atom pairs int ap_num = 0; int max_m = mgrids_num_; @@ -401,12 +401,12 @@ void PhiOperatorGpu::phi_mul_dm( template void PhiOperatorGpu::phi_dot_phi( const Real* phi_i_d, - const Real* phi_j_d, - Real* rho_d) const + const double* phi_j_d, + double* rho_d) const { dim3 grid_dim(mgrids_num_, bgrid_batch_->get_batch_size()); dim3 threads_per_block(64); - phi_dot_phi_kernel<<>>( + phi_dot_phi_kernel<<>>( phi_i_d, phi_j_d, mgrids_num_, diff --git a/source/source_lcao/module_gint/kernel/phi_operator_gpu.h b/source/source_lcao/module_gint/kernel/phi_operator_gpu.h index ed440f499d6..fbdbc95352a 100644 --- a/source/source_lcao/module_gint/kernel/phi_operator_gpu.h +++ b/source/source_lcao/module_gint/kernel/phi_operator_gpu.h @@ -33,23 +33,28 @@ class PhiOperatorGpu const Real* phi_d, Real* result_d) const; + // All GEMM accumulators (hr in phi_mul_phi, phi_dm in phi_mul_dm) are + // double-typed regardless of Real: when Real=float the multiplies stay in + // fp32 (cheap) but per-block reductions and device-side atomicAdd run in + // fp64 so the global reductions don't drift. void phi_mul_phi( const Real* phi_d, const Real* phi_vldr3_d, - HContainer& hRGint, - Real* hr_d) const; - + HContainer& hRGint, + double* hr_d) const; + void phi_mul_dm( const Real* phi_d, const Real* dm_d, const HContainer& dm, const bool is_symm, - Real* phi_dm_d); + double* phi_dm_d); + // phi_j_d is the output of phi_mul_dm and therefore always double. void phi_dot_phi( const Real* phi_i_d, - const Real* phi_j_d, - Real* rho_d) const; + const double* phi_j_d, + double* rho_d) const; // These remain double-only (for force/stress paths) void phi_dot_dphi( @@ -105,7 +110,10 @@ class PhiOperatorGpu mutable CudaMemWrapper gemm_ldc_; mutable CudaMemWrapper gemm_A_; mutable CudaMemWrapper gemm_B_; - mutable CudaMemWrapper gemm_C_; + // Single C-pointer buffer: both phi_mul_phi (output hr) and phi_mul_dm + // (output phi_dm) write into double* accumulators, so a single shared + // gemm_C_ device buffer can serve both call sites. + mutable CudaMemWrapper gemm_C_; mutable CudaMemWrapper gemm_alpha_; }; diff --git a/source/source_lcao/module_gint/kernel/phi_operator_kernel.cu b/source/source_lcao/module_gint/kernel/phi_operator_kernel.cu index aa2dcc9fb7d..f8caa4307d1 100644 --- a/source/source_lcao/module_gint/kernel/phi_operator_kernel.cu +++ b/source/source_lcao/module_gint/kernel/phi_operator_kernel.cu @@ -388,30 +388,32 @@ template __global__ void phi_mul_vldr3_kernel( const int*, const int*, const int*, float*); // rho(ir) = \sum_{iwt} \phi_i(ir,iwt) * \phi_j^*(ir,iwt) -// each block calculate the dot product of phi_i and phi_j of a meshgrid -template +// each block calculate the dot product of phi_i and phi_j of a meshgrid. +// The per-thread/warp/block reduction is in double regardless of input types +// so that fp32 inputs are summed without catastrophic precision loss. +template __global__ void phi_dot_phi_kernel( - const Real* __restrict__ phi_i, - const Real* __restrict__ phi_j, + const Tin_a* __restrict__ phi_i, + const Tin_b* __restrict__ phi_j, const int mgrids_per_bgrid, const int* __restrict__ mgrid_lidx, const int* __restrict__ bgrid_phi_len, const int* __restrict__ bgrid_phi_start, - Real* __restrict__ rho) + double* __restrict__ rho) { - __shared__ Real s_data[32]; // the length of s_data equals the max warp num of a block + __shared__ double s_data[32]; // the length of s_data equals the max warp num of a block const int bgrid_id = blockIdx.y; const int mgrid_id = blockIdx.x; const int phi_len = bgrid_phi_len[bgrid_id]; const int phi_start = bgrid_phi_start[bgrid_id] + mgrid_id * phi_len; - const Real* phi_i_mgrid = phi_i + phi_start; - const Real* phi_j_mgrid = phi_j + phi_start; + const Tin_a* phi_i_mgrid = phi_i + phi_start; + const Tin_b* phi_j_mgrid = phi_j + phi_start; const int batch_mgrid_id = bgrid_id * mgrids_per_bgrid + mgrid_id; const int mgrid_local_idx = mgrid_lidx[batch_mgrid_id]; const int tid = threadIdx.x; const int warp_id = tid / 32; const int lane_id = tid % 32; - Real tmp_sum = Real(0.0); + double tmp_sum = 0.0; for (int i = tid; i < phi_len; i += blockDim.x) { @@ -426,7 +428,7 @@ __global__ void phi_dot_phi_kernel( } __syncthreads(); - tmp_sum = (tid < blockDim.x / 32) ? s_data[tid] : Real(0.0); + tmp_sum = (tid < blockDim.x / 32) ? s_data[tid] : 0.0; if(warp_id == 0) { tmp_sum = warpReduceSum(tmp_sum); @@ -439,12 +441,12 @@ __global__ void phi_dot_phi_kernel( } // Explicit instantiations for phi_dot_phi_kernel -template __global__ void phi_dot_phi_kernel( +template __global__ void phi_dot_phi_kernel( const double*, const double*, const int, const int*, const int*, const int*, double*); -template __global__ void phi_dot_phi_kernel( - const float*, const float*, const int, - const int*, const int*, const int*, float*); +template __global__ void phi_dot_phi_kernel( + const float*, const double*, const int, + const int*, const int*, const int*, double*); __global__ void phi_dot_dphi_kernel( const double* __restrict__ phi, diff --git a/source/source_lcao/module_gint/kernel/phi_operator_kernel.cuh b/source/source_lcao/module_gint/kernel/phi_operator_kernel.cuh index 9e51ae97c19..23ca51c2810 100644 --- a/source/source_lcao/module_gint/kernel/phi_operator_kernel.cuh +++ b/source/source_lcao/module_gint/kernel/phi_operator_kernel.cuh @@ -92,16 +92,19 @@ __global__ void phi_mul_vldr3_kernel( Real* __restrict__ result); // rho(ir) = \sum_{iwt} \phi_i(ir,iwt) * \phi_j^*(ir,iwt) -// each block calculate the dot product of phi_i and phi_j of a meshgrid -template +// each block calculate the dot product of phi_i and phi_j of a meshgrid. +// Inputs phi_i and phi_j can have different element types: in the rho path +// phi_i is fp32 (Real) while phi_j (phi_dm) is fp64; in the tau path both are +// fp64. The per-block reduction and atomicAdd to rho run in fp64 regardless. +template __global__ void phi_dot_phi_kernel( - const Real* __restrict__ phi_i, // phi_i(ir,iwt) - const Real* __restrict__ phi_j, // phi_j(ir,iwt) + const Tin_a* __restrict__ phi_i, // phi_i(ir,iwt) + const Tin_b* __restrict__ phi_j, // phi_j(ir,iwt) const int mgrids_per_bgrid, // the number of mgrids of each biggrid const int* __restrict__ mgrid_lidx, // the idx of mgrid in local cell const int* __restrict__ bgrid_phi_len, // the length of phi on a mgrid of a biggrid const int* __restrict__ bgrid_phi_start, // the start idx in phi of each biggrid - Real* __restrict__ rho); // rho(ir) + double* __restrict__ rho); // rho(ir) __global__ void phi_dot_dphi_kernel( const double* __restrict__ phi, diff --git a/source/source_lcao/module_gint/phi_operator.h b/source/source_lcao/module_gint/phi_operator.h index 83df19e9d26..f96d5fc35bd 100644 --- a/source/source_lcao/module_gint/phi_operator.h +++ b/source/source_lcao/module_gint/phi_operator.h @@ -48,12 +48,17 @@ class PhiOperator double* ddphi_yy, double* ddphi_yz, double* ddphi_zz) const; // phi_dm(ir,iwt_2) = \sum_{iwt_1} phi(ir,iwt_1) * dm(iwt_1,iwt_2) - template + // The phi_dm output is always double regardless of Tin: this keeps the + // downstream phi_dot_phi reading a uniform fp64 right-hand side. When + // Tin=float the GEMM still runs in fp32 (K is small ~10, so per-call + // accumulation in fp32 is fine); only the final write into phi_dm casts + // to double. + template void phi_mul_dm( - const T*const phi, // phi(ir,iwt) - const HContainer& dm, // dm(iwt_1,iwt_2) + const Tin*const phi, // phi(ir,iwt) + const HContainer& dm, // dm(iwt_1,iwt_2) const bool is_symm, - T*const phi_dm) const; // phi_dm(ir,iwt) + double*const phi_dm) const; // phi_dm(ir,iwt) // result(ir,iwt) = phi(ir,iwt) * vl(ir) template @@ -64,20 +69,25 @@ class PhiOperator T*const result) const; // result(ir,iwt) // hr(iwt_i,iwt_j) = \sum_{ir} phi_i(ir,iwt_i) * phi_i(ir,iwt_j) - // this is a thread-safe function - template + // this is a thread-safe function. + // The accumulator hr is always double: when Tin=float we want fp32 multiplies + // but fp64 accumulation across many biggrids to avoid catastrophic precision loss. + template void phi_mul_phi( - const T*const phi_i, // phi_i(ir,iwt) - const T*const phi_j, // phi_j(ir,iwt) - HContainer& hr, // hr(iwt_i,iwt_j) + const Tin*const phi_i, // phi_i(ir,iwt) + const Tin*const phi_j, // phi_j(ir,iwt) + HContainer& hr, // hr(iwt_i,iwt_j) const TriPart part) const; // rho(ir) = \sum_{iwt} \phi_i(ir,iwt) * \phi_j(ir,iwt) - template + // phi_j is always double-typed (it is the output of phi_mul_dm, which now + // accumulates into a double buffer regardless of Tin). The inner dot + // product is accumulated in double too: fp32 phi_i + fp64 phi_j → fp64. + template void phi_dot_phi( const Tin*const phi_i, // phi_i(ir,iwt) - const Tin*const phi_j, // phi_j(ir,iwt) - Tout*const rho) const; // rho(ir) + const double*const phi_j, // phi_j(ir,iwt) + double*const rho) const; // rho(ir) void phi_dot_dphi( const double* phi, diff --git a/source/source_lcao/module_gint/phi_operator.hpp b/source/source_lcao/module_gint/phi_operator.hpp index fefafb3f8f3..071a775d8ee 100644 --- a/source/source_lcao/module_gint/phi_operator.hpp +++ b/source/source_lcao/module_gint/phi_operator.hpp @@ -18,15 +18,46 @@ void PhiOperator::set_phi(T* phi) const } } +// Helpers to dispatch a Tin-typed BLAS-GEMM target buffer: +// - For Tin=double, write the GEMM result directly into the caller's +// double* phi_dm (no scratch, no cast). +// - For Tin=float, allocate a fp32 scratch buffer and cast the result into +// phi_dm at the end. K of each individual GEMM is small (~atom nw) so +// accumulating in fp32 inside the scratch is fine. +inline double* phi_mul_dm_scratch_(double* phi_dm, std::vector& /*scratch*/, int /*size*/) +{ + return phi_dm; +} + +template +inline Tin* phi_mul_dm_scratch_(double* /*phi_dm*/, std::vector& scratch, int size) +{ + scratch.assign(size, Tin(0)); + return scratch.data(); +} + +inline void phi_mul_dm_finalize_(double* /*phi_dm*/, const std::vector& /*scratch*/, int /*size*/) {} + +template +inline void phi_mul_dm_finalize_(double* phi_dm, const std::vector& scratch, int size) +{ + for (int k = 0; k < size; ++k) + { + phi_dm[k] = static_cast(scratch[k]); + } +} + // phi_dm(ir,iwt_2) = \sum_{iwt_1} phi(ir,iwt_1) * dm(iwt_1,iwt_2) -template +template void PhiOperator::phi_mul_dm( - const T*const phi, // phi(ir,iwt) - const HContainer& dm, // dm(iwt_1,iwt_2) + const Tin*const phi, // phi(ir,iwt) + const HContainer& dm, // dm(iwt_1,iwt_2) const bool is_symm, - T*const phi_dm) const // phi_dm(ir,iwt) + double*const phi_dm) const // phi_dm(ir,iwt) { - ModuleBase::GlobalFunc::ZEROS(phi_dm, rows_ * cols_); + std::vector scratch; + Tin* target = phi_mul_dm_scratch_(phi_dm, scratch, rows_ * cols_); + ModuleBase::GlobalFunc::ZEROS(target, rows_ * cols_); for(int i = 0; i < biggrid_->get_atoms_num(); ++i) { @@ -36,14 +67,14 @@ void PhiOperator::phi_mul_dm( if(is_symm) { const auto dm_mat = dm.find_matrix(atom_i->get_iat(), atom_i->get_iat(), 0, 0, 0); - constexpr T alpha = 1.0; - constexpr T beta = 1.0; + constexpr Tin alpha = 1.0; + constexpr Tin beta = 1.0; BlasConnector::symm_cm( 'L', 'U', atoms_phi_len_[i], rows_, alpha, dm_mat->get_pointer(), atoms_phi_len_[i], &phi[0 * cols_ + atoms_startidx_[i]], cols_, - beta, &phi_dm[0 * cols_ + atoms_startidx_[i]], cols_); + beta, &target[0 * cols_ + atoms_startidx_[i]], cols_); } const int start = is_symm ? i + 1 : 0; @@ -71,16 +102,18 @@ void PhiOperator::phi_mul_dm( continue; } - const T alpha = is_symm ? 2.0 : 1.0; - constexpr T beta = 1.0; + const Tin alpha = is_symm ? 2.0 : 1.0; + constexpr Tin beta = 1.0; BlasConnector::gemm( 'N', 'N', len, atoms_phi_len_[j], atoms_phi_len_[i], alpha, &phi[start_idx * cols_ + atoms_startidx_[i]], cols_, dm_mat->get_pointer(), atoms_phi_len_[j], - beta, &phi_dm[start_idx * cols_ + atoms_startidx_[j]], cols_); + beta, &target[start_idx * cols_ + atoms_startidx_[j]], cols_); } } + + phi_mul_dm_finalize_(phi_dm, scratch, rows_ * cols_); } // result(ir) = phi(ir) * vl(ir) @@ -104,15 +137,17 @@ void PhiOperator::phi_mul_vldr3( } // hr(iwt_i,iwt_j) += \sum_{ir} phi_i(ir,iwt_i) * phi_i(ir,iwt_j) -// this is a thread-safe function -template +// this is a thread-safe function. +// The per-biggrid GEMM accumulator tmp_hr stays in Tin (small K, no significant +// precision loss); only the global add into HContainer is widened. +template void PhiOperator::phi_mul_phi( - const T*const phi_i, // phi_i(ir,iwt) - const T*const phi_j, // phi_j(ir,iwt) - HContainer& hr, // hr(iwt_i,iwt_j) + const Tin*const phi_i, // phi_i(ir,iwt) + const Tin*const phi_j, // phi_j(ir,iwt) + HContainer& hr, // hr(iwt_i,iwt_j) const TriPart part) const { - std::vector tmp_hr; + std::vector tmp_hr; for(int i = 0; i < biggrid_->get_atoms_num(); ++i) { const auto atom_i = biggrid_->get_atom(i); @@ -158,7 +193,7 @@ void PhiOperator::phi_mul_phi( tmp_hr.resize(n_i * n_j); ModuleBase::GlobalFunc::ZEROS(tmp_hr.data(), n_i*n_j); - constexpr T alpha=1, beta=1; + constexpr Tin alpha=1, beta=1; BlasConnector::gemm( 'T', 'N', n_i, n_j, len, alpha, phi_i + start_idx * cols_ + atoms_startidx_[i], cols_, @@ -171,18 +206,38 @@ void PhiOperator::phi_mul_phi( } } +// Mixed-precision dotc wrapper. Accepts (double, double) or (double, float); +// when y is fp32 it is upcast into the caller-provided fp64 scratch buffer +// before dispatching to BlasConnector::dotc (which requires uniform types). +// The buffer is grown on demand and meant to be reused across many calls. +inline double dotc_mixed(int n, const double* x, const double* y, + std::vector& /*buf*/) +{ + return BlasConnector::dotc(n, x, 1, y, 1); +} + +inline double dotc_mixed(int n, const double* x, const float* y, + std::vector& buf) +{ + if (static_cast(buf.size()) < n) { buf.resize(n); } + for (int k = 0; k < n; ++k) { buf[k] = static_cast(y[k]); } + return BlasConnector::dotc(n, x, 1, buf.data(), 1); +} + // rho(ir) = \sum_{iwt} \phi_i(ir,iwt) * \phi_j^*(ir,iwt) -template +// phi_j is always double (output of phi_mul_dm); phi_i may be fp32. dotc_mixed +// keeps the inner product in fp64 in either case. +template void PhiOperator::phi_dot_phi( const Tin*const phi_i, // phi_i(ir,iwt) - const Tin*const phi_j, // phi_j(ir,iwt) - Tout*const rho) const // rho(ir) + const double*const phi_j, // phi_j(ir,iwt) + double*const rho) const // rho(ir) { - constexpr int inc = 1; + std::vector buf; for(int i = 0; i < biggrid_->get_mgrids_num(); ++i) { - rho[mgrid_lidx_[i]] += static_cast( - BlasConnector::dotc(cols_, phi_j + i * cols_, inc, phi_i + i * cols_, inc)); + rho[mgrid_lidx_[i]] += dotc_mixed( + cols_, phi_j + i * cols_, phi_i + i * cols_, buf); } } diff --git a/source/source_lcao/module_hcontainer/base_matrix.h b/source/source_lcao/module_hcontainer/base_matrix.h index 1c540e6d2de..36044ee187f 100644 --- a/source/source_lcao/module_hcontainer/base_matrix.h +++ b/source/source_lcao/module_hcontainer/base_matrix.h @@ -119,6 +119,20 @@ class BaseMatrix } } + // Cross-type variant: cast each element of `array` to T before accumulating. + // Selected by overload resolution only when TSrc != T (the non-template + // overload is preferred when types match). + template + void add_array_ts(const TSrc* array) + { + std::lock_guard lock(mtx); + const int size = nrow_local * ncol_local; + for (int i = 0; i < size; ++i) + { + value_begin[i] += static_cast(array[i]); + } + } + private: bool allocated = false;