Feature: EXX PW GPU Support

source/CMakeLists.txt

@@ -54,6 +54,10 @@ list(APPEND device_srcs
   source_pw/module_pwdft/kernels/vnl_op.cpp
   source_base/kernels/math_ylm_op.cpp
   source_hamilt/module_xc/kernels/xc_functional_op.cpp
+  source_pw/module_pwdft/kernels/cal_density_real_op.cpp
+  source_pw/module_pwdft/kernels/mul_potential_op.cpp
+  source_pw/module_pwdft/kernels/vec_mul_vec_complex_op.cpp
+  source_pw/module_pwdft/kernels/exx_cal_energy_op.cpp
 )
 
 if(USE_CUDA)
@@ -80,6 +84,10 @@ if(USE_CUDA)
     source_base/kernels/cuda/math_kernel_op.cu
     source_base/kernels/cuda/math_kernel_op_vec.cu
     source_hamilt/module_xc/kernels/cuda/xc_functional_op.cu
+    source_pw/module_pwdft/kernels/cuda/cal_density_real_op.cu
+    source_pw/module_pwdft/kernels/cuda/mul_potential_op.cu
+    source_pw/module_pwdft/kernels/cuda/vec_mul_vec_complex.cu
+    source_pw/module_pwdft/kernels/cuda/exx_cal_energy_op.cu
   )
 endif()
 

source/Makefile.Objects

@@ -342,6 +342,10 @@ OBJS_HAMILT=hamilt_pw.o\
     velocity_pw.o\
     radial_proj.o\
     exx_helper.o\
+    vec_mul_vec_complex_op.o\
+    exx_cal_energy_op.o\
+    cal_density_real_op.o\
+    mul_potential_op.o\
 
 OBJS_HAMILT_OF=kedf_tf.o\
     kedf_vw.o\

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

@@ -70,8 +70,6 @@ struct lapack_trtri<T, DEVICE_GPU> {
     {
         // TODO: trtri is not implemented in this method yet
         // Cause the trtri in cuSolver is not stable for ABACUS!
-        // But why?! trtri and potri are different routines for different job! 
-        // How can BPCG work without using a proper routine? 
         cuSolverConnector::trtri(cusolver_handle, uplo, diag, dim, Mat, lda);
         // cuSolverConnector::potri(cusolver_handle, uplo, diag, dim, Mat, lda);
     }

source/source_base/module_container/base/third_party/cusolver.h

@@ -87,45 +87,57 @@ static inline
 void potrf (cusolverDnHandle_t& cusolver_handle, const char& uplo, const int& n, float * A, const int& lda)
 {
     int lwork;
+    int *info = nullptr;
+    cudaErrcheck(cudaMalloc((void**)&info, 1 * sizeof(int)));
     cusolverErrcheck(cusolverDnSpotrf_bufferSize(cusolver_handle, cublas_fill_mode(uplo), n, A, n, &lwork));
     float* work;
     cudaErrcheck(cudaMalloc((void**)&work, lwork * sizeof(float)));
     // Perform Cholesky decomposition
-    cusolverErrcheck(cusolverDnSpotrf(cusolver_handle, cublas_fill_mode(uplo), n, A, n, work, lwork, nullptr));
+    cusolverErrcheck(cusolverDnSpotrf(cusolver_handle, cublas_fill_mode(uplo), n, A, n, work, lwork, info));
     cudaErrcheck(cudaFree(work));
+    cudaErrcheck(cudaFree(info));
 }
 static inline
 void potrf (cusolverDnHandle_t& cusolver_handle, const char& uplo, const int& n, double * A, const int& lda)
 {
     int lwork;
+    int *info = nullptr;
+    cudaErrcheck(cudaMalloc((void**)&info, 1 * sizeof(int)));
     cusolverErrcheck(cusolverDnDpotrf_bufferSize(cusolver_handle, cublas_fill_mode(uplo), n, A, n, &lwork));
     double* work;
     cudaErrcheck(cudaMalloc((void**)&work, lwork * sizeof(double)));
     // Perform Cholesky decomposition
-    cusolverErrcheck(cusolverDnDpotrf(cusolver_handle, cublas_fill_mode(uplo), n, A, n, work, lwork, nullptr));
+    cusolverErrcheck(cusolverDnDpotrf(cusolver_handle, cublas_fill_mode(uplo), n, A, n, work, lwork, info));
     cudaErrcheck(cudaFree(work));
+    cudaErrcheck(cudaFree(info));
 }
 static inline
 void potrf (cusolverDnHandle_t& cusolver_handle, const char& uplo, const int& n, std::complex<float> * A, const int& lda)
 {
     int lwork;
-    cusolverErrcheck(cusolverDnCpotrf_bufferSize(cusolver_handle, cublas_fill_mode(uplo), n, reinterpret_cast<cuComplex*>(A), n, &lwork));
+    int *info = nullptr;
+    cudaErrcheck(cudaMalloc((void**)&info, 1 * sizeof(int)));
+    cusolverErrcheck(cusolverDnCpotrf_bufferSize(cusolver_handle, cublas_fill_mode(uplo), n, reinterpret_cast<cuComplex*>(A), lda, &lwork));
     cuComplex* work;
     cudaErrcheck(cudaMalloc((void**)&work, lwork * sizeof(cuComplex)));
     // Perform Cholesky decomposition
-    cusolverErrcheck(cusolverDnCpotrf(cusolver_handle, cublas_fill_mode(uplo), n, reinterpret_cast<cuComplex*>(A), n, work, lwork, nullptr));
+    cusolverErrcheck(cusolverDnCpotrf(cusolver_handle, cublas_fill_mode(uplo), n, reinterpret_cast<cuComplex*>(A), lda, work, lwork, info));
     cudaErrcheck(cudaFree(work));
+    cudaErrcheck(cudaFree(info));
 }
 static inline
 void potrf (cusolverDnHandle_t& cusolver_handle, const char& uplo, const int& n, std::complex<double> * A, const int& lda)
 {
     int lwork;
-    cusolverErrcheck(cusolverDnZpotrf_bufferSize(cusolver_handle, cublas_fill_mode(uplo), n, reinterpret_cast<cuDoubleComplex*>(A), n, &lwork));
+    int *info = nullptr;
+    cudaErrcheck(cudaMalloc((void**)&info, 1 * sizeof(int)));
+    cusolverErrcheck(cusolverDnZpotrf_bufferSize(cusolver_handle, cublas_fill_mode(uplo), n, reinterpret_cast<cuDoubleComplex*>(A), lda, &lwork));
     cuDoubleComplex* work;
     cudaErrcheck(cudaMalloc((void**)&work, lwork * sizeof(cuDoubleComplex)));
     // Perform Cholesky decomposition
-    cusolverErrcheck(cusolverDnZpotrf(cusolver_handle, cublas_fill_mode(uplo), n, reinterpret_cast<cuDoubleComplex*>(A), n, work, lwork, nullptr));
+    cusolverErrcheck(cusolverDnZpotrf(cusolver_handle, cublas_fill_mode(uplo), n, reinterpret_cast<cuDoubleComplex*>(A), lda, work, lwork, info));
     cudaErrcheck(cudaFree(work));
+    cudaErrcheck(cudaFree(info));
 }
 
 

source/source_basis/module_pw/pw_basis.h

@@ -432,6 +432,9 @@ class PW_Basis
     void set_device(std::string device_);
     void set_precision(std::string precision_);
 
+    std::string get_device() const { return device; }
+    std::string get_precision() const { return precision; }
+
 protected:
 
   std::string device = "cpu";       ///< cpu or gpu

source/source_esolver/esolver_ks_pw.cpp

@@ -619,6 +619,7 @@ void ESolver_KS_PW<T, Device>::iter_finish(UnitCell& ucell, const int istep, int
             {
                 auto start = std::chrono::high_resolution_clock::now();
                 exx_helper.set_firstiter(false);
+                exx_helper.op_exx->first_iter = false;
                 exx_helper.set_psi(this->kspw_psi);
 
                 conv_esolver = exx_helper.exx_after_converge(iter);

source/source_io/input_conv.cpp

@@ -488,19 +488,15 @@ void Input_Conv::Convert()
     {
         if (ModuleSymmetry::Symmetry::symm_flag != -1)
         {
-            ModuleBase::WARNING("Input_Conv", "EXX PW works only with symmetry=-1");
-            ModuleSymmetry::Symmetry::symm_flag = -1;
+            ModuleBase::WARNING_QUIT("Input_Conv", "EXX PW works only with symmetry=-1");
+            // ModuleSymmetry::Symmetry::symm_flag = -1;
         }
 
         if (PARAM.inp.nspin != 1 && PARAM.inp.nspin != 2)
         {
             ModuleBase::WARNING_QUIT("Input_Conv", "EXX PW works only with nspin=1 and 2");
         }
 
-        if (PARAM.inp.device != "cpu")
-        {
-            ModuleBase::WARNING_QUIT("Input_Conv", "EXX PW works only with device=cpu");
-        }
     }
 
     //----------------------------------------------------------

source/source_pw/module_pwdft/CMakeLists.txt

@@ -26,6 +26,7 @@ list(APPEND objects
     stress_func_nl.cpp
     stress_func_us.cpp
     stress_func_onsite.cpp
+    stress_func_exx.cpp
     stress_pw.cpp
     VL_in_pw.cpp
     VNL_in_pw.cpp
@@ -47,7 +48,6 @@ add_library(
     module_pwdft
     OBJECT
     ${objects}
-        stress_func_exx.cpp
 )
 
 if(ENABLE_COVERAGE)

source/source_pw/module_pwdft/kernels/cal_density_real_op.cpp

@@ -0,0 +1,25 @@
+#include "source_pw/module_pwdft/kernels/cal_density_real_op.h"
+#include "source_psi/psi.h"
+namespace hamilt
+{
+template <typename T>
+struct cal_density_real_op<T, base_device::DEVICE_CPU>
+{
+    void operator()(const T *in1, const T *in2, T *out, double omega, int nrxx)
+    {
+#ifdef _OPENMP
+#pragma omp parallel for schedule(static)
+#endif
+        for (int ir = 0; ir < nrxx; ir++)
+        {
+            // assert(is_finite(psi_nk_real[ir]));
+            // assert(is_finite(psi_mq_real[ir]));
+            out[ir] = in1[ir] * std::conj(in2[ir]) / static_cast<T>(omega); // Phase e^(i(q-k)r)
+        }
+    }
+
+};
+
+template struct cal_density_real_op<std::complex<float>, base_device::DEVICE_CPU>;
+template struct cal_density_real_op<std::complex<double>, base_device::DEVICE_CPU>;
+}

source/source_pw/module_pwdft/kernels/cal_density_real_op.h

@@ -0,0 +1,14 @@
+#include "source_base/macros.h"
+
+#ifndef CAL_DENSITY_REAL_OP_H
+#define CAL_DENSITY_REAL_OP_H
+namespace hamilt
+{
+template <typename T, typename Device>
+struct cal_density_real_op
+{
+    using Real = typename GetTypeReal<T>::type;
+    void operator()(const T *psi1, const T* psi2, T *out, double omega, int nrxx);
+};
+}
+#endif //CAL_DENSITY_REAL_OP_H

source/source_pw/module_pwdft/kernels/cuda/cal_density_real_op.cu

@@ -0,0 +1,48 @@
+#include "source_pw/module_pwdft/kernels/cal_density_real_op.h"
+#include "source_psi/psi.h"
+
+#include <thrust/complex.h>
+
+namespace hamilt
+{
+template <typename FPTYPE>
+__global__ void cal_density_real_kernel(
+    const thrust::complex<FPTYPE> *in1,
+    const thrust::complex<FPTYPE> *in2,
+    thrust::complex<FPTYPE> *out,
+    const FPTYPE omega,
+    int nrxx)
+{
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < nrxx)
+    {
+        out[idx] = in1[idx] * thrust::conj(in2[idx]) / static_cast<thrust::complex<FPTYPE>>(omega);
+    }
+}
+
+template <typename FPTYPE>
+struct cal_density_real_op<std::complex<FPTYPE>, base_device::DEVICE_GPU>
+{
+    using T = std::complex<FPTYPE>;
+    void operator()(const T *psi1, const T *psi2, T *out, double omega, int nrxx)
+    {
+        int threads_per_block = 256;
+        int num_blocks = (nrxx + threads_per_block - 1) / threads_per_block;
+
+        cal_density_real_kernel<FPTYPE><<<num_blocks, threads_per_block>>>(
+            reinterpret_cast<const thrust::complex<FPTYPE> *>(psi1),
+            reinterpret_cast<const thrust::complex<FPTYPE> *>(psi2),
+            reinterpret_cast<thrust::complex<FPTYPE> *>(out),
+            static_cast<FPTYPE>(omega), nrxx);
+
+        cudaError_t err = cudaGetLastError();
+        if (err != cudaSuccess)
+        {
+            throw std::runtime_error("CUDA error in cal_density_real_kernel: " + std::string(cudaGetErrorString(err)));
+        }
+    }
+};
+
+template struct cal_density_real_op<std::complex<float>, base_device::DEVICE_GPU>;
+template struct cal_density_real_op<std::complex<double>, base_device::DEVICE_GPU>;
+} // namespace hamilt

source/source_pw/module_pwdft/kernels/cuda/exx_cal_energy_op.cu

@@ -0,0 +1,84 @@
+#include "source_pw/module_pwdft/kernels/exx_cal_energy_op.h"
+#include "source_psi/psi.h"
+
+#include <thrust/complex.h>
+
+namespace hamilt
+{
+
+// #ifdef _OPENMP
+// #pragma omp parallel for reduction(+:Eexx_ik_real)
+// #endif
+// for (int ig = 0; ig < rhopw_dev->npw; ig++)
+// {
+//     int nks = wfcpw->nks;
+//     int npw = rhopw_dev->npw;
+//     int nk = nks / nk_fac;
+//     Real Fac = pot[ik * nks * npw + iq * npw + ig];
+
+// Eexx_ik_real += Fac * (density_recip[ig] * std::conj(density_recip[ig])).real()
+//                 * wg_iqb_real / nqs * wg_ikb_real / kv->wk[ik];
+// }
+
+template <typename FPTYPE>
+__global__ void cal_vec_norm_kernel(
+    const thrust::complex<FPTYPE> *den,
+    const FPTYPE *pot,
+    FPTYPE *result,
+    int npw)
+{
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx < npw)
+    {
+        // atomicAdd(result, (den[idx].real() * den[idx].real() + den[idx].imag() * den[idx].imag()) * pot[idx]);
+        FPTYPE tmp =(den[idx] * thrust::conj(den[idx])).real() * pot[idx];
+        atomicAdd(result, tmp);
+    }
+    __syncthreads();
+}
+
+template <typename FPTYPE>
+struct exx_cal_energy_op<std::complex<FPTYPE>, base_device::DEVICE_GPU>
+{
+    using T = std::complex<FPTYPE>;
+    FPTYPE operator()(const T *den, const FPTYPE *pot, FPTYPE scalar, int npw)
+    {
+        // T *den_cpu = new T[npw];
+        // FPTYPE *pot_cpu = new FPTYPE[npw];
+        // cudaMemcpy(den_cpu, den, npw * sizeof(T), cudaMemcpyDeviceToHost);
+        // cudaMemcpy(pot_cpu, pot, npw * sizeof(FPTYPE), cudaMemcpyDeviceToHost);
+        // FPTYPE result = exx_cal_energy_op<std::complex<FPTYPE>, base_device::DEVICE_CPU>()(den_cpu, pot_cpu, scalar, npw);
+        // delete[] den_cpu;
+        // delete[] pot_cpu;
+        // return result;
+        FPTYPE result = 0.0;
+
+        int threads_per_block = 256;
+        int num_blocks = (npw + threads_per_block - 1) / threads_per_block;
+
+        FPTYPE *d_result;
+        cudaMalloc(&d_result, sizeof(FPTYPE));
+        cudaMemset(d_result, 0, sizeof(FPTYPE));
+
+        cal_vec_norm_kernel<FPTYPE><<<num_blocks, threads_per_block>>>(
+            reinterpret_cast<const thrust::complex<FPTYPE> *>(den),
+            pot,
+            d_result,
+            npw);
+
+        cudaError_t err = cudaGetLastError();
+        if (err != cudaSuccess)
+        {
+            throw std::runtime_error("CUDA error in cal_vec_norm_kernel: " + std::string(cudaGetErrorString(err)));
+        }
+
+        cudaMemcpy(&result, d_result, sizeof(FPTYPE), cudaMemcpyDeviceToHost);
+        cudaFree(d_result);
+
+        return scalar * result;
+    }
+};
+
+template struct exx_cal_energy_op<std::complex<float>, base_device::DEVICE_GPU>;
+template struct exx_cal_energy_op<std::complex<double>, base_device::DEVICE_GPU>;
+} // namespace hamilt

source/source_pw/module_pwdft/kernels/cuda/mul_potential_op.cu

@@ -0,0 +1,53 @@
+#include "source_pw/module_pwdft/kernels/mul_potential_op.h"
+#include "source_io/module_parameter/parameter.h"
+#include "source_psi/psi.h"
+
+#include <thrust/complex.h>
+
+namespace hamilt {
+template <typename FPTYPE>
+__global__ void mul_potential_kernel(
+    const FPTYPE *pot_shifted,
+    thrust::complex<FPTYPE> *density_recip,
+    int npw)
+{
+    int ig = blockIdx.x * blockDim.x + threadIdx.x;
+    if (ig < npw)
+    {
+        density_recip[ig] *= pot_shifted[ig];
+    }
+}
+
+template <typename FPTYPE>
+struct mul_potential_op<std::complex<FPTYPE>, base_device::DEVICE_GPU>
+{
+    using T = std::complex<FPTYPE>;
+    void operator()(const FPTYPE *pot, T *density_recip, int npw, int nks, int ik, int iq)
+    {
+// #ifdef _OPENMP
+// #pragma omp parallel for schedule(static)
+// #endif
+//         for (int ig = 0; ig < npw; ig++)
+//         {
+//             int ig_kq = ik * nks * npw + iq * npw + ig;
+//             density_recip[ig] *= pot[ig_kq];
+//
+//         }
+        int threads_per_block = 256;
+        int num_blocks = (npw + threads_per_block - 1) / threads_per_block;
+
+        mul_potential_kernel<<<num_blocks, threads_per_block>>>(
+            pot + ik * nks * npw + iq * npw,
+            reinterpret_cast<thrust::complex<FPTYPE>*>(density_recip),
+            npw);
+
+        cudaError_t err = cudaGetLastError();
+        if (err != cudaSuccess)
+        {
+            throw std::runtime_error("CUDA error in mul_potential_kernel: " + std::string(cudaGetErrorString(err)));
+        }
+    }
+};
+template struct mul_potential_op<std::complex<float>, base_device::DEVICE_GPU>;
+template struct mul_potential_op<std::complex<double>, base_device::DEVICE_GPU>;
+} // hamilt