Feature: Support GPU diagHsubspace in CG method

source/module_base/lapack_connector.h

@@ -159,6 +159,20 @@ class LapackConnector
         return aux;
     }
 
+    static inline
+    std::complex<double>* transpose(const std::complex<double>* a, const int n, const int lda, const int nbase_x)
+    {
+        std::complex<double>* aux = new std::complex<double>[lda*n];
+        for (int i = 0; i < n; ++i)
+        {
+            for (int j = 0; j < lda; ++j)
+            {
+                aux[j * n + i] = a[i * nbase_x + j];
+            }
+        }
+        return aux;
+    }
+
     // Transpose the fortran-form real-std::complex array to the std::complex matrix.
     static inline
     void transpose(const std::complex<double>* aux, ModuleBase::ComplexMatrix& a, const int n, const int lda)
@@ -224,6 +238,31 @@ class LapackConnector
         delete[] aux;
         delete[] bux;
     }
+    // wrap function of fortran lapack routine zhegv.
+    static inline
+    void zhegv(	const int itype, const char jobz, const char uplo, const int n, std::complex<double>* a,
+                const int lda, const std::complex<double>* b, const int ldb, double* w, std::complex<double>* work,
+                int lwork, double* rwork, int info, int ld_real)
+    {	// Transpose the std::complex matrix to the fortran-form real-std::complex array.
+        std::complex<double>* aux = LapackConnector::transpose(a, n, lda, ld_real);
+        std::complex<double>* bux = LapackConnector::transpose(b, n, ldb, ld_real);
+
+        // call the fortran routine
+        zhegv_(&itype, &jobz, &uplo, &n, aux, &lda, bux, &ldb, w, work, &lwork, rwork, &info);
+        // Transpose the fortran-form real-std::complex array to the std::complex matrix.
+        // LapackConnector::transpose(aux, a, n, lda);
+        // LapackConnector::transpose(bux, b, n, ldb);
+        for (int i = 0; i < n; ++i)
+        {
+            for (int j = 0; j < lda; ++j)
+            {
+                a[j * ld_real + i] = aux[i*lda+j];
+            }
+        }
+        // free the memory.
+        delete[] aux;
+        delete[] bux;
+    }
 
     // wrap function of fortran lapack routine zheev.
     static inline
@@ -252,6 +291,38 @@ class LapackConnector
         delete[] bux;
         delete[] zux;
 
+    }
+    // wrap function of fortran lapack routine zheev.
+    static inline
+    void zhegvx( const int itype, const char jobz, const char range, const char uplo,
+                 const int n, const std::complex<double>* a, const int lda, const std::complex<double>* b,
+                 const int ldb, const double vl, const double vu, const int il, const int iu,
+                 const double abstol, const int m, double* w, std::complex<double>* z, const int ldz,
+                 std::complex<double>* work, const int lwork, double* rwork, int* iwork,
+                 int* ifail, int info, int nbase_x)
+    {
+        // Transpose the std::complex matrix to the fortran-form real-std::complex array.
+        std::complex<double>* aux = LapackConnector::transpose(a, n, lda, nbase_x);
+        std::complex<double>* bux = LapackConnector::transpose(b, n, ldb, nbase_x);
+        std::complex<double>* zux = new std::complex<double>[n*iu];// mohan modify 2009-08-02
+
+        // call the fortran routine
+        zhegvx_(&itype, &jobz, &range, &uplo, &n, aux, &lda, bux, &ldb, &vl,
+                &vu, &il,&iu, &abstol, &m, w, zux, &ldz, work, &lwork, rwork, iwork, ifail, &info);
+
+        // Transpose the fortran-form real-std::complex array to the std::complex matrix
+        for (int i = 0; i < iu; ++i)
+        {
+            for (int j = 0; j < n; ++j)
+            {
+                z[j * nbase_x + i] = zux[i*n+j];
+            }
+        }
+
+        // free the memory.
+        delete[] aux;
+        delete[] bux;
+        delete[] zux;
     }
 
 	// calculate the eigenvalues and eigenfunctions of a real symmetric matrix.

source/module_hsolver/CMakeLists.txt

@@ -5,6 +5,7 @@ list(APPEND objects
     hsolver_pw_sdft.cpp
     diago_iter_assist.cpp
     src/math_kernel.cpp
+    src/dngvd_op.cpp
 )
 
 if(ENABLE_LCAO)
@@ -16,7 +17,10 @@ if(ENABLE_LCAO)
 endif()
 
 if (USE_CUDA)
-  list(APPEND objects src/cuda/math_kernel.cu)
+  list(APPEND objects 
+      src/cuda/math_kernel.cu
+      src/cuda/dngvd_op.cu
+  )
 elseif(USE_ROCM)
   list(APPEND objects src/rocm/math_kernel.cu)
 endif()

source/module_hsolver/diago_cg.cpp

@@ -683,23 +683,7 @@ void DiagoCG<FPTYPE, Device>::diag(hamilt::Hamilt<FPTYPE, Device> *phm_in, psi::
     {
         if(DiagoIterAssist<FPTYPE>::need_subspace || ntry > 0)
         {
-            #if defined(__CUDA) || defined(__ROCM)
-            psi::Psi<std::complex<FPTYPE>, psi::DEVICE_CPU> cpu_psi = psi;
-            // hamilt::Hamilt<FPTYPE>* h_phm_in = new hamilt::HamiltPW<FPTYPE>(reinterpret_cast<hamilt::HamiltPW<FPTYPE>*>(phm_in));
-            hamilt::Hamilt<FPTYPE, psi::DEVICE_CPU>* h_phm_in =
-                    new hamilt::HamiltPW<FPTYPE, psi::DEVICE_CPU>(
-                            reinterpret_cast<hamilt::HamiltPW<FPTYPE, psi::DEVICE_GPU>*>(phm_in));
-            DiagoIterAssist<FPTYPE>::diagH_subspace(h_phm_in, cpu_psi, cpu_psi, eigenvalue_in);
-            psi::memory::synchronize_memory_op<std::complex<FPTYPE>, Device, psi::DEVICE_CPU>()(
-                    psi.get_device(),
-                    cpu_psi.get_device(),
-                    psi.get_pointer() - psi.get_psi_bias(),
-                    cpu_psi.get_pointer() - cpu_psi.get_psi_bias(),
-                    psi.size());
-            delete reinterpret_cast<hamilt::HamiltPW<FPTYPE, psi::DEVICE_CPU>*>(h_phm_in);
-            #else
-            DiagoIterAssist<FPTYPE>::diagH_subspace(phm_in, psi, psi, eigenvalue_in);
-            #endif
+            DiagoIterAssist<FPTYPE, Device>::diagH_subspace(phm_in, psi, psi, eigenvalue_in);
         }
 
         DiagoIterAssist<FPTYPE>::avg_iter += 1.0;