Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 8 additions & 64 deletions source/source_lcao/module_gint/gint_rho.cpp
Original file line number Diff line number Diff line change
@@ -1,6 +1,3 @@
#include <algorithm>
#include <type_traits>

#include "source_base/global_function.h"
#include "gint_rho.h"
#include "gint_common.h"
Expand Down Expand Up @@ -30,15 +27,14 @@ template<typename Real>
void Gint_rho::cal_gint_impl_()
{
std::vector<HContainer<Real>> dm_gint_vec = init_dm_gint_<Real>();
std::vector<std::vector<Real>> rho_cache(nspin_);
std::vector<Real*> rho_data(nspin_);
// rho_[is] is always double; phi_dot_phi accumulates into it directly.
std::vector<double*> rho_data(nspin_);
for (int is = 0; is < nspin_; ++is)
{
rho_data[is] = get_rho_data_<Real>(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_<Real>(rho_cache);
}

template<typename Real>
Expand All @@ -52,64 +48,19 @@ std::vector<HContainer<Real>> 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<std::vector<double>>& /*rho_cache*/)
{
return rho[is];
}

template<typename Real>
Real* get_rho_data(double* const* rho, int is, int local_mgrid_num,
std::vector<std::vector<Real>>& 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<Real>(value);
});
return rho_cache[is].data();
}

inline void transfer_rho_back(double* const* /*rho*/, int /*nspin*/, int /*local_mgrid_num*/,
const std::vector<std::vector<double>>& /*rho_cache*/)
{
// Nothing to do: double rho was written directly.
}

template<typename Real>
void transfer_rho_back(double* const* rho, int nspin, int local_mgrid_num,
const std::vector<std::vector<Real>>& rho_cache)
{
for (int is = 0; is < nspin; ++is)
{
for (int ir = 0; ir < local_mgrid_num; ++ir)
{
rho[is][ir] = static_cast<double>(rho_cache[is][ir]);
}
}
}



template<typename Real>
Real* Gint_rho::get_rho_data_(int is, std::vector<std::vector<Real>>& rho_cache) const
{
return get_rho_data(
rho_, is, gint_info_->get_local_mgrid_num(), rho_cache);
}

template<typename Real>
void Gint_rho::cal_rho_(
const std::vector<HContainer<Real>>& dm_gint_vec,
const std::vector<Real*>& rho_data) const
const std::vector<double*>& rho_data) const
{
#pragma omp parallel
{
PhiOperator phi_op;
std::vector<Real> phi;
std::vector<Real> 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<double> phi_dm;
#pragma omp for schedule(dynamic)
for (int i = 0; i < gint_info_->get_bgrids_num(); i++)
{
Expand All @@ -132,11 +83,4 @@ void Gint_rho::cal_rho_(
}
}

template<typename Real>
void Gint_rho::transfer_rho_cache_(const std::vector<std::vector<Real>>& rho_cache) const
{
transfer_rho_back(
rho_, nspin_, gint_info_->get_local_mgrid_num(), rho_cache);
}

}
10 changes: 3 additions & 7 deletions source/source_lcao/module_gint/gint_rho.h
Original file line number Diff line number Diff line change
Expand Up @@ -28,16 +28,12 @@ class Gint_rho : public Gint
template<typename Real>
std::vector<HContainer<Real>> init_dm_gint_() const;

template<typename Real>
Real* get_rho_data_(int is, std::vector<std::vector<Real>>& 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<typename Real>
void cal_rho_(
const std::vector<HContainer<Real>>& dm_gint_vec,
const std::vector<Real*>& rho_data) const;

template<typename Real>
void transfer_rho_cache_(const std::vector<std::vector<Real>>& rho_cache) const;
const std::vector<double*>& rho_data) const;

// input
const std::vector<HContainer<double>*> dm_vec_;
Expand Down
26 changes: 11 additions & 15 deletions source/source_lcao/module_gint/gint_rho_gpu.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,6 @@
#include "kernel/phi_operator_gpu.h"
#include "source_base/module_device/device_check.h"

#include <algorithm>

namespace ModuleGint
{

Expand Down Expand Up @@ -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<CudaMemWrapper<Real>> dm_gint_d_vec(nspin_);
std::vector<CudaMemWrapper<Real>> rho_d_vec(nspin_);
std::vector<CudaMemWrapper<double>> rho_d_vec(nspin_);
for (int is = 0; is < nspin_; is++)
{
dm_gint_d_vec[is] = CudaMemWrapper<Real>(dm_gint_vec[is].get_nnr(), 0, false);
rho_d_vec[is] = CudaMemWrapper<Real>(gint_info_->get_local_mgrid_num(), 0, false);
rho_d_vec[is] = CudaMemWrapper<double>(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));
}
Expand All @@ -61,7 +60,9 @@ void Gint_rho_gpu::cal_gint_impl_()
CHECK_CUDA(cudaStreamCreate(&stream));
PhiOperatorGpu<Real> phi_op(gint_info_->get_gpu_vars(), stream);
CudaMemWrapper<Real> phi(BatchBigGrid::get_max_phi_len(), stream, false);
CudaMemWrapper<Real> 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<double> 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)
{
Expand All @@ -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<Real> 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<double>(rho_tmp[ir]);
}
CHECK_CUDA(cudaMemcpy(rho_[is], rho_d_vec[is].get_device_ptr(),
local_mgrid_num * sizeof(double), cudaMemcpyDeviceToHost));
}
}

} // namespace ModuleGint
} // namespace ModuleGint
30 changes: 6 additions & 24 deletions source/source_lcao/module_gint/gint_vl.cpp
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
#include <algorithm>
#include <type_traits>

#include "gint_common.h"
#include "gint_vl.h"
Expand Down Expand Up @@ -30,12 +29,6 @@ void Gint_vl::cal_gint()
// Private functions
//========================

template<typename Real>
HContainer<Real> Gint_vl::init_hr_gint_() const
{
return gint_info_->get_hr<Real>();
}

// 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).
Expand All @@ -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<double>& hr_gint, HContainer<double>* hR)
{
compose_hr_gint(hr_gint);
hr_gint_to_hR(hr_gint, *hR);
}

template<typename Real>
void finalize_hr_gint_(HContainer<Real>& hr_gint, HContainer<double>* hR)
{
HContainer<double> hr_gint_dp = make_cast_hcontainer<double>(hr_gint);
compose_hr_gint(hr_gint_dp);
hr_gint_to_hR(hr_gint_dp, *hR);
}

template<typename Real>
void Gint_vl::cal_gint_impl_()
{
HContainer<Real> hr_gint = init_hr_gint_<Real>();
// hr_gint is always allocated as HContainer<double>: 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<double> hr_gint = gint_info_->get_hr<double>();
std::vector<Real> vr_eff_buffer;
const Real* vr_eff = get_vr_eff_data_(
vr_eff_, gint_info_->get_local_mgrid_num(), vr_eff_buffer);
Expand Down Expand Up @@ -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

3 changes: 0 additions & 3 deletions source/source_lcao/module_gint/gint_vl.h
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,6 @@ class Gint_vl : public Gint
template<typename Real>
void cal_gint_impl_();

template<typename Real>
HContainer<Real> init_hr_gint_() const;

// input
const double* vr_eff_ = nullptr;

Expand Down
39 changes: 13 additions & 26 deletions source/source_lcao/module_gint/gint_vl_gpu.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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<double>& hr_gint, HContainer<double>* 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<typename Real>
void finalize_hr_gint_gpu_(HContainer<Real>& hr_gint, HContainer<double>* hR)
{
HContainer<double> hr_gint_dp = make_cast_hcontainer<double>(hr_gint);
compose_hr_gint(hr_gint_dp);
hr_gint_to_hR(hr_gint_dp, *hR);
}

template<typename Real>
void Gint_vl_gpu::cal_gint_impl_()
{
// 1. Initialize hr_gint as HContainer<Real>
HContainer<Real> hr_gint = gint_info_->get_hr<Real>();
// hr_gint is always allocated as HContainer<double>: 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<double> hr_gint = gint_info_->get_hr<double>();

// 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<Real> vr_eff_d(local_mgrid_num, 0, false);
CudaMemWrapper<Real> hr_gint_d(hr_gint.get_nnr(), 0, false);
CudaMemWrapper<double> hr_gint_d(hr_gint.get_nnr(), 0, false);

if (std::is_same<Real, double>::value)
{
Expand All @@ -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
Expand Down Expand Up @@ -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_);
}

}
}
8 changes: 4 additions & 4 deletions source/source_lcao/module_gint/kernel/dgemm_vbatch.cu
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand All @@ -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)
{
Expand All @@ -49,7 +49,7 @@ template void gemm_nn_vbatch<double>(
template void gemm_nn_vbatch<float>(
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<double>(
int, int, int, const int*, const int*, const int*,
Expand All @@ -59,4 +59,4 @@ template void gemm_tn_vbatch<double>(
template void gemm_tn_vbatch<float>(
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*);
Loading
Loading