Add: use cuSolver for LCAO in GPU (#799)

Co-authored-by: Xia, Yu <[email protected]>

Author: haxushu

CMakeLists.txt

@@ -13,7 +13,8 @@ project(ABACUS
 
 option(ENABLE_DEEPKS "Enable DeePKS functionality" OFF)
 option(ENABLE_LIBXC "Enable LibXC functionality" OFF)
-option(USE_CUDA "Enable support to CUDA." OFF)
+option(USE_CUDA "Enable support to CUDA for PW." OFF)
+option(USE_CUSOLVER_LCAO "Enable support to CUSOLVER for LCAO." OFF)
 option(USE_ROCM "Enable support to ROCm." OFF)
 option(USE_OPENMP " Enable OpenMP in abacus." ON)
 option(ENABLE_ASAN "Enable AddressSanitizer" OFF)
@@ -58,21 +59,32 @@ set(CMAKE_CXX_STANDARD 11)
 include(CheckLanguage)
 check_language(CUDA)
 if(CMAKE_CUDA_COMPILER)
-  if(NOT DEFINED USE_CUDA)
-    message("CUDA components detected. \nWill build the CUDA version of ABACUS.")
-    set(USE_CUDA ON)
+  if(NOT DEFINED USE_CUDA OR NOT DEFINED USE_CUSOLVER_LCAO)
+    if (NOT DEFINED USE_CUDA AND NOT DEFINED USE_CUSOLVER_LCAO)
+      message("CUDA components detected. \nWill build the CUDA for PW version of ABACUS by default.")
+      set(USE_CUDA ON)
+      set(USE_CUSOLVER_LCAO OFF)
+    elseif (NOT DEFINED USE_CUDA)
+      set(USE_CUDA OFF)
+    else()
+      set(USE_CUSOLVER_LCAO OFF)
+    endif()
   else()
-    if(NOT USE_CUDA)
-      message(WARNING "CUDA components detected, but USE_CUDA set to OFF. \nNOT building CUDA version of ABACUS.")
+    if(NOT USE_CUDA AND NOT USE_CUSOLVER_LCAO)
+      message(WARNING "CUDA components detected, but both USE_CUDA and USE_CUSOLVER_LCAO set to OFF. \nNOT building CUDA version of ABACUS.")
+    elseif (USE_CUDA AND USE_CUSOLVER_LCAO)
+      message(FATAL_ERROR "USE_CUDA and USE_CUSOLVER_LCAO set, but now they not allowed to coexist.")
     endif()
   endif()
 else() # CUDA not found
-  if (USE_CUDA)
-    message(FATAL_ERROR "USE_CUDA set but no CUDA components found.")
+  if (USE_CUDA OR USE_CUSOLVER_LCAO)
+    message(FATAL_ERROR "USE_CUDA or USE_CUSOLVER_LCAO set but no CUDA components found.")
     set(USE_CUDA OFF)
+    set(USE_CUSOLVER_LCAO OFF)
   endif()
 endif()
-if(USE_CUDA)
+
+if(USE_CUDA OR USE_CUSOLVER_LCAO)
   set(CMAKE_CXX_STANDARD 14)
   set(CMAKE_CXX_EXTENSIONS ON)
   set(CMAKE_CXX_STANDARD_REQUIRED ON)
@@ -92,11 +104,18 @@ if(USE_CUDA)
   60 # P100
   70 # V100
   75 # T4
+  80 # A100
   )
   include_directories(${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
-  add_compile_definitions(__CUDA)
+  if (USE_CUDA)
+    add_compile_definitions(__CUDA)
+  endif()
+  if (USE_CUSOLVER_LCAO)
+    add_compile_definitions(__CUSOLVER_LCAO)
+  endif()
 endif()
 
+
 # Warning: CMake add support to HIP in version 3.21. This is rather a new version.
 # Use cmake with AMD-ROCm: https://rocmdocs.amd.com/en/latest/Installation_Guide/Using-CMake-with-AMD-ROCm.html
 if(USE_ROCM)

docs/features.md

@@ -74,7 +74,7 @@ equation. For PW basis, there are CG and Blocked Davidson methods for solving th
 equation for each basis.
 
 - PW: ks_solver = ‘cg’ or ‘dav’
-- LCAO: ks_solver = ‘hpseps’ , ‘genelpa’ or ‘lapack’
+- LCAO: ks_solver = ‘hpseps’ , ‘genelpa’ , ‘scalapack_gvx’ or 'cusolver'
 - LCAO_in_PW: ks_solver = ‘lapack’
 
 If you set ks_solver=‘hpseps’ for basis_type=‘pw’, the program will be stopped with an error

docs/input-main.md

@@ -441,6 +441,7 @@ calculations.
   - genelpa: This method should be used if you choose localized orbitals.
   - hpseps: old method, still used.
   - lapack: lapack can be used for localized orbitals, but is only used for single processor.
+  - cusolver: this method needs building with the cusolver component for lcao and at least one gpu is available.
 
    If you set ks_solver=`hpseps` for basis_type=`pw`, the program will be stopped with an error message:
 

docs/install.md

@@ -75,6 +75,9 @@ You can also choose to build with which components.
 ```bash
 cmake -B build -DUSE_LIBXC=1 -DUSE_CUDA=1
 ```
+```bash
+cmake -B build -DUSE_CUSOLVER_LCAO=1
+```
 
 If Libxc is not installed in standard path (i.e. installed with a custom prefix path), you may add the installation prefix of `FindLibxc.cmake` to `CMAKE_MODULE_PATH` environment variable, or set `Libxc_DIR` to the directory containing the file.
 

source/input.cpp

@@ -2485,6 +2485,12 @@ void Input::Check(void)
         {
             ModuleBase::WARNING_QUIT("Input", "not ready for linear_scaling method in lcao .");
         }
+        else if (ks_solver == "cusolver")
+        {
+#ifndef __MPI
+            ModuleBase::WARNING_QUIT("Input","Cusolver can not be used for series version.");
+#endif
+        }
         else
         {
             ModuleBase::WARNING_QUIT("Input", "please check the ks_solver parameter!");
@@ -2666,7 +2672,7 @@ void Input::Check(void)
             if (!(calculation == "nscf"))
                 ModuleBase::WARNING_QUIT("Input", "calculate berry phase, please set calculation = nscf");
         }
-        else if (basis_type == "lcao" && (ks_solver == "genelpa" || ks_solver == "scalapack_gvx"))
+        else if (basis_type == "lcao" && ks_solver == "genelpa" || ks_solver == "scalapack_gvx" || ks_solver == "cusolver")
         {
             if (!(calculation == "nscf"))
                 ModuleBase::WARNING_QUIT("Input", "calculate berry phase, please set calculation = nscf");

source/module_base/global_function.h

@@ -326,6 +326,16 @@ double ddot_real(
         const std::complex<double>* psi_R,
         const bool reduce = true) ;
 
+//==========================================================
+// GLOBAL FUNCTION :
+// NAME : IS_COLUMN_MAJOR_KS_SOLVER
+// check ks_solver requires column major or not
+//==========================================================
+static inline bool IS_COLUMN_MAJOR_KS_SOLVER()
+{
+    return GlobalV::KS_SOLVER=="genelpa" || GlobalV::KS_SOLVER=="scalapack_gvx" || GlobalV::KS_SOLVER=="cusolver";
+}
+
 }//namespace GlobalFunc
 }//namespace ModuleBase
 

source/module_deepks/LCAO_deepks_vdelta.cpp

@@ -113,7 +113,7 @@ void LCAO_Deepks::add_v_delta(const UnitCell_pseudo &ucell,
 
                             int iic;
 
-                            if(GlobalV::KS_SOLVER=="genelpa" || GlobalV::KS_SOLVER=="scalapack_gvx")  // save the matrix as column major format
+                            if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER())
                             {
                                 iic=iw1_local+iw2_local*nrow;
                             }
@@ -142,7 +142,7 @@ void LCAO_Deepks::check_v_delta(const int nrow, const int ncol)
 		for (int icol=0;icol<ncol;icol++)
 		{
 			int iic;
-			if(GlobalV::KS_SOLVER=="genelpa" || GlobalV::KS_SOLVER=="scalapack_gvx")  // save the matrix as column major format
+            if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER())
 			{
 				iic=irow+icol*nrow;
 			}
@@ -429,7 +429,7 @@ void LCAO_Deepks::cal_e_delta_band_k(const std::vector<ModuleBase::ComplexMatrix
             if (mu >= 0 && nu >= 0)
             {                
                 int iic;
-                if(GlobalV::KS_SOLVER=="genelpa" || GlobalV::KS_SOLVER=="scalapack_gvx")  // save the matrix as column major format
+                if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER())
                 {
                     iic=mu+nu*nrow;
                 }
@@ -452,4 +452,4 @@ void LCAO_Deepks::cal_e_delta_band_k(const std::vector<ModuleBase::ComplexMatrix
     return;
 }
 
-#endif
+#endif

source/module_deepks/test/nnr.cpp

@@ -212,7 +212,7 @@ void test_deepks::folding_nnr(const Test_Deepks::K_Vectors &kv)
 
 								if(nu<0)continue;
 								int iic;
-								if(GlobalV::KS_SOLVER=="genelpa" || GlobalV::KS_SOLVER=="scalapack_gvx")  // save the matrix as column major format
+                                if (ModuleBase::GlobalFunc::IS_COLUMN_MAJOR_KS_SOLVER() )
 								{
 									iic=mu+nu*ParaO.nrow;
 								}
@@ -232,4 +232,4 @@ void test_deepks::folding_nnr(const Test_Deepks::K_Vectors &kv)
 		} // end T1
 		assert(index==this->nnr);
 	}
-}
+}

source/module_orbital/ORB_control.cpp

@@ -167,7 +167,7 @@ void ORB_control::setup_2d_division(std::ofstream& ofs_running,
 
     // (1) calculate nrow, ncol, nloc.
     if (ks_solver == "genelpa" || ks_solver == "hpseps" || ks_solver == "scalpack"
-        || ks_solver == "selinv" || ks_solver == "scalapack_gvx")
+        || ks_solver == "selinv" || ks_solver == "scalapack_gvx" || ks_solver == "cusolver")
     {
         ofs_running << " divide the H&S matrix using 2D block algorithms." << std::endl;
 #ifdef __MPI

source/module_orbital/parallel_orbitals.cpp

@@ -98,7 +98,7 @@ void ORB_control::set_trace(std::ofstream& ofs_running)
 	}
 #ifdef __MPI
     else if(ks_solver=="scalpack" || ks_solver=="genelpa" || ks_solver=="hpseps" 
-		|| ks_solver=="selinv" || ks_solver=="scalapack_gvx") //xiaohui add 2013-09-02
+		|| ks_solver=="selinv" || ks_solver=="scalapack_gvx" || ks_solver=="cusolver") //xiaohui add 2013-09-02
     {
         // ofs_running << " nrow=" << nrow << std::endl;
         for (int irow=0; irow< pv->nrow; irow++)
@@ -205,6 +205,9 @@ void ORB_control::divide_HS_2d
 	// get the 2D index of computer.
 	pv->dim0 = (int)sqrt((double)dsize); //mohan update 2012/01/13
 	//while (GlobalV::NPROC_IN_POOL%dim0!=0)
+
+    if (ks_solver=="cusolver") pv->dim0 = 1; // Xu Shu add 2022-03-25
+
 	while (dsize%pv->dim0!=0)
 	{
 		pv->dim0 = pv->dim0 - 1;
@@ -227,6 +230,8 @@ void ORB_control::divide_HS_2d
 	{
 		pv->nb = nb2d; // mohan add 2010-06-28
 	}
+
+    if (ks_solver=="cusolver") pv->nb = 1; // Xu Shu add 2022-03-25
 	ModuleBase::GlobalFunc::OUT(ofs_running,"nb2d", pv->nb);
 
     this->set_parameters(ofs_running, ofs_warning);
@@ -252,7 +257,7 @@ void ORB_control::divide_HS_2d
 	pv->nloc = pv->MatrixInfo.col_num * pv->MatrixInfo.row_num;
 
 	// init blacs context for genelpa
-    if(ks_solver=="genelpa" || ks_solver=="scalapack_gvx")
+    if (ks_solver == "genelpa" || ks_solver == "scalapack_gvx" || ks_solver == "cusolver")
     {
         pv->blacs_ctxt = cart2blacs(pv->comm_2D, pv->dim0, pv->dim1,
             nlocal, nbands, pv->nb, pv->nrow, pv->desc, pv->desc_wfc);