Feature：Now sDFT support GPU (#5423)

* make chebyshev support GPU

* make part of stoiter support GPU

* make scf support GPU

* add in GPU class implement

* fix GPU bug

* fix compile

* fix bug

* fix compile

* try fix pytest

Author: Qianruipku

python/pyabacus/src/ModuleBase/CMakeLists.txt

@@ -2,6 +2,7 @@ list(APPEND pymodule_base
     ${PROJECT_SOURCE_DIR}/src/ModuleBase/py_base_math.cpp
     ${BASE_PATH}/kernels/math_op.cpp
     ${BASE_PATH}/module_device/memory_op.cpp
+    ${BASE_PATH}/module_device/device.cpp
     )
 
 pybind11_add_module(_base_pack MODULE ${pymodule_base})

python/pyabacus/src/ModuleNAO/CMakeLists.txt

@@ -17,6 +17,7 @@ list(APPEND _naos
     ${ABACUS_SOURCE_DIR}/module_base/kernels/math_op.cpp
     # ${ABACUS_SOURCE_DIR}/module_psi/kernels/psi_memory_op.cpp
     ${ABACUS_SOURCE_DIR}/module_base/module_device/memory_op.cpp
+    ${ABACUS_SOURCE_DIR}/module_base/module_device/device.cpp
 )
 add_library(naopack SHARED 
     ${_naos}

source/CMakeLists.txt

@@ -30,6 +30,7 @@ list(APPEND device_srcs
   module_hamilt_pw/hamilt_pwdft/kernels/veff_op.cpp
   module_hamilt_pw/hamilt_pwdft/kernels/ekinetic_op.cpp
   module_hamilt_pw/hamilt_pwdft/kernels/meta_op.cpp
+  module_hamilt_pw/hamilt_stodft/kernels/hpsi_norm_op.cpp
   module_basis/module_pw/kernels/pw_op.cpp
   module_hsolver/kernels/dngvd_op.cpp
   module_hsolver/kernels/math_kernel_op.cpp
@@ -55,6 +56,7 @@ if(USE_CUDA)
     module_hamilt_pw/hamilt_pwdft/kernels/cuda/veff_op.cu
     module_hamilt_pw/hamilt_pwdft/kernels/cuda/ekinetic_op.cu
     module_hamilt_pw/hamilt_pwdft/kernels/cuda/meta_op.cu
+    module_hamilt_pw/hamilt_stodft/kernels/cuda/hpsi_norm_op.cu
     module_basis/module_pw/kernels/cuda/pw_op.cu
     module_hsolver/kernels/cuda/dngvd_op.cu
     module_hsolver/kernels/cuda/math_kernel_op.cu
@@ -78,6 +80,7 @@ if(USE_ROCM)
     module_hamilt_pw/hamilt_pwdft/kernels/rocm/veff_op.hip.cu
     module_hamilt_pw/hamilt_pwdft/kernels/rocm/ekinetic_op.hip.cu
     module_hamilt_pw/hamilt_pwdft/kernels/rocm/meta_op.hip.cu
+    module_hamilt_pw/hamilt_stodft/kernels/rocm/hpsi_norm_op.hip.cu
     module_basis/module_pw/kernels/rocm/pw_op.hip.cu
     module_hsolver/kernels/rocm/dngvd_op.hip.cu
     module_hsolver/kernels/rocm/math_kernel_op.hip.cu

source/Makefile.Objects

@@ -49,6 +49,7 @@ VPATH=./src_global:\
 ./module_hamilt_pw/hamilt_stodft:\
 ./module_hamilt_pw/hamilt_pwdft/operator_pw:\
 ./module_hamilt_pw/hamilt_pwdft/kernels:\
+./module_hamilt_pw/hamilt_stodft/kernels:\
 ./module_hamilt_lcao/module_hcontainer:\
 ./module_hamilt_lcao/hamilt_lcaodft:\
 ./module_hamilt_lcao/module_tddft:\
@@ -296,6 +297,7 @@ OBJS_HAMILT=hamilt_pw.o\
     operator_pw.o\
     ekinetic_pw.o\
     ekinetic_op.o\
+    hpsi_norm_op.o\
     veff_pw.o\
     veff_op.o\
     nonlocal_pw.o\

source/module_base/CMakeLists.txt

@@ -61,6 +61,8 @@ add_library(
     ${LIBM_SRC}
 )
 
+target_link_libraries(base PUBLIC container)
+
 add_subdirectory(module_container)
 
 if(ENABLE_COVERAGE)

source/module_base/blas_connector.cpp

@@ -154,6 +154,15 @@ void BlasConnector::gemm(const char transa, const char transb, const int m, cons
 	#endif
 }
 
+void BlasConnector::gemv(const char trans, const int m, const int n,
+    const float alpha, const float* A, const int lda, const float* X, const int incx,
+    const float beta, float* Y, const int incy, base_device::AbacusDevice_t device_type)
+{
+	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
+    	sgemv_(&trans, &m, &n, &alpha, A, &lda, X, &incx, &beta, Y, &incy);
+}
+}
+
 void BlasConnector::gemv(const char trans, const int m, const int n,
     const double alpha, const double* A, const int lda, const double* X, const int incx,
     const double beta, double* Y, const int incy, base_device::AbacusDevice_t device_type)

source/module_base/blas_connector.h

@@ -40,6 +40,9 @@ extern "C"
 	double dznrm2_( const int *n, const std::complex<double> *X, const int *incX );
 
 	// level 2: matrix-std::vector operations, O(n^2) data and O(n^2) work.
+	void sgemv_(const char*const transa, const int*const m, const int*const n,
+		const float*const alpha, const float*const a, const int*const lda, const float*const x, const int*const incx,
+		const float*const beta, float*const y, const int*const incy);
 	void dgemv_(const char*const transa, const int*const m, const int*const n,
 		const double*const alpha, const double*const a, const int*const lda, const double*const x, const int*const incx,
 		const double*const beta, double*const y, const int*const incy);
@@ -178,6 +181,11 @@ class BlasConnector
 		const std::complex<double> alpha, const std::complex<double> *a, const int lda, const std::complex<double> *b, const int ldb,
 		const std::complex<double> beta, std::complex<double> *c, const int ldc, base_device::AbacusDevice_t device_type = base_device::AbacusDevice_t::CpuDevice);
 
+	static
+	void gemv(const char trans, const int m, const int n,
+        const float alpha, const float* A, const int lda, const float* X, const int incx,
+        const float beta, float* Y, const int incy, base_device::AbacusDevice_t device_type = base_device::AbacusDevice_t::CpuDevice);
+
     static
     void gemv(const char trans, const int m, const int n,
         const double alpha, const double* A, const int lda, const double* X, const int incx,