deepmodeling · denghuilu · Nov 22, 2022 · Nov 15, 2022 · Nov 16, 2022 · Nov 17, 2022
diff --git a/docs/advanced/acceleration/cuda.md b/docs/advanced/acceleration/cuda.md
@@ -0,0 +1,53 @@
+# CUDA GPU Implementations
+
+In ABACUS, we provide the option to use the GPU devices to accelerate the performance.
+And it has the following general features:
+
+- **Full gpu implementations**: During the SCF progress, `Psi`, `Hamilt`, `Hsolver`, and `DiagCG` classes are stored or calculated by the GPU devices.
+
+- **Electronic state data**: (e.g. electronic density) are moved from the GPU to the CPU(s) every scf step.
+
+- **Acclerated by the NVIDIA libraries**: `cuBLAS` for common linear algebra calculations, `cuSolver` for eigen values/vectors, and `cuFFT` for the conversions between the real and recip spaces.
+
+- **Multi GPU supprted**: Using multiple MPI tasks will often give the best performance. Note each MPI task will be bind to a GPU device with automatically computing load balancing.
+
+- **Parallel strategy**: K point parallel. 
+
+## Required hardware/software
+
+To compile and use ABACUS in CUDA mode, you currently need to have an NVIDIA GPU and install the corresponding NVIDIA CUDA toolkit software on your system (this is only tested on Linux and unsupported on Windows):
+
+- Check if you have an NVIDIA GPU: cat /proc/driver/nvidia/gpus/*/information
+
+- Go to https://developer.nvidia.com/cuda-downloads
+
+- Install a driver and toolkit appropriate for your system (SDK is not necessary)
+
+
+## Building ABACUS with the GPU support:
+
+Check the [Advanced Installation Options](https://abacus-rtd.readthedocs.io/en/latest/advanced/install.html#build-with-cuda-support) for the installation of CUDA version support.
+
+## Run with the GPU support by editing the INPUT script:
+
+In `INPUT` file we need to set the value keyword [device](../input_files/input-main.md#device) to be `gpu`.
+
+## Examples
+We provides [examples](https://github.com/deepmodeling/abacus-develop/tree/develop/examples/gpu) of gpu calculations.
+
+## Known limitations
+
+- Only CG method is supported, so the input keyword `ks_solver` can only take the value `cg`,
+- Only PW basis is supported, so the input keyword `basis_type` can only take the value `pw`,
+- Only k point parallelization is supported, so the input keyword `kpar` will be set to match the number of MPI tasks automatically.
+- Supported CUDA architectures:
+  - 60 # P100, 1080ti
+  - 70 # V100
+  - 75 # T4
+  - 80 # A100, 3090
+
+## FAQ
+```
+Q: Does the GPU implementations support atomic orbital basis sets?
+A: Currently no.
+```
diff --git a/docs/advanced/acceleration/index.rst b/docs/advanced/acceleration/index.rst
@@ -0,0 +1,12 @@
+=================================================
+Accelerate Performance
+=================================================
+
+This section describes various methods for improving ABACUS performance for different classes of problems running on different kinds of devices.
+
+Accelerated versions of CUDA GPU implementations have been added to ABACUS, which will typically run faster than the standard non-accelerated versions. This requires appropriate hardware to be present on your system, e.g. NVIDIA GPUs.
+
+.. toctree::
+   :maxdepth: 2
+
+   cuda
diff --git a/docs/advanced/input_files/input-main.md b/docs/advanced/input_files/input-main.md
@@ -2,7 +2,7 @@
 
 - [System variables](#system-variables)
 
-  [suffix](#suffix) | [ntype](#ntype) | [calculation](#calculation) | [esolver_type](#esolver_type) | [symmetry](#symmetry) | [kpar](#kpar) | [bndpar](#bndpar) | [latname](#latname) | [init_wfc](#init_wfc) | [init_chg](#init_chg) | [init_vel](#init_vel) | [nelec](#nelec) | [tot_magnetization](#tot_magnetization) | [dft_functional](#dft_functional) | [xc_temperature](#xc_temperature) | [pseudo_rcut](#pseudo_rcut) | [pseudo_mesh](#pseudo_mesh) | [mem_saver](#mem_saver) | [diago_proc](#diago_proc) | [nbspline](#nbspline) | [kspacing](#kspacing)  | [min_dist_coef](#min_dist_coef) | [symmetry_prec](#symmetry_prec)
+  [suffix](#suffix) | [ntype](#ntype) | [calculation](#calculation) | [esolver_type](#esolver_type) | [symmetry](#symmetry) | [kpar](#kpar) | [bndpar](#bndpar) | [latname](#latname) | [init_wfc](#init_wfc) | [init_chg](#init_chg) | [init_vel](#init_vel) | [nelec](#nelec) | [tot_magnetization](#tot_magnetization) | [dft_functional](#dft_functional) | [xc_temperature](#xc_temperature) | [pseudo_rcut](#pseudo_rcut) | [pseudo_mesh](#pseudo_mesh) | [mem_saver](#mem_saver) | [diago_proc](#diago_proc) | [nbspline](#nbspline) | [kspacing](#kspacing)  | [min_dist_coef](#min_dist_coef) | [symmetry_prec](#symmetry_prec) | [device](#device)
 - [Variables related to input files](#variables-related-to-input-files)
 
   [stru_file](#stru_file) | [kpoint_file](#kpoint_file) | [pseudo_dir](#pseudo_dir) | [orbital_dir](#orbital_dir) | [read_file_dir](#read_file_dir) | [wannier_card](#wannier_card)
@@ -261,6 +261,19 @@ This part of variables are used to control general system parameters.
 - **Descrption**: The accuracy for symmetry judgment. The unit is Bohr.
 - **Default**: 1.0e-5
 
+### device
+- **Type**: String
+- **Descrption**: Specifies the computing device for ABACUS.
+
+  Available options are:
+  - `cpu`: for CPUs via Intel, AMD, or Other supported CPU devices
+  - `gpu`: for GPUs via CUDA.
+
+  Known limitations: 
+  - `pw basis`: required by the `gpu` acceleraion options
+  - `cg ks_solver`: required by the `gpu` acceleraion options
+- **Default**: `cpu`
+
 [back to top](#full-list-of-input-keywords)
 
 ## Variables related to input files

diff --git a/docs/advanced/install.md b/docs/advanced/install.md
@@ -47,6 +47,17 @@ To build tests for ABACUS, define `BUILD_TESTING` flag. You can also specify pat
 ```bash
 cmake -B build -DBUILD_TESTING=1
 ```
+
+## Build with CUDA support
+
+### Extra prerequisites
+- [CUDA-Toolkit](https://developer.nvidia.com/cuda-toolkit)
+
+To build cuda gpu support for ABACUS, define `USE_CUDA` flag. You can also specify path to local installation of cuda toolkit by setting `CUDA_TOOLKIT_ROOT_DIR` flags.
+```bash
+cmake -B build -DUSE_CUDA=1
+```
+
 ## Build ABACUS with make
 
 > Note: We suggest using CMake to configure and compile.

diff --git a/docs/index.rst b/docs/index.rst
@@ -38,6 +38,7 @@ the developments and applications of novel machine-learning-assisted DFT methods
    advanced/pp_orb
    advanced/opt
    advanced/md
+   advanced/acceleration/index
    advanced/elec_properties/index
    advanced/interface/index
    advanced/input_files/index

diff --git a/examples/gpu/si16_pw/INPUT b/examples/gpu/si16_pw/INPUT
@@ -0,0 +1,25 @@
+INPUT_PARAMETERS
+#Parameters (1.General)
+suffix			autotest
+calculation		scf
+ntype			1
+#nbands			8
+symmetry		1
+device			cpu
+
+#Parameters (2.Iteration)
+ecutwfc			60
+scf_thr			1e-8
+scf_nmax		100
+cal_force 		1
+cal_stress 		1
+#Parameters (3.Basis)
+basis_type		pw
+
+#Parameters (4.Smearing)
+smearing_method		gauss
+smearing_sigma		0.002
+
+#Parameters (5.Mixing)
+mixing_type		pulay
+mixing_beta		0.3
diff --git a/examples/gpu/si16_pw/KPT b/examples/gpu/si16_pw/KPT
@@ -0,0 +1,4 @@
+K_POINTS
+0
+Gamma
+5 5 5 0 0 0
diff --git a/examples/gpu/si16_pw/STRU b/examples/gpu/si16_pw/STRU
@@ -0,0 +1,34 @@
+ATOMIC_SPECIES
+Si 14 ../../../tests/PP_ORB/Si_ONCV_PBE-1.0.upf
+
+LATTICE_CONSTANT
+1
+
+LATTICE_VECTORS
+ 10.2000000  10.2000000   0.0000000
+ 10.2000000   0.0000000  10.2000000
+  0.0000000  10.2000000  10.2000000
+
+ATOMIC_POSITIONS
+Direct
+
+Si
+0.0
+16
+  0.0000000   0.0000000   0.0000000 1 1 1
+  0.1250000   0.1250000   0.1250000 1 1 1
+  0.0000000   0.0000000   0.5000000 1 1 1
+  0.1250000   0.1250000   0.6250000 1 1 1
+  0.0000000   0.5000000   0.0000000 1 1 1
+  0.1250000   0.6250000   0.1250000 1 1 1
+  0.0000000   0.5000000   0.5000000 1 1 1
+  0.1250000   0.6250000   0.6250000 1 1 1
+  0.5000000   0.0000000   0.0000000 1 1 1
+  0.6250000   0.1250000   0.1250000 1 1 1
+  0.5000000   0.0000000   0.5000000 1 1 1
+  0.6250000   0.1250000   0.6250000 1 1 1
+  0.5000000   0.5000000   0.0000000 1 1 1
+  0.6250000   0.6250000   0.1250000 1 1 1
+  0.5000000   0.5000000   0.5000000 1 1 1
+  0.6250000   0.6250000   0.6250000 1 1 1
+
diff --git a/source/input.cpp b/source/input.cpp
@@ -477,6 +477,10 @@ void Input::Default(void)
     of_read_kernel = false;
     of_kernel_file = "WTkernel.txt";
 
+    //==========================================================
+    //    OFDFT sunliang added on 2022-11-15
+    //==========================================================
+    device = "cpu";
     return;
 }
 
@@ -1751,6 +1755,12 @@ bool Input::Read(const std::string &fn)
             read_value(ifs, of_kernel_file);
         }
         //----------------------------------------------------------------------------------
+        //    device control denghui added on 2022-11-05
+        //----------------------------------------------------------------------------------     
+        else if (strcmp("device", word) == 0) {
+            read_value(ifs, device);
+        }
+        //----------------------------------------------------------------------------------
         else
         {
             // xiaohui add 2015-09-15
@@ -2505,6 +2515,10 @@ void Input::Bcast()
     Parallel_Common::bcast_int(of_full_pw_dim);
     Parallel_Common::bcast_bool(of_read_kernel);
     Parallel_Common::bcast_string(of_kernel_file);
+    //----------------------------------------------------------------------------------
+    //    device control denghui added on 2022-11-05
+    //----------------------------------------------------------------------------------
+    Parallel_Common::bcast_string(device);
 
     return;
 }

diff --git a/source/input.h b/source/input.h
@@ -463,6 +463,11 @@ class Input
     bool of_read_kernel; // If set to 1, the kernel of WT KEDF will be filled from file of_kernel_file, not from formula. Only usable for WT KEDF.
     string of_kernel_file; // The name of WT kernel file.
 
+    //==========================================================
+    //    device control denghui added on 2022-11-15
+    //==========================================================
+    std::string device;
+
     //==========================================================
     // variables for test only
     //==========================================================

diff --git a/source/input_conv.cpp b/source/input_conv.cpp
@@ -30,6 +30,7 @@
 #include "module_hsolver/hsolver_lcao.h"
 #include "module_elecstate/potentials/efield.h"
 #include "module_elecstate/potentials/gatefield.h"
+#include "module_psi/include/device.h"
 
 void Input_Conv::Convert(void)
 {
@@ -81,18 +82,19 @@ void Input_Conv::Convert(void)
     GlobalV::NBANDS = INPUT.nbands;
     GlobalC::wf.pw_seed = INPUT.pw_seed;
     GlobalV::NBANDS_ISTATE = INPUT.nbands_istate;
-#if ((defined __CUDA) || (defined __ROCM))
-    int temp_nproc;
-    MPI_Comm_size(MPI_COMM_WORLD, &temp_nproc);
-    if (temp_nproc != INPUT.kpar)
-    {
-        ModuleBase::WARNING("Input_conv", "None kpar set in INPUT file, auto set kpar value.");
+    GlobalV::device_flag = 
+        psi::device::get_device_flag(
+            INPUT.device,
+            INPUT.ks_solver, 
+            INPUT.basis_type);
+
+    if (GlobalV::device_flag == "gpu") {
+        GlobalV::KPAR = psi::device::get_device_kpar(INPUT.kpar);
+    }
+    else {
+        GlobalV::KPAR = INPUT.kpar;
+        GlobalV::NSTOGROUP = INPUT.bndpar;
     }
-    GlobalV::KPAR = temp_nproc;
-#else
-    GlobalV::KPAR = INPUT.kpar;
-    GlobalV::NSTOGROUP = INPUT.bndpar;
-#endif
     GlobalV::CALCULATION = INPUT.calculation;
     GlobalV::ESOLVER_TYPE = INPUT.esolver_type;
 

diff --git a/source/module_base/global_variable.cpp b/source/module_base/global_variable.cpp
@@ -239,6 +239,8 @@ int of_full_pw_dim = 0;
 bool of_read_kernel = false;
 std::string of_kernel_file = "WTkernel.txt";
 
+// denghui added for device flag
+std::string device_flag = "unknown";
 std::string chg_extrap = "";
 int out_pot = 0;
 

diff --git a/source/module_base/global_variable.h b/source/module_base/global_variable.h
@@ -268,6 +268,10 @@ extern int of_full_pw_dim;  // If of_full_pw = 1, the dimention of FFT will be t
 extern bool of_read_kernel; // If set to 1, the kernel of WT KEDF will be filled from file of_kernel_file, not from formula. Only usable for WT KEDF.
 extern std::string of_kernel_file; // The name of WT kernel file.
 
+//==========================================================
+// device flags
+//==========================================================
+extern std::string device_flag;
 extern std::string chg_extrap;
 extern int out_pot;