Merge branch 'develop' into fix/bpcg-gemm-insteadof-einsum

Author: Cstandardlib

docs/advanced/input_files/input-main.md

@@ -11,7 +11,6 @@
     - [kpar](#kpar)
     - [bndpar](#bndpar)
     - [latname](#latname)
-    - [psi\_initializer](#psi_initializer)
     - [init\_wfc](#init_wfc)
     - [init\_chg](#init_chg)
     - [init\_vel](#init_vel)
@@ -93,6 +92,7 @@
     - [scf\_os\_stop](#scf_os_stop)
     - [scf\_os\_thr](#scf_os_thr)
     - [scf\_os\_ndim](#scf_os_ndim)
+    - [sc\_os\_ndim](#sc_os_ndim)
     - [chg\_extrap](#chg_extrap)
     - [lspinorb](#lspinorb)
     - [noncolin](#noncolin)
@@ -467,7 +467,7 @@
     - [abs\_broadening](#abs_broadening)
     - [ri\_hartree\_benchmark](#ri_hartree_benchmark)
     - [aims\_nbasis](#aims_nbasis)
-  - [Reduced Density Matrix Functional Theory](#Reduced-Density-Matrix-Functional-Theory)
+  - [Reduced Density Matrix Functional Theory](#reduced-density-matrix-functional-theory)
     - [rdmft](#rdmft)
     - [rdmft\_power\_alpha](#rdmft_power_alpha)
 
@@ -580,17 +580,6 @@ These variables are used to control general system parameters.
   - triclinic: triclinic (14)
 - **Default**: none
 
-### psi_initializer
-
-- **Type**: Integer
-- **Description**: enable the experimental feature psi_initializer, to support use numerical atomic orbitals initialize wavefunction (`basis_type pw` case).
-
-  NOTE: this feature is not well-implemented for `nspin 4` case (closed presently), and cannot use with `calculation nscf`/`esolver_type sdft` cases.
-  Available options are:
-  - 0: disable psi_initializer
-  - 1: enable psi_initializer
-- **Default**: 0
-
 ### init_wfc
 
 - **Type**: String
@@ -602,8 +591,6 @@ These variables are used to control general system parameters.
   - atomic+random: add small random numbers on atomic pseudo-wavefunctions
   - file: from binary files `WAVEFUNC*.dat`, which are output by setting [out_wfc_pw](#out_wfc_pw) to `2`.
   - random: random numbers
-
-  with `psi_initializer 1`, two more options are supported:
   - nao: from numerical atomic orbitals. If they are not enough, other wave functions are initialized with random numbers.
   - nao+random: add small random numbers on numerical atomic orbitals
 - **Default**: atomic

source/Makefile.Objects

@@ -69,7 +69,7 @@ VPATH=./src_global:\
 ./module_ri:\
 ./module_parameter:\
 ./module_lr:\
-./module_lr/AX:\
+./module_lr/ao_to_mo_transformer:\
 ./module_lr/dm_trans:\
 ./module_lr/operator_casida:\
 ./module_lr/potentials:\
@@ -189,6 +189,8 @@ OBJS_CELL=atom_pseudo.o\
     check_atomic_stru.o\
     update_cell.o\
     bcast_cell.o\
+    read_stru.o\
+    read_atom_species.o\
 
 OBJS_DEEPKS=LCAO_deepks.o\
         deepks_force.o\
@@ -398,6 +400,7 @@ OBJS_PSI=psi.o\
 
 OBJS_PSI_INITIALIZER=psi_initializer.o\
                      psi_initializer_random.o\
+                     psi_initializer_file.o\
                      psi_initializer_atomic.o\
                      psi_initializer_atomic_random.o\
                      psi_initializer_nao.o\
@@ -494,6 +497,7 @@ OBJS_IO=input_conv.o\
     to_wannier90_lcao.o\
     fR_overlap.o\
     unk_overlap_pw.o\
+    write_pao.o\
     write_wfc_pw.o\
     winput.o\
     write_cube.o\
@@ -669,8 +673,6 @@ OBJS_SRCPW=H_Ewald_pw.o\
     of_stress_pw.o\
     symmetry_rho.o\
     symmetry_rhog.o\
-    wavefunc.o\
-    wf_atomic.o\
     psi_init.o\
     elecond.o\
     sto_tool.o\
@@ -723,8 +725,8 @@ OBJS_TENSOR=tensor.o\
 
     OBJS_LR=lr_util.o\
     lr_util_hcontainer.o\
-    AX_parallel.o\
-    AX_serial.o\
+    ao_to_mo_parallel.o\
+    ao_to_mo_serial.o\
     dm_trans_parallel.o\
     dm_trans_serial.o\
     dmr_complex.o\

source/module_base/blas_connector.cpp

@@ -50,6 +50,36 @@ namespace BlasUtils{
 		return CUBLAS_OP_N;
 	}
 
+	cublasSideMode_t judge_side(const char& trans)
+	{
+		if (trans == 'L')
+		{
+			return CUBLAS_SIDE_LEFT;
+		}
+		else if (trans == 'R')
+		{
+			return CUBLAS_SIDE_RIGHT;
+		}
+		return CUBLAS_SIDE_LEFT;
+	}
+
+	cublasFillMode_t judge_fill(const char& trans)
+	{
+		if (trans == 'F')
+		{
+			return CUBLAS_FILL_MODE_FULL;
+		}
+		else if (trans == 'U')
+		{
+			return CUBLAS_FILL_MODE_UPPER;
+		}
+		else if (trans == 'D')
+		{
+			return CUBLAS_FILL_MODE_LOWER;
+		}
+		return CUBLAS_FILL_MODE_FULL;
+	}
+
 } // namespace BlasUtils
 
 #endif
@@ -398,6 +428,13 @@ void BlasConnector::symm_cm(const char side, const char uplo, const int m, const
 		&alpha, a, &lda, b, &ldb,
 		&beta, c, &ldc);
 	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		cublasSideMode_t sideMode = BlasUtils::judge_side(side);
+		cublasFillMode_t fillMode = BlasUtils::judge_fill(uplo);
+		cublasErrcheck(cublasSsymm(BlasUtils::cublas_handle, sideMode, fillMode, m, n, &alpha, a, lda, b, ldb, &beta, c, ldc));
+#endif
+	}
 }
 
 void BlasConnector::symm_cm(const char side, const char uplo, const int m, const int n,
@@ -409,6 +446,13 @@ void BlasConnector::symm_cm(const char side, const char uplo, const int m, const
 		&alpha, a, &lda, b, &ldb,
 		&beta, c, &ldc);
 	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		cublasSideMode_t sideMode = BlasUtils::judge_side(side);
+		cublasFillMode_t fillMode = BlasUtils::judge_fill(uplo);
+		cublasErrcheck(cublasDsymm(BlasUtils::cublas_handle, sideMode, fillMode, m, n, &alpha, a, lda, b, ldb, &beta, c, ldc));
+#endif
+	}
 }
 
 void BlasConnector::symm_cm(const char side, const char uplo, const int m, const int n,
@@ -420,6 +464,13 @@ void BlasConnector::symm_cm(const char side, const char uplo, const int m, const
 		&alpha, a, &lda, b, &ldb,
 		&beta, c, &ldc);
 	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		cublasSideMode_t sideMode = BlasUtils::judge_side(side);
+		cublasFillMode_t fillMode = BlasUtils::judge_fill(uplo);
+		cublasErrcheck(cublasCsymm(BlasUtils::cublas_handle, sideMode, fillMode, m, n, (float2*)&alpha, (float2*)a, lda, (float2*)b, ldb, (float2*)&beta, (float2*)c, ldc));
+#endif
+	}
 }
 
 void BlasConnector::symm_cm(const char side, const char uplo, const int m, const int n,
@@ -431,6 +482,13 @@ void BlasConnector::symm_cm(const char side, const char uplo, const int m, const
 		&alpha, a, &lda, b, &ldb,
 		&beta, c, &ldc);
 	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		cublasSideMode_t sideMode = BlasUtils::judge_side(side);
+		cublasFillMode_t fillMode = BlasUtils::judge_fill(uplo);
+		cublasErrcheck(cublasZsymm(BlasUtils::cublas_handle, sideMode, fillMode, m, n, (double2*)&alpha, (double2*)a, lda, (double2*)b, ldb, (double2*)&beta, (double2*)c, ldc));
+#endif
+	}
 }
 
 void BlasConnector::hemm_cm(char side, char uplo, int m, int n,
@@ -442,6 +500,13 @@ void BlasConnector::hemm_cm(char side, char uplo, int m, int n,
 		&alpha, a, &lda, b, &ldb,
 		&beta, c, &ldc);
 	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		cublasSideMode_t sideMode = BlasUtils::judge_side(side);
+		cublasFillMode_t fillMode = BlasUtils::judge_fill(uplo);
+		cublasErrcheck(cublasChemm(BlasUtils::cublas_handle, sideMode, fillMode, m, n, (float2*)&alpha, (float2*)a, lda, (float2*)b, ldb, (float2*)&beta, (float2*)c, ldc));
+#endif
+	}
 }
 
 void BlasConnector::hemm_cm(char side, char uplo, int m, int n,
@@ -453,6 +518,13 @@ void BlasConnector::hemm_cm(char side, char uplo, int m, int n,
 		&alpha, a, &lda, b, &ldb,
 		&beta, c, &ldc);
 	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		cublasSideMode_t sideMode = BlasUtils::judge_side(side);
+		cublasFillMode_t fillMode = BlasUtils::judge_fill(uplo);
+		cublasErrcheck(cublasZhemm(BlasUtils::cublas_handle, sideMode, fillMode, m, n, (double2*)&alpha, (double2*)a, lda, (double2*)b, ldb, (double2*)&beta, (double2*)c, ldc));
+#endif
+	}
 }
 
 void BlasConnector::gemv(const char trans, const int m, const int n,
@@ -461,7 +533,13 @@ void BlasConnector::gemv(const char trans, const int m, const int n,
 {
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
     	sgemv_(&trans, &m, &n, &alpha, A, &lda, X, &incx, &beta, Y, &incy);
-}
+	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		cublasOperation_t cutransA = BlasUtils::judge_trans(false, trans, "gemv_op");
+		cublasErrcheck(cublasSgemv(BlasUtils::cublas_handle, cutransA, m, n, &alpha, A, lda, X, incx, &beta, Y, incy));
+#endif
+	}
 }
 
 void BlasConnector::gemv(const char trans, const int m, const int n,
@@ -470,7 +548,13 @@ void BlasConnector::gemv(const char trans, const int m, const int n,
 {
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
     	dgemv_(&trans, &m, &n, &alpha, A, &lda, X, &incx, &beta, Y, &incy);
-}
+	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		cublasOperation_t cutransA = BlasUtils::judge_trans(false, trans, "gemv_op");
+		cublasErrcheck(cublasDgemv(BlasUtils::cublas_handle, cutransA, m, n, &alpha, A, lda, X, incx, &beta, Y, incy));
+#endif
+	}
 }
 
 void BlasConnector::gemv(const char trans, const int m, const int n,
@@ -479,7 +563,15 @@ void BlasConnector::gemv(const char trans, const int m, const int n,
 {
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
     	cgemv_(&trans, &m, &n, &alpha, A, &lda, X, &incx, &beta, Y, &incy);
-}
+	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		cuFloatComplex alpha_cu = make_cuFloatComplex(alpha.real(), alpha.imag());
+    	cuFloatComplex beta_cu = make_cuFloatComplex(beta.real(), beta.imag());
+		cublasOperation_t cutransA = BlasUtils::judge_trans(true, trans, "gemv_op");
+		cublasErrcheck(cublasCgemv(BlasUtils::cublas_handle, cutransA, m, n, &alpha_cu, (cuFloatComplex*)A, lda, (cuFloatComplex*)X, incx, &beta_cu, (cuFloatComplex*)Y, incy));
+#endif
+	}
 }
 
 void BlasConnector::gemv(const char trans, const int m, const int n,
@@ -488,7 +580,15 @@ void BlasConnector::gemv(const char trans, const int m, const int n,
 {
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
     	zgemv_(&trans, &m, &n, &alpha, A, &lda, X, &incx, &beta, Y, &incy);
-}
+	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		cuDoubleComplex alpha_cu = make_cuDoubleComplex(alpha.real(), alpha.imag());
+    	cuDoubleComplex beta_cu = make_cuDoubleComplex(beta.real(), beta.imag());
+		cublasOperation_t cutransA = BlasUtils::judge_trans(true, trans, "gemv_op");
+		cublasErrcheck(cublasZgemv(BlasUtils::cublas_handle, cutransA, m, n, &alpha_cu, (cuDoubleComplex*)A, lda, (cuDoubleComplex*)X, incx, &beta_cu, (cuDoubleComplex*)Y, incy));
+#endif
+	}
 }
 
 // out = ||x||_2
@@ -497,6 +597,13 @@ float BlasConnector::nrm2( const int n, const float *X, const int incX, base_dev
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
 		return snrm2_( &n, X, &incX );
 	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		float result = 0.0;
+		cublasErrcheck(cublasSnrm2(BlasUtils::cublas_handle, n, X, incX, &result));
+		return result;
+#endif
+	}
 	return snrm2_( &n, X, &incX );
 }
 
@@ -506,6 +613,13 @@ double BlasConnector::nrm2( const int n, const double *X, const int incX, base_d
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
 		return dnrm2_( &n, X, &incX );
 	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		double result = 0.0;
+		cublasErrcheck(cublasDnrm2(BlasUtils::cublas_handle, n, X, incX, &result));
+		return result;
+#endif
+	}
 	return dnrm2_( &n, X, &incX );
 }
 
@@ -515,6 +629,13 @@ double BlasConnector::nrm2( const int n, const std::complex<double> *X, const in
 	if (device_type == base_device::AbacusDevice_t::CpuDevice) {
 		return dznrm2_( &n, X, &incX );
 	}
+	else if (device_type == base_device::AbacusDevice_t::GpuDevice){
+#ifdef __CUDA
+		double result = 0.0;
+		cublasErrcheck(cublasDznrm2(BlasUtils::cublas_handle, n, (double2*)X, incX, &result));
+		return result;
+#endif
+	}
 	return dznrm2_( &n, X, &incX );
 }
 

source/module_base/module_container/ATen/ops/test/CMakeLists.txt

@@ -1,7 +1,7 @@
 AddTest(
   TARGET container_ops_uts
   LIBS parameter ${math_libs} 
-  SOURCES einsum_op_test.cpp linalg_op_test.cpp
+  SOURCES einsum_op_test.cpp linalg_op_test.cpp ../../kernels/lapack.cpp
 )
 
 target_link_libraries(container_ops_uts container base device)

source/module_basis/module_pw/pw_basis_k.cpp

@@ -22,7 +22,6 @@ PW_Basis_K::~PW_Basis_K()
     delete[] igl2isz_k;
     delete[] igl2ig_k;
     delete[] gk2;
-    delete[] ig2ixyz_k_;
 #if defined(__CUDA) || defined(__ROCM)
     if (this->device == "gpu") {
         if (this->precision == "single") {
@@ -169,6 +168,7 @@ void PW_Basis_K::setupIndGk()
         syncmem_int_h2d_op()(gpu_ctx, cpu_ctx, this->d_igl2isz_k, this->igl2isz_k, this->npwk_max * this->nks);
     }
 #endif
+    this->get_ig2ixyz_k();
     return;
 }
 
@@ -334,8 +334,12 @@ int& PW_Basis_K::getigl2ig(const int ik, const int igl) const
 
 void PW_Basis_K::get_ig2ixyz_k()
 {
-    delete[] this->ig2ixyz_k_;
-    this->ig2ixyz_k_ = new int [this->npwk_max * this->nks];
+    if (this->device != "gpu")
+    {
+        //only GPU need to get ig2ixyz_k
+        return;
+    }
+    int * ig2ixyz_k_cpu = new int [this->npwk_max * this->nks];
     ModuleBase::Memory::record("PW_B_K::ig2ixyz", sizeof(int) * this->npwk_max * this->nks);
     assert(gamma_only == false); //We only finish non-gamma_only fft on GPU temperarily.
     for(int ik = 0; ik < this->nks; ++ik)
@@ -348,15 +352,12 @@ void PW_Basis_K::get_ig2ixyz_k()
             int ixy = this->is2fftixy[is];
             int iy = ixy % this->ny;
             int ix = ixy / this->ny;
-            ig2ixyz_k_[igl + ik * npwk_max] = iz + iy * nz + ix * ny * nz;
+            ig2ixyz_k_cpu[igl + ik * npwk_max] = iz + iy * nz + ix * ny * nz;
         }
     }
-#if defined(__CUDA) || defined(__ROCM)
-    if (this->device == "gpu") {
-        resmem_int_op()(gpu_ctx, ig2ixyz_k, this->npwk_max * this->nks);
-        syncmem_int_h2d_op()(gpu_ctx, cpu_ctx, this->ig2ixyz_k, this->ig2ixyz_k_, this->npwk_max * this->nks);
-    }
-#endif
+    resmem_int_op()(gpu_ctx, ig2ixyz_k, this->npwk_max * this->nks);
+    syncmem_int_h2d_op()(gpu_ctx, cpu_ctx, this->ig2ixyz_k, ig2ixyz_k_cpu, this->npwk_max * this->nks);
+    delete[] ig2ixyz_k_cpu;
 }
 
 std::vector<int> PW_Basis_K::get_ig2ix(const int ik) const