diff --git a/docs/advanced/acceleration/cuda.md b/docs/advanced/acceleration/cuda.md new file mode 100644 index 00000000000..8836f910139 --- /dev/null +++ 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 new file mode 100644 index 00000000000..740d05de9ab --- /dev/null +++ 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 \ No newline at end of file diff --git a/docs/advanced/input_files/input-main.md b/docs/advanced/input_files/input-main.md index a23e1707f43..810761dc1ef 100644 --- 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 index 0d4b2c75d5e..8c7dc44cfdf 100644 --- 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 index 596eac40837..ac3ef1c24e0 100644 --- 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 new file mode 100644 index 00000000000..df01c6a69f3 --- /dev/null +++ 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 new file mode 100644 index 00000000000..abc57daebed --- /dev/null +++ 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 new file mode 100644 index 00000000000..05e97c427f4 --- /dev/null +++ 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 index e7f2a7c720c..b853e022ce0 100644 --- 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 index 577a84a5cf2..f922bdde297 100644 --- 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 index c93b013fb9f..6df294de8cb 100644 --- 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 index 689035e2763..b706ef37cf1 100644 --- 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 index 8d9afdda652..176ea63d106 100644 --- 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; diff --git a/source/module_elecstate/elecstate_pw.cpp b/source/module_elecstate/elecstate_pw.cpp index 3114a739bd0..bb63a118b2d 100644 --- a/source/module_elecstate/elecstate_pw.cpp +++ b/source/module_elecstate/elecstate_pw.cpp @@ -13,20 +13,20 @@ template ElecStatePW::ElecStatePW(ModulePW::PW_Basis_K *wfc_basis_in, Charge* chg_in, K_Vectors *pkv_in, int nbands_in) : basis(wfc_basis_in) { this->classname = "ElecStatePW"; - init(chg_in, pkv_in, pkv_in->nks, nbands_in); + this->init(chg_in, pkv_in, pkv_in->nks, nbands_in); } template ElecStatePW::~ElecStatePW() { if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { - delete_memory_var_op()(this->ctx, this->rho_data); + delmem_var_op()(this->ctx, this->rho_data); if (XC_Functional::get_func_type() == 3) { - delete_memory_var_op()(this->ctx, this->kin_r_data); + delmem_var_op()(this->ctx, this->kin_r_data); } } - delete_memory_complex_op()(this->ctx, this->wfcr); - delete_memory_complex_op()(this->ctx, this->wfcr_another_spin); + delmem_complex_op()(this->ctx, this->wfcr); + delmem_complex_op()(this->ctx, this->wfcr_another_spin); } template @@ -34,13 +34,13 @@ void ElecStatePW::init_rho_data() { if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { this->rho = new FPTYPE*[this->charge->nspin]; - resize_memory_var_op()(this->ctx, this->rho_data, this->charge->nspin * this->charge->nrxx); + resmem_var_op()(this->ctx, this->rho_data, this->charge->nspin * this->charge->nrxx); for (int ii = 0; ii < this->charge->nspin; ii++) { this->rho[ii] = this->rho_data + ii * this->charge->nrxx; } if (XC_Functional::get_func_type() == 3) { this->kin_r = new FPTYPE*[this->charge->nspin]; - resize_memory_var_op()(this->ctx, this->kin_r_data, this->charge->nspin * this->charge->nrxx); + resmem_var_op()(this->ctx, this->kin_r_data, this->charge->nspin * this->charge->nrxx); for (int ii = 0; ii < this->charge->nspin; ii++) { this->kin_r[ii] = this->kin_r_data + ii * this->charge->nrxx; } @@ -52,8 +52,8 @@ void ElecStatePW::init_rho_data() this->kin_r = this->charge->kin_r; } } - resize_memory_complex_op()(this->ctx, this->wfcr, this->basis->nmaxgr); - resize_memory_complex_op()(this->ctx, this->wfcr_another_spin, this->charge->nrxx); + resmem_complex_op()(this->ctx, this->wfcr, this->basis->nmaxgr); + resmem_complex_op()(this->ctx, this->wfcr_another_spin, this->charge->nrxx); this->init_rho = true; } @@ -74,11 +74,11 @@ void ElecStatePW::psiToRho(const psi::Psi, { // denghui replaced at 20221110 // ModuleBase::GlobalFunc::ZEROS(this->rho[is], this->charge->nrxx); - set_memory_var_op()(this->ctx, this->rho[is], 0, this->charge->nrxx); + setmem_var_op()(this->ctx, this->rho[is], 0, this->charge->nrxx); if (XC_Functional::get_func_type() == 3) { // ModuleBase::GlobalFunc::ZEROS(this->charge->kin_r[is], this->charge->nrxx); - set_memory_var_op()(this->ctx, this->kin_r[is], 0, this->charge->nrxx); + setmem_var_op()(this->ctx, this->kin_r[is], 0, this->charge->nrxx); } } @@ -87,6 +87,16 @@ void ElecStatePW::psiToRho(const psi::Psi, psi.fix_k(ik); this->updateRhoK(psi); } + if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { + for (int ii = 0; ii < GlobalV::NSPIN; ii++) { + syncmem_var_d2h_op()( + this->cpu_ctx, + this->ctx, + this->charge->rho[ii], + this->rho[ii], + this->charge->nrxx); + } + } this->parallelK(); ModuleBase::timer::tick("ElecStatePW", "psiToRho"); } @@ -97,16 +107,11 @@ void ElecStatePW::updateRhoK(const psi::Psi this->rhoBandK(psi); } -/*void ElecStatePW::getNewRho() -{ - return; -}*/ - template void ElecStatePW::parallelK() { #ifdef __MPI - charge->rho_mpi(); + this->charge->rho_mpi(); if(GlobalV::ESOLVER_TYPE == "sdft") //qinarui add it 2021-7-21 { this->eband /= GlobalV::NPROC_IN_POOL; @@ -144,14 +149,14 @@ void ElecStatePW::rhoBandK(const psi::Psi, if (GlobalV::NSPIN == 4) { int npwx = npw / 2; - npw = this->klist->ngk[ik]; for (int ibnd = 0; ibnd < nbands; ibnd++) { /// /// only occupied band should be calculated. /// - if (this->wg(ik, ibnd) < ModuleBase::threshold_wg) + if (this->wg(ik, ibnd) < ModuleBase::threshold_wg) { continue; + } this->basis->recip_to_real(this->ctx, &psi(ibnd,0), this->wfcr, ik); @@ -210,8 +215,9 @@ void ElecStatePW::rhoBandK(const psi::Psi, /// /// only occupied band should be calculated. /// - if (this->wg(ik, ibnd) < ModuleBase::threshold_wg) + if (this->wg(ik, ibnd) < ModuleBase::threshold_wg) { continue; + } this->basis->recip_to_real(this->ctx, &psi(ibnd,0), this->wfcr, ik); diff --git a/source/module_elecstate/elecstate_pw.h b/source/module_elecstate/elecstate_pw.h index 57118079bb9..6eaec83888f 100644 --- a/source/module_elecstate/elecstate_pw.h +++ b/source/module_elecstate/elecstate_pw.h @@ -38,17 +38,21 @@ class ElecStatePW : public ElecState void init_rho_data(); Device * ctx = {}; + psi::DEVICE_CPU * cpu_ctx = {}; bool init_rho = false; FPTYPE ** rho = nullptr, ** kin_r = nullptr; FPTYPE * rho_data = nullptr, * kin_r_data = nullptr; std::complex *wfcr = nullptr, *wfcr_another_spin = nullptr; using elecstate_pw_op = elecstate::elecstate_pw_op; - using set_memory_var_op = psi::memory::set_memory_op; - using resize_memory_var_op = psi::memory::resize_memory_op; - using delete_memory_var_op = psi::memory::delete_memory_op; - using resize_memory_complex_op = psi::memory::resize_memory_op, Device>; - using delete_memory_complex_op = psi::memory::delete_memory_op, Device>; + + using setmem_var_op = psi::memory::set_memory_op; + using resmem_var_op = psi::memory::resize_memory_op; + using delmem_var_op = psi::memory::delete_memory_op; + using syncmem_var_d2h_op = psi::memory::synchronize_memory_op; + + using resmem_complex_op = psi::memory::resize_memory_op, Device>; + using delmem_complex_op = psi::memory::delete_memory_op, Device>; }; } // namespace elecstate diff --git a/source/module_elecstate/potentials/potential_new.cpp b/source/module_elecstate/potentials/potential_new.cpp index 0412df10d2e..185dd2eec7f 100644 --- a/source/module_elecstate/potentials/potential_new.cpp +++ b/source/module_elecstate/potentials/potential_new.cpp @@ -47,6 +47,11 @@ Potential::~Potential() } this->components.clear(); } + #if (defined(__CUDA) || defined(__ROCM)) + if (GlobalV::device_flag == "gpu") { + delmem_var_op()(this->gpu_ctx, d_v_effective); + } + #endif } void Potential::pot_register(std::vector& components_list) @@ -124,6 +129,11 @@ void Potential::allocate() this->vofk_effective.create(GlobalV::NSPIN, nrxx); ModuleBase::Memory::record("Potential", "vofk", GlobalV::NSPIN * nrxx, "double"); } + #if (defined(__CUDA) || defined(__ROCM)) + if (GlobalV::device_flag == "gpu") { + resmem_var_op()(this->gpu_ctx, d_v_effective, GlobalV::NSPIN * nrxx); + } + #endif } void Potential::update_from_charge(const Charge* chg, const UnitCell* ucell) @@ -138,6 +148,12 @@ void Potential::update_from_charge(const Charge* chg, const UnitCell* ucell) this->cal_v_eff(chg, ucell, this->v_effective); + #if (defined(__CUDA) || defined(__ROCM)) + if (GlobalV::device_flag == "gpu") { + syncmem_var_h2d_op()(this->gpu_ctx, this->cpu_ctx, d_v_effective, this->v_effective.c, this->v_effective.nr * this->v_effective.nc); + } + #endif + ModuleBase::timer::tick("Potential", "update_from_charge"); } diff --git a/source/module_elecstate/potentials/potential_new.h b/source/module_elecstate/potentials/potential_new.h index 15f400a9dac..5904def711b 100644 --- a/source/module_elecstate/potentials/potential_new.h +++ b/source/module_elecstate/potentials/potential_new.h @@ -76,6 +76,38 @@ class Potential : public PotBase { return this->v_effective; } + double* get_effective_v_data() + { + if (this->v_effective.nc > 0) + { + #if (defined(__CUDA) || defined(__ROCM)) + if (GlobalV::device_flag == "gpu") { + return this->d_v_effective; + } + #endif + return &(this->v_effective(0, 0)); + } + else + { + return nullptr; + } + } + const double* get_effective_v_data() const + { + if (this->v_effective.nc > 0) + { + #if (defined(__CUDA) || defined(__ROCM)) + if (GlobalV::device_flag == "gpu") { + return this->d_v_effective; + } + #endif + return &(this->v_effective(0, 0)); + } + else + { + return nullptr; + } + } double* get_effective_v(int is) { if (this->v_effective.nc > 0) @@ -159,6 +191,15 @@ class Potential : public PotBase std::vector v_effective_fixed; ModuleBase::matrix v_effective; +#if (defined(__CUDA) || defined(__ROCM)) + double * d_v_effective = nullptr; + psi::DEVICE_CPU * cpu_ctx = nullptr; + psi::DEVICE_GPU * gpu_ctx = nullptr; + using resmem_var_op = psi::memory::resize_memory_op; + using delmem_var_op = psi::memory::delete_memory_op; + using syncmem_var_h2d_op = psi::memory::synchronize_memory_op; +#endif + ModuleBase::matrix vofk_effective; bool fixed_done = false; diff --git a/source/module_esolver/esolver.cpp b/source/module_esolver/esolver.cpp index 0b2e9deb907..43b56533470 100644 --- a/source/module_esolver/esolver.cpp +++ b/source/module_esolver/esolver.cpp @@ -92,7 +92,13 @@ namespace ModuleESolver //initialize the corresponding Esolver child class if (esolver_type == "ksdft_pw") { - p_esolver = new ESolver_KS_PW(); + #if ((defined __CUDA) || (defined __ROCM)) + if (GlobalV::device_flag == "gpu") { + p_esolver = new ESolver_KS_PW(); + return; + } + #endif + p_esolver = new ESolver_KS_PW(); } #ifdef __LCAO else if (esolver_type == "ksdft_lcao") diff --git a/source/module_esolver/esolver_ks.cpp b/source/module_esolver/esolver_ks.cpp index a9f6b0de2cd..5f9512a8944 100644 --- a/source/module_esolver/esolver_ks.cpp +++ b/source/module_esolver/esolver_ks.cpp @@ -18,7 +18,8 @@ namespace ModuleESolver { - ESolver_KS::ESolver_KS() + template + ESolver_KS::ESolver_KS() { classname = "ESolver_KS"; basisname = "PLEASE ADD BASISNAME FOR CURRENT ESOLVER."; @@ -36,14 +37,16 @@ namespace ModuleESolver tmp->setbxyz(INPUT.bx,INPUT.by,INPUT.bz); } - ESolver_KS::~ESolver_KS() + template + ESolver_KS::~ESolver_KS() { delete this->pw_wfc; delete this->p_hamilt; delete this->phsol; } - void ESolver_KS::Init(Input& inp, UnitCell& ucell) + template + void ESolver_KS::Init(Input& inp, UnitCell& ucell) { ESolver_FP::Init(inp,ucell); chr.cal_nelec(); @@ -117,7 +120,8 @@ namespace ModuleESolver CE.Init_CE(); } - void ESolver_KS::hamilt2density(const int istep, const int iter, const double ethr) + template + void ESolver_KS::hamilt2density(const int istep, const int iter, const FPTYPE ethr) { ModuleBase::timer::tick(this->classname, "hamilt2density"); //Temporarily, before HSolver is constructed, it should be overrided by @@ -128,7 +132,8 @@ namespace ModuleESolver ModuleBase::timer::tick(this->classname, "hamilt2density"); } - void ESolver_KS::print_wfcfft(Input& inp, ofstream &ofs) + template + void ESolver_KS::print_wfcfft(Input& inp, ofstream &ofs) { ofs << "\n\n\n\n"; ofs << " >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>" << std::endl; @@ -143,7 +148,7 @@ namespace ModuleESolver ofs << " <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<" << std::endl; ofs << "\n\n\n\n"; ofs << "\n SETUP PLANE WAVES FOR WAVE FUNCTIONS" << std::endl; - double ecut = INPUT.ecutwfc; + FPTYPE ecut = INPUT.ecutwfc; if(abs(ecut-this->pw_wfc->gk_ecut * this->pw_wfc->tpiba2) > 1e-6) { ecut = this->pw_wfc->gk_ecut * this->pw_wfc->tpiba2; @@ -166,7 +171,8 @@ namespace ModuleESolver ModuleBase::GlobalFunc::DONE(ofs, "INIT PLANEWAVE"); } - void ESolver_KS::Run(const int istep, UnitCell& ucell) + template + void ESolver_KS::Run(const int istep, UnitCell& ucell) { if (!(GlobalV::CALCULATION == "scf" || GlobalV::CALCULATION == "md" || GlobalV::CALCULATION == "relax" || GlobalV::CALCULATION == "cell-relax")) @@ -192,7 +198,7 @@ namespace ModuleESolver #else auto iterstart = std::chrono::system_clock::now(); #endif - double diag_ethr = this->phsol->set_diagethr(istep, iter, drho); + FPTYPE diag_ethr = this->phsol->set_diagethr(istep, iter, drho); eachiterinit(istep, iter); this->hamilt2density(istep, iter, diag_ethr); @@ -203,10 +209,11 @@ namespace ModuleESolver //they do not occupy all processors, for example wavefunctions uses 20 processors while density uses 10. if(GlobalV::MY_STOGROUP == 0) { - // double drho = this->estate.caldr2(); + // FPTYPE drho = this->estate.caldr2(); // EState should be used after it is constructed. + drho = GlobalC::CHR_MIX.get_drho(pelec->charge, GlobalV::nelec); - double hsolver_error = 0.0; + FPTYPE hsolver_error = 0.0; if (firstscf) { firstscf = false; @@ -246,9 +253,9 @@ namespace ModuleESolver updatepot(istep, iter); eachiterfinish(iter); #ifdef __MPI - double duration = (double)(MPI_Wtime() - iterstart); + FPTYPE duration = (FPTYPE)(MPI_Wtime() - iterstart); #else - double duration = (std::chrono::system_clock::now() - iterstart).count() / CLOCKS_PER_SEC; + FPTYPE duration = (std::chrono::system_clock::now() - iterstart).count() / CLOCKS_PER_SEC; #endif printiter(iter, drho, duration, diag_ethr); if (this->conv_elec) @@ -266,7 +273,8 @@ namespace ModuleESolver return; }; - void ESolver_KS::printhead() + template + void ESolver_KS::printhead() { std::cout << " " << std::setw(7) << "ITER"; if (GlobalV::NSPIN == 2) @@ -280,12 +288,14 @@ namespace ModuleESolver std::cout << std::setw(11) << "TIME(s)" << std::endl; } - void ESolver_KS::printiter(const int iter, const double drho, const double duration, const double ethr) + template + void ESolver_KS::printiter(const int iter, const FPTYPE drho, const FPTYPE duration, const FPTYPE ethr) { GlobalC::en.print_etot(this->conv_elec, iter, drho, duration, ethr); } - void ESolver_KS::writehead(std::ofstream& ofs_running, const int istep, const int iter) + template + void ESolver_KS::writehead(std::ofstream& ofs_running, const int istep, const int iter) { ofs_running << "\n " @@ -295,10 +305,14 @@ namespace ModuleESolver << "--------------------------------\n"; } - int ESolver_KS::getniter() + template + int ESolver_KS::getniter() { return this->niter; } - +template class ESolver_KS; +#if ((defined __CUDA) || (defined __ROCM)) +template class ESolver_KS; +#endif } \ No newline at end of file diff --git a/source/module_esolver/esolver_ks.h b/source/module_esolver/esolver_ks.h index 245790515f0..ee22377758d 100644 --- a/source/module_esolver/esolver_ks.h +++ b/source/module_esolver/esolver_ks.h @@ -13,14 +13,15 @@ namespace ModuleESolver { + template class ESolver_KS : public ESolver_FP { public: ESolver_KS(); virtual ~ESolver_KS(); // HSolver* phsol; - double scf_thr; // scf threshold - double drho; // the difference between rho_in (before HSolver) and rho_out (After HSolver) + FPTYPE scf_thr; // scf threshold + FPTYPE drho; // the difference between rho_in (before HSolver) and rho_out (After HSolver) int maxniter; // maximum iter steps for scf int niter; // iter steps actually used in scf bool conv_elec; // If electron density is converged in scf. @@ -30,10 +31,10 @@ namespace ModuleESolver virtual void Run(const int istep, UnitCell& cell) override; // calculate electron density from a specific Hamiltonian - virtual void hamilt2density(const int istep, const int iter, const double ethr); + virtual void hamilt2density(const int istep, const int iter, const FPTYPE ethr); // calculate electron states from a specific Hamiltonian - virtual void hamilt2estates(const double ethr){}; + virtual void hamilt2estates(const FPTYPE ethr){}; // get current step of Ionic simulation virtual int getniter() override; @@ -59,17 +60,17 @@ namespace ModuleESolver void printhead(); // Print inforamtion in each iter // G1 -3.435545e+03 0.000000e+00 3.607e-01 2.862e-01 - void printiter(const int iter, const double drho, const double duration, const double ethr); + void printiter(const int iter, const FPTYPE drho, const FPTYPE duration, const FPTYPE ethr); // Write the headline in the running_log file // "PW/LCAO" ALGORITHM --------------- ION= 1 ELEC= 1-------------------------------- void writehead(std::ofstream& ofs_running, const int istep, const int iter); +// TODO: control single precision at input files - hsolver::HSolver* phsol = nullptr; - hamilt::Hamilt* p_hamilt = nullptr; - ModulePW::PW_Basis_K* pw_wfc = nullptr; - - Charge_Extra CE; + hsolver::HSolver* phsol = nullptr; + hamilt::Hamilt* p_hamilt = nullptr; + ModulePW::PW_Basis_K* pw_wfc = nullptr; + Charge_Extra CE; protected: std::string basisname; //PW or LCAO diff --git a/source/module_esolver/esolver_ks_lcao.h b/source/module_esolver/esolver_ks_lcao.h index 1dd817a0c35..1a7c6f3e3d0 100644 --- a/source/module_esolver/esolver_ks_lcao.h +++ b/source/module_esolver/esolver_ks_lcao.h @@ -11,7 +11,7 @@ namespace ModuleESolver { - class ESolver_KS_LCAO : public ESolver_KS + class ESolver_KS_LCAO : public ESolver_KS { public: ESolver_KS_LCAO(); diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp index 14ecdff98f1..365f2a572dc 100644 --- a/source/module_esolver/esolver_ks_pw.cpp +++ b/source/module_esolver/esolver_ks_pw.cpp @@ -4,10 +4,6 @@ //--------------temporary---------------------------- #include "../src_pw/global.h" -#include "../module_base/global_function.h" -#include "../module_symmetry/symmetry.h" -#include "../src_pw/pw_complement.h" -#include "../src_pw/structure_factor.h" #include "../src_pw/symmetry_rho.h" #include "../src_io/print_info.h" #include "../src_pw/H_Ewald_pw.h" @@ -34,37 +30,55 @@ #include "src_io/numerical_basis.h" #include "src_io/to_wannier90.h" #include "src_io/berryphase.h" +#include "module_psi/include/device.h" +#include "module_hsolver/include/math_kernel.h" namespace ModuleESolver { - ESolver_KS_PW::ESolver_KS_PW() + template + ESolver_KS_PW::ESolver_KS_PW() { - classname = "ESolver_KS_PW"; - basisname = "PW"; + this->classname = "ESolver_KS_PW"; + this->basisname = "PW"; + this->device = psi::device::get_device_type(this->ctx); + #if ((defined __CUDA) || (defined __ROCM)) + if (this->device == psi::GpuDevice) { + hsolver::createBLAShandle(); + } + #endif } - ESolver_KS_PW::~ESolver_KS_PW() + + template + ESolver_KS_PW::~ESolver_KS_PW() { //delete HSolver and ElecState if(this->phsol != nullptr) { - delete (hsolver::HSolverPW*)this->phsol; + delete reinterpret_cast*>(this->phsol); this->phsol = nullptr; } if(this->pelec != nullptr) { - delete (elecstate::ElecStatePW*)this->pelec; + delete reinterpret_cast*>(this->pelec); this->pelec = nullptr; } //delete Hamilt if(this->p_hamilt != nullptr) { - delete (hamilt::HamiltPW*)this->p_hamilt; + delete reinterpret_cast*>(this->p_hamilt); this->p_hamilt = nullptr; } + #if ((defined __CUDA) || (defined __ROCM)) + if (this->device == psi::GpuDevice) { + delete reinterpret_cast, Device>*>(this->kspw_psi); + hsolver::destoryBLAShandle(); + } + #endif } - void ESolver_KS_PW::Init_GlobalC(Input& inp, UnitCell& cell) + template + void ESolver_KS_PW::Init_GlobalC(Input& inp, UnitCell& cell) { this->psi = GlobalC::wf.allocate(GlobalC::kv.nks); @@ -114,23 +128,30 @@ namespace ModuleESolver GlobalC::wf.wfcinit(this->psi); } + // denghui added 20221116 + this->kspw_psi = this->device == psi::GpuDevice ? + new psi::Psi, Device>(this->psi[0]) : + reinterpret_cast, Device>*> (this->psi); + ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "INIT BASIS"); } - void ESolver_KS_PW::Init(Input& inp, UnitCell& ucell) + + template + void ESolver_KS_PW::Init(Input& inp, UnitCell& ucell) { - ESolver_KS::Init(inp,ucell); + ESolver_KS::Init(inp,ucell); //init HSolver if(this->phsol == nullptr) { - this->phsol = new hsolver::HSolverPW(GlobalC::wfcpw); + this->phsol = new hsolver::HSolverPW(GlobalC::wfcpw); } //init ElecState, if(this->pelec == nullptr) { - this->pelec = new elecstate::ElecStatePW( GlobalC::wfcpw, &(chr), (K_Vectors*)(&(GlobalC::kv)), GlobalV::NBANDS); + this->pelec = new elecstate::ElecStatePW( GlobalC::wfcpw, &(this->chr), (K_Vectors*)(&(GlobalC::kv)), GlobalV::NBANDS); } // Inititlize the charge density. @@ -159,16 +180,17 @@ namespace ModuleESolver } } - void ESolver_KS_PW::beforescf(int istep) + template + void ESolver_KS_PW::beforescf(int istep) { ModuleBase::TITLE("ESolver_KS_PW", "beforescf"); // Temporary, md and relax will merge later liuyu add 2022-11-07 if(GlobalV::CALCULATION == "md" && istep) { - CE.update_istep(); - CE.save_pos_next(GlobalC::ucell); - CE.extrapolate_charge(pelec->charge); + this->CE.update_istep(); + this->CE.save_pos_next(GlobalC::ucell); + this->CE.extrapolate_charge(this->pelec->charge); if(GlobalC::ucell.cell_parameter_updated) { @@ -184,10 +206,10 @@ namespace ModuleESolver { GlobalV::ofs_running << " Setup the extrapolated charge." << std::endl; // charge extrapolation if istep>0. - CE.update_istep(); - CE.update_all_pos(GlobalC::ucell); - CE.extrapolate_charge(pelec->charge); - CE.save_pos_next(GlobalC::ucell); + this->CE.update_istep(); + this->CE.update_all_pos(GlobalC::ucell); + this->CE.extrapolate_charge(this->pelec->charge); + this->CE.save_pos_next(GlobalC::ucell); GlobalV::ofs_running << " Setup the Vl+Vh+Vxc according to new structure factor and new charge." << std::endl; // calculate the new potential accordint to @@ -208,13 +230,13 @@ namespace ModuleESolver //delete Hamilt if not first scf if(this->p_hamilt != nullptr) { - delete (hamilt::HamiltPW*)this->p_hamilt; + delete reinterpret_cast*>(this->p_hamilt); this->p_hamilt = nullptr; } //allocate HamiltPW if(this->p_hamilt == nullptr) { - this->p_hamilt = new hamilt::HamiltPW(this->pelec->pot); + this->p_hamilt = new hamilt::HamiltPW(this->pelec->pot); } //---------------------------------------------------------- @@ -240,12 +262,13 @@ namespace ModuleESolver Symmetry_rho srho; for (int is = 0; is < GlobalV::NSPIN; is++) { - srho.begin(is, *(pelec->charge), GlobalC::rhopw, GlobalC::Pgrid, GlobalC::symm); + srho.begin(is, *(this->pelec->charge), GlobalC::rhopw, GlobalC::Pgrid, GlobalC::symm); } } - void ESolver_KS_PW::othercalculation(const int istep) + template + void ESolver_KS_PW::othercalculation(const int istep) { ModuleBase::TITLE("ESolver_KS_PW", "othercalculation"); ModuleBase::timer::tick("ESolver_KS_PW", "othercalculation"); @@ -278,7 +301,8 @@ namespace ModuleESolver return; } - void ESolver_KS_PW::eachiterinit(const int istep, const int iter) + template + void ESolver_KS_PW::eachiterinit(const int istep, const int iter) { // mohan add 2010-07-16 if (iter == 1) GlobalC::CHR_MIX.reset(); @@ -291,12 +315,13 @@ namespace ModuleESolver // prepared fox mixing. if(GlobalV::MY_STOGROUP == 0) { - pelec->charge->save_rho_before_sum_band(); + this->pelec->charge->save_rho_before_sum_band(); } } //Temporary, it should be replaced by hsolver later. - void ESolver_KS_PW:: hamilt2density(const int istep, const int iter, const double ethr) + template + void ESolver_KS_PW::hamilt2density(const int istep, const int iter, const FPTYPE ethr) { if(this->phsol != nullptr) { @@ -310,17 +335,16 @@ namespace ModuleESolver // be careful that istep start from 0 and iter start from 1 if((istep==0||istep==1)&&iter==1) { - hsolver::DiagoIterAssist::need_subspace = false; + hsolver::DiagoIterAssist::need_subspace = false; } else { - hsolver::DiagoIterAssist::need_subspace = true; + hsolver::DiagoIterAssist::need_subspace = true; } - hsolver::DiagoIterAssist::PW_DIAG_THR = ethr; - hsolver::DiagoIterAssist::PW_DIAG_NMAX = GlobalV::PW_DIAG_NMAX; - this->phsol->solve(this->p_hamilt, this->psi[0], this->pelec, GlobalV::KS_SOLVER); - + hsolver::DiagoIterAssist::PW_DIAG_THR = ethr; + hsolver::DiagoIterAssist::PW_DIAG_NMAX = GlobalV::PW_DIAG_NMAX; + this->phsol->solve(this->p_hamilt, this->kspw_psi[0], this->pelec, GlobalV::KS_SOLVER); // transform energy for print GlobalC::en.eband = this->pelec->eband; GlobalC::en.demet = this->pelec->demet; @@ -344,11 +368,11 @@ namespace ModuleESolver Symmetry_rho srho; for (int is = 0; is < GlobalV::NSPIN; is++) { - srho.begin(is, *(pelec->charge), GlobalC::rhopw, GlobalC::Pgrid, GlobalC::symm); + srho.begin(is, *(this->pelec->charge), GlobalC::rhopw, GlobalC::Pgrid, GlobalC::symm); } // compute magnetization, only for LSDA(spin==2) - GlobalC::ucell.magnet.compute_magnetization(pelec->charge); + GlobalC::ucell.magnet.compute_magnetization(this->pelec->charge); // deband is calculated from "output" charge density calculated // in sum_band // need 'rho(out)' and 'vr (v_h(in) and v_xc(in))' @@ -358,7 +382,8 @@ namespace ModuleESolver } //Temporary, it should be rewritten with Hamilt class. - void ESolver_KS_PW::updatepot(const int istep, const int iter) + template + void ESolver_KS_PW::updatepot(const int istep, const int iter) { if (!this->conv_elec) { @@ -371,7 +396,8 @@ namespace ModuleESolver } } - void ESolver_KS_PW::eachiterfinish(const int iter) + template + void ESolver_KS_PW::eachiterfinish(const int iter) { //print_eigenvalue(GlobalV::ofs_running); GlobalC::en.calculate_etot(); @@ -395,9 +421,9 @@ namespace ModuleESolver std::stringstream ssc; std::stringstream ss1; ssc << GlobalV::global_out_dir << "tmp" << "_SPIN" << is + 1 << "_CHG"; - pelec->charge->write_rho(pelec->charge->rho_save[is], is, iter, ssc.str(), 3);//mohan add 2007-10-17 + this->pelec->charge->write_rho(this->pelec->charge->rho_save[is], is, iter, ssc.str(), 3);//mohan add 2007-10-17 ss1 << GlobalV::global_out_dir << "tmp" << "_SPIN" << is + 1 << "_CHG.cube"; - pelec->charge->write_rho_cube(pelec->charge->rho_save[is], is, ss1.str(), 3); + this->pelec->charge->write_rho_cube(this->pelec->charge->rho_save[is], is, ss1.str(), 3); } } //output wavefunctions @@ -410,17 +436,14 @@ namespace ModuleESolver WF_io::write_wfc(ssw.str(), this->psi[0], &GlobalC::kv, GlobalC::wfcpw); //ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running,"write wave functions into file WAVEFUNC.dat"); } - } - - } - - void ESolver_KS_PW::afterscf(const int istep) + template + void ESolver_KS_PW::afterscf(const int istep) { // Temporary liuyu add 2022-11-07 - CE.update_all_pos(GlobalC::ucell); + this->CE.update_all_pos(GlobalC::ucell); #ifdef __LCAO if (GlobalC::chi0_hilbert.epsilon) // pengfei 2016-11-23 @@ -458,8 +481,8 @@ namespace ModuleESolver std::stringstream ss1; ssc << GlobalV::global_out_dir << "SPIN" << is + 1 << "_CHG"; ss1 << GlobalV::global_out_dir << "SPIN" << is + 1 << "_CHG.cube"; - pelec->charge->write_rho(pelec->charge->rho_save[is], is, 0, ssc.str());//mohan add 2007-10-17 - pelec->charge->write_rho_cube(pelec->charge->rho_save[is], is, ss1.str(), 3); + this->pelec->charge->write_rho(this->pelec->charge->rho_save[is], is, 0, ssc.str());//mohan add 2007-10-17 + this->pelec->charge->write_rho_cube(this->pelec->charge->rho_save[is], is, ss1.str(), 3); } if (this->conv_elec) { @@ -477,16 +500,25 @@ namespace ModuleESolver std::stringstream ssp_ave; ssp << GlobalV::global_out_dir << "ElecStaticPot"; ssp_ave << GlobalV::global_out_dir << "ElecStaticPot_AVE"; - this->pelec->pot->write_elecstat_pot(ssp.str(), ssp_ave.str(), GlobalC::rhopw, pelec->charge); //output 'Hartree + local pseudopot' + this->pelec->pot->write_elecstat_pot(ssp.str(), ssp_ave.str(), GlobalC::rhopw, this->pelec->charge); //output 'Hartree + local pseudopot' } if (GlobalV::OUT_LEVEL != "m") { this->print_eigenvalue(GlobalV::ofs_running); } + if (this->device == psi::GpuDevice) { + syncmem_complex_d2h_op()( + this->psi[0].get_device(), + this->kspw_psi[0].get_device(), + this->psi[0].get_pointer() - this->psi[0].get_psi_bias(), + this->kspw_psi[0].get_pointer() - this->kspw_psi[0].get_psi_bias(), + this->psi[0].size()); + } } - void ESolver_KS_PW::print_eigenvalue(std::ofstream& ofs) + template + void ESolver_KS_PW::print_eigenvalue(std::ofstream& ofs) { bool wrong = false; for (int ik = 0; ik < GlobalC::kv.nks; ++ik) @@ -591,26 +623,29 @@ namespace ModuleESolver - void ESolver_KS_PW::cal_Energy(double& etot) + template + void ESolver_KS_PW::cal_Energy(FPTYPE& etot) { etot = GlobalC::en.etot; } - void ESolver_KS_PW::cal_Force(ModuleBase::matrix& force) + template + void ESolver_KS_PW::cal_Force(ModuleBase::matrix& force) { Forces ff; - ff.init(force, this->pelec->wg, pelec->charge, this->psi); + ff.init(force, this->pelec->wg, this->pelec->charge, this->psi); } - void ESolver_KS_PW::cal_Stress(ModuleBase::matrix& stress) + template + void ESolver_KS_PW::cal_Stress(ModuleBase::matrix& stress) { Stress_PW ss(this->pelec); ss.cal_stress(stress, this->psi); //external stress - double unit_transform = 0.0; + FPTYPE unit_transform = 0.0; unit_transform = ModuleBase::RYDBERG_SI / pow(ModuleBase::BOHR_RADIUS_SI,3) * 1.0e-8; - double external_stress[3] = {GlobalV::PRESS1,GlobalV::PRESS2,GlobalV::PRESS3}; + FPTYPE external_stress[3] = {GlobalV::PRESS1,GlobalV::PRESS2,GlobalV::PRESS3}; for(int i=0;i<3;i++) { stress(i,i) -= external_stress[i]/unit_transform; @@ -618,7 +653,8 @@ namespace ModuleESolver GlobalV::PRESSURE = (stress(0,0)+stress(1,1)+stress(2,2))/3; } - void ESolver_KS_PW::postprocess() + template + void ESolver_KS_PW::postprocess() { GlobalV::ofs_running << "\n\n --------------------------------------------" << std::endl; @@ -627,6 +663,7 @@ namespace ModuleESolver GlobalV::ofs_running << " --------------------------------------------\n\n" << std::endl; //print occupation in istate.info + GlobalC::en.print_occ(this->pelec); // compute density of states GlobalC::en.perform_dos_pw(this->pelec); @@ -684,13 +721,14 @@ namespace ModuleESolver } } - void ESolver_KS_PW::hamilt2estates(const double ethr) + template + void ESolver_KS_PW::hamilt2estates(const FPTYPE ethr) { if(this->phsol != nullptr) { - hsolver::DiagoIterAssist::need_subspace = false; - hsolver::DiagoIterAssist::PW_DIAG_THR = ethr; - this->phsol->solve(this->p_hamilt, this->psi[0], this->pelec, GlobalV::KS_SOLVER, true); + hsolver::DiagoIterAssist::need_subspace = false; + hsolver::DiagoIterAssist::PW_DIAG_THR = ethr; + this->phsol->solve(this->p_hamilt, this->kspw_psi[0], this->pelec, GlobalV::KS_SOLVER, true); } else { @@ -698,7 +736,8 @@ namespace ModuleESolver } } - void ESolver_KS_PW::nscf() + template + void ESolver_KS_PW::nscf() { ModuleBase::TITLE("ESolver_KS_PW","nscf"); ModuleBase::timer::tick("ESolver_KS_PW","nscf"); @@ -707,7 +746,7 @@ namespace ModuleESolver //======================================== // diagonalization of the KS hamiltonian // ======================================= - double diag_ethr = GlobalV::PW_DIAG_THR; + FPTYPE diag_ethr = GlobalV::PW_DIAG_THR; if(diag_ethr - 1e-2 > -1e-5) diag_ethr = std::max(1e-13, 0.1*std::min(1e-2,GlobalV::SCF_THR / GlobalV::nelec)); GlobalV::ofs_running << " PW_DIAG_THR = "<< diag_ethr << std::endl; @@ -764,4 +803,8 @@ namespace ModuleESolver return; } +template class ESolver_KS_PW; +#if ((defined __CUDA) || (defined __ROCM)) +template class ESolver_KS_PW; +#endif } diff --git a/source/module_esolver/esolver_ks_pw.h b/source/module_esolver/esolver_ks_pw.h index 6cf29b2035a..0ecd71ed7e4 100644 --- a/source/module_esolver/esolver_ks_pw.h +++ b/source/module_esolver/esolver_ks_pw.h @@ -10,22 +10,23 @@ namespace ModuleESolver { - class ESolver_KS_PW : public ESolver_KS + template + class ESolver_KS_PW : public ESolver_KS { public: ESolver_KS_PW(); ~ESolver_KS_PW(); void Init(Input& inp, UnitCell& cell) override; - void cal_Energy(double& etot) override; + void cal_Energy(FPTYPE& etot) override; void cal_Force(ModuleBase::matrix& force) override; void cal_Stress(ModuleBase::matrix& stress) override; - virtual void hamilt2density(const int istep, const int iter, const double ethr) override; - virtual void hamilt2estates(const double ethr) override; + virtual void hamilt2density(const int istep, const int iter, const FPTYPE ethr) override; + virtual void hamilt2estates(const FPTYPE ethr) override; virtual void nscf() override; void postprocess() override; //calculate conductivities with Kubo-Greenwood formula - void KG(const int nche_KG, const double fwhmin, const double wcut, - const double dw_in, const int times, ModuleBase::matrix& wg); + void KG(const int nche_KG, const FPTYPE fwhmin, const FPTYPE wcut, + const FPTYPE dw_in, const int times, ModuleBase::matrix& wg); protected: virtual void beforescf(const int istep) override; @@ -39,14 +40,18 @@ namespace ModuleESolver //Init Global class void Init_GlobalC(Input& inp, UnitCell& cell); //calculate conductivities from j-j correlation function - void calcondw(const int nt,const double dt,const double fwhmin,const double wcut,const double dw_in,double*ct11,double*ct12,double *ct22); + void calcondw(const int nt,const FPTYPE dt, const FPTYPE fwhmin, const FPTYPE wcut, const FPTYPE dw_in, FPTYPE *ct11, FPTYPE *ct12, FPTYPE *ct22); private: // It copies the function in Threshold_Elec class. // After all ESolver, HSolver are constructed, Class Electrons and Threshold_Elec should be deleted. void print_eigenvalue(std::ofstream& ofs); - + + Device * ctx = {}; + psi::AbacusDevice_t device = {}; + psi::Psi, Device>* kspw_psi = nullptr; + using syncmem_complex_d2h_op = psi::memory::synchronize_memory_op, psi::DEVICE_CPU, Device>; }; -} +} // namespace ModuleESolver #endif diff --git a/source/module_esolver/esolver_ks_pw_tool.cpp b/source/module_esolver/esolver_ks_pw_tool.cpp index 5a95a5f2054..b53a003c455 100644 --- a/source/module_esolver/esolver_ks_pw_tool.cpp +++ b/source/module_esolver/esolver_ks_pw_tool.cpp @@ -23,7 +23,9 @@ namespace ModuleESolver //------------------------------------------------------------------ #define TWOSQRT2LN2 2.354820045030949 // FWHM = 2sqrt(2ln2) * \sigma #define FACTOR 1.839939223835727e7 -void ESolver_KS_PW::KG(const int nche_KG, const double fwhmin, const double wcut, const double dw_in, const int times, ModuleBase::matrix& wg) +template +void ESolver_KS_PW::KG(const int nche_KG, const FPTYPE fwhmin, const FPTYPE wcut, + const FPTYPE dw_in, const int times, ModuleBase::matrix& wg) { //----------------------------------------------------------- // KS conductivity @@ -32,9 +34,9 @@ void ESolver_KS_PW::KG(const int nche_KG, const double fwhmin, const double wcut char transn = 'N'; char transc = 'C'; int nw = ceil(wcut / dw_in); - double dw = dw_in / ModuleBase::Ry_to_eV; // converge unit in eV to Ry - double sigma = fwhmin / TWOSQRT2LN2 / ModuleBase::Ry_to_eV; - double dt = ModuleBase::PI / (dw * nw) / times; // unit in a.u., 1 a.u. = 4.837771834548454e-17 s + FPTYPE dw = dw_in / ModuleBase::Ry_to_eV; // converge unit in eV to Ry + FPTYPE sigma = fwhmin / TWOSQRT2LN2 / ModuleBase::Ry_to_eV; + FPTYPE dt = ModuleBase::PI / (dw * nw) / times; // unit in a.u., 1 a.u. = 4.837771834548454e-17 s int nt = ceil(sqrt(20) / sigma / dt); cout << "nw: " << nw << " ; dw: " << dw * ModuleBase::Ry_to_eV << " eV" << endl; cout << "nt: " << nt << " ; dt: " << dt << " a.u.(ry^-1)" << endl; @@ -43,13 +45,13 @@ void ESolver_KS_PW::KG(const int nche_KG, const double fwhmin, const double wcut const int nk = GlobalC::kv.nks; const int ndim = 3; const int npwx = GlobalC::wf.npwx; - const double tpiba = GlobalC::ucell.tpiba; + const FPTYPE tpiba = GlobalC::ucell.tpiba; const int nbands = GlobalV::NBANDS; - const double ef = GlobalC::en.ef; + const FPTYPE ef = GlobalC::en.ef; - double *ct11 = new double[nt]; - double *ct12 = new double[nt]; - double *ct22 = new double[nt]; + FPTYPE *ct11 = new FPTYPE[nt]; + FPTYPE *ct12 = new FPTYPE[nt]; + FPTYPE *ct22 = new FPTYPE[nt]; ModuleBase::GlobalFunc::ZEROS(ct11, nt); ModuleBase::GlobalFunc::ZEROS(ct12, nt); ModuleBase::GlobalFunc::ZEROS(ct22, nt); @@ -59,14 +61,14 @@ void ESolver_KS_PW::KG(const int nche_KG, const double fwhmin, const double wcut { velop.init(ik); const int npw = GlobalC::kv.ngk[ik]; - complex *levc = &(this->psi[0](ik, 0, 0)); - complex *prevc = new complex[3 * npwx * nbands]; + complex *levc = &(this->psi[0](ik, 0, 0)); + complex *prevc = new complex[3 * npwx * nbands]; // px|right> velop.act(this->psi, nbands*GlobalV::NPOL, levc, prevc); for (int id = 0; id < ndim; ++id) { this->p_hamilt->updateHk(ik); - complex *pij = new complex[nbands * nbands]; + complex *pij = new complex[nbands * nbands]; zgemm_(&transc, &transn, &nbands, @@ -98,19 +100,19 @@ void ESolver_KS_PW::KG(const int nche_KG, const double fwhmin, const double wcut for (int it = itstart; it < itstart + ntper; ++it) // for(int it = 0 ; it < nt; ++it) { - double tmct11 = 0; - double tmct12 = 0; - double tmct22 = 0; - double *enb = &(this->pelec->ekb(ik, 0)); + FPTYPE tmct11 = 0; + FPTYPE tmct12 = 0; + FPTYPE tmct22 = 0; + FPTYPE *enb = &(this->pelec->ekb(ik, 0)); for (int ib = 0; ib < nbands; ++ib) { - double ei = enb[ib]; - double fi = wg(ik, ib); + FPTYPE ei = enb[ib]; + FPTYPE fi = wg(ik, ib); for (int jb = ib + 1; jb < nbands; ++jb) { - double ej = enb[jb]; - double fj = wg(ik, jb); - double tmct = sin((ej - ei) * (it)*dt) * (fi - fj) * norm(pij[ib * nbands + jb]); + FPTYPE ej = enb[jb]; + FPTYPE fj = wg(ik, jb); + FPTYPE tmct = sin((ej - ei) * (it)*dt) * (fi - fj) * norm(pij[ib * nbands + jb]); tmct11 += tmct; tmct12 += -tmct * ((ei + ej) / 2 - ef); tmct22 += tmct * pow((ei + ej) / 2 - ef, 2); @@ -142,33 +144,34 @@ void ESolver_KS_PW::KG(const int nche_KG, const double fwhmin, const double wcut delete[] ct22; } -void ESolver_KS_PW::calcondw(const int nt, - const double dt, - const double fwhmin, - const double wcut, - const double dw_in, - double *ct11, - double *ct12, - double *ct22) +template +void ESolver_KS_PW::calcondw(const int nt, + const FPTYPE dt, + const FPTYPE fwhmin, + const FPTYPE wcut, + const FPTYPE dw_in, + FPTYPE *ct11, + FPTYPE *ct12, + FPTYPE *ct22) { - double factor = FACTOR; + FPTYPE factor = FACTOR; const int ndim = 3; int nw = ceil(wcut / dw_in); - double dw = dw_in / ModuleBase::Ry_to_eV; // converge unit in eV to Ry - double sigma = fwhmin / TWOSQRT2LN2 / ModuleBase::Ry_to_eV; + FPTYPE dw = dw_in / ModuleBase::Ry_to_eV; // converge unit in eV to Ry + FPTYPE sigma = fwhmin / TWOSQRT2LN2 / ModuleBase::Ry_to_eV; ofstream ofscond("je-je.txt"); ofscond << setw(8) << "#t(a.u.)" << setw(15) << "c11(t)" << setw(15) << "c12(t)" << setw(15) << "c22(t)" << setw(15) << "decay" << endl; for (int it = 0; it < nt; ++it) { ofscond << setw(8) << (it)*dt << setw(15) << -2 * ct11[it] << setw(15) << -2 * ct12[it] << setw(15) - << -2 * ct22[it] << setw(15) << exp(-double(1) / 2 * sigma * sigma * pow((it)*dt, 2)) << endl; + << -2 * ct22[it] << setw(15) << exp(-FPTYPE(1) / 2 * sigma * sigma * pow((it)*dt, 2)) << endl; } ofscond.close(); - double *cw11 = new double[nw]; - double *cw12 = new double[nw]; - double *cw22 = new double[nw]; - double *kappa = new double[nw]; + FPTYPE *cw11 = new FPTYPE[nw]; + FPTYPE *cw12 = new FPTYPE[nw]; + FPTYPE *cw22 = new FPTYPE[nw]; + FPTYPE *kappa = new FPTYPE[nw]; ModuleBase::GlobalFunc::ZEROS(cw11, nw); ModuleBase::GlobalFunc::ZEROS(cw12, nw); ModuleBase::GlobalFunc::ZEROS(cw22, nw); @@ -177,11 +180,11 @@ void ESolver_KS_PW::calcondw(const int nt, for (int it = 0; it < nt; ++it) { cw11[iw] += -2 * ct11[it] * sin(-(iw + 0.5) * dw * it * dt) - * exp(-double(1) / 2 * sigma * sigma * pow((it)*dt, 2)) / (iw + 0.5) / dw * dt; + * exp(-FPTYPE(1) / 2 * sigma * sigma * pow((it)*dt, 2)) / (iw + 0.5) / dw * dt; cw12[iw] += -2 * ct12[it] * sin(-(iw + 0.5) * dw * it * dt) - * exp(-double(1) / 2 * sigma * sigma * pow((it)*dt, 2)) / (iw + 0.5) / dw * dt; + * exp(-FPTYPE(1) / 2 * sigma * sigma * pow((it)*dt, 2)) / (iw + 0.5) / dw * dt; cw22[iw] += -2 * ct22[it] * sin(-(iw + 0.5) * dw * it * dt) - * exp(-double(1) / 2 * sigma * sigma * pow((it)*dt, 2)) / (iw + 0.5) / dw * dt; + * exp(-FPTYPE(1) / 2 * sigma * sigma * pow((it)*dt, 2)) / (iw + 0.5) / dw * dt; } } ofscond.open("Onsager.txt"); @@ -189,10 +192,10 @@ void ESolver_KS_PW::calcondw(const int nt, << "L12/e(Am^-1)" << setw(20) << "L22/e^2(Wm^-1)" << endl; for (int iw = 0; iw < nw; ++iw) { - cw11[iw] *= double(2) / ndim / GlobalC::ucell.omega * factor; // unit in Sm^-1 + cw11[iw] *= FPTYPE(2) / ndim / GlobalC::ucell.omega * factor; // unit in Sm^-1 cw12[iw] - *= double(2) / ndim / GlobalC::ucell.omega * factor * 2.17987092759e-18 / 1.6021766208e-19; // unit in Am^-1 - cw22[iw] *= double(2) / ndim / GlobalC::ucell.omega * factor + *= FPTYPE(2) / ndim / GlobalC::ucell.omega * factor * 2.17987092759e-18 / 1.6021766208e-19; // unit in Am^-1 + cw22[iw] *= FPTYPE(2) / ndim / GlobalC::ucell.omega * factor * pow(2.17987092759e-18 / 1.6021766208e-19, 2); // unit in Wm^-1 kappa[iw] = (cw22[iw] - pow(cw12[iw], 2) / cw11[iw]) / Occupy::gaussian_parameter / ModuleBase::Ry_to_eV / 11604.518026; @@ -210,4 +213,9 @@ void ESolver_KS_PW::calcondw(const int nt, delete[] cw22; delete[] kappa; } + +template class ESolver_KS_PW; +#if ((defined __CUDA) || (defined __ROCM)) +template class ESolver_KS_PW; +#endif } // namespace ModuleESolver \ No newline at end of file diff --git a/source/module_esolver/esolver_sdft_pw.h b/source/module_esolver/esolver_sdft_pw.h index 9805f955c4f..5a671890b9f 100644 --- a/source/module_esolver/esolver_sdft_pw.h +++ b/source/module_esolver/esolver_sdft_pw.h @@ -9,7 +9,7 @@ namespace ModuleESolver { -class ESolver_SDFT_PW: public ESolver_KS_PW +class ESolver_SDFT_PW: public ESolver_KS_PW { public: ESolver_SDFT_PW(); diff --git a/source/module_hamilt/hamilt_pw.cpp b/source/module_hamilt/hamilt_pw.cpp index 7ab1f4204a3..60b364519ef 100644 --- a/source/module_hamilt/hamilt_pw.cpp +++ b/source/module_hamilt/hamilt_pw.cpp @@ -72,7 +72,9 @@ HamiltPW::HamiltPW(elecstate::Potential* pot_in) pot_in->pot_register(pot_register_in); Operator, Device>* veff = new Veff>( isk, - &(pot_in->get_effective_v()), + pot_in->get_effective_v_data(), + pot_in->get_effective_v().nr, + pot_in->get_effective_v().nc, GlobalC::wfcpw ); if(this->ops == nullptr) @@ -203,10 +205,10 @@ HamiltPW::HamiltPW(const HamiltPW *hamilt) } template class HamiltPW; +template HamiltPW::HamiltPW(const HamiltPW *hamilt); #if ((defined __CUDA) || (defined __ROCM)) template class HamiltPW; -template HamiltPW::HamiltPW(const HamiltPW *hamilt); -template HamiltPW::HamiltPW(const HamiltPW *hamilt); +template HamiltPW::HamiltPW(const HamiltPW *hamilt); #endif } // namespace hamilt \ No newline at end of file diff --git a/source/module_hamilt/ks_pw/ekinetic_pw.cpp b/source/module_hamilt/ks_pw/ekinetic_pw.cpp index 743aff29bc9..19d7cba60b1 100644 --- a/source/module_hamilt/ks_pw/ekinetic_pw.cpp +++ b/source/module_hamilt/ks_pw/ekinetic_pw.cpp @@ -2,6 +2,7 @@ #include "module_base/timer.h" #include "module_base/tool_quit.h" +#include "module_psi/include/device.h" using hamilt::Ekinetic; using hamilt::OperatorPW; @@ -16,13 +17,17 @@ Ekinetic>::Ekinetic( this->classname = "Ekinetic"; this->cal_type = pw_ekinetic; this->tpiba2 = tpiba2_in; + this->gk2 = gk2_in; + this->gk2_in = gk2_in; this->gk2_row = gk2_row; this->gk2_col = gk2_col; + this->device = psi::device::get_device_type(this->ctx); #if ((defined __CUDA) || (defined __ROCM)) - resize_memory_op()(this->ctx, this->gk2, this->gk2_row * this->gk2_col); - synchronize_memory_op()(this->ctx, this->cpu_ctx, this->gk2, gk2_in, this->gk2_row * this->gk2_col); -#else - this->gk2 = gk2_in; + if (this->device == psi::GpuDevice) { + resmem_var_op()(this->ctx, this->_gk2, this->gk2_row * this->gk2_col); + syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, this->_gk2, gk2_in, this->gk2_row * this->gk2_col); + this->gk2 = this->_gk2; + } #endif if( this->tpiba2 < 1e-10 || this->gk2 == nullptr) { ModuleBase::WARNING_QUIT("EkineticPW", "Constuctor of Operator::EkineticPW is failed, please check your code!"); @@ -32,7 +37,9 @@ Ekinetic>::Ekinetic( template Ekinetic>::~Ekinetic() { #if ((defined __CUDA) || (defined __ROCM)) - delete_memory_op()(this->ctx, this->gk2); + if (this->device == psi::GpuDevice) { + delmem_var_op()(this->ctx, this->_gk2); + } #endif // __CUDA || __ROCM } @@ -70,14 +77,20 @@ hamilt::Ekinetic>::Ekinetic(const Ekineticcal_type = pw_ekinetic; this->ik = ekinetic->get_ik(); this->tpiba2 = ekinetic->get_tpiba2(); + this->gk2 = ekinetic->get_gk2(); this->gk2_row = ekinetic->get_gk2_row(); this->gk2_col = ekinetic->get_gk2_col(); - resize_memory_op()(this->ctx, this->gk2, this->gk2_row * this->gk2_col); - psi::memory::synchronize_memory_op()( - this->ctx, ekinetic->get_ctx(), - this->gk2, ekinetic->get_gk2(), - this->gk2_row * this->gk2_col); - + this->device = psi::device::get_device_type(this->ctx); +#if ((defined __CUDA) || (defined __ROCM)) + if (this->device == psi::GpuDevice) { + resmem_var_op()(this->ctx, this->_gk2, this->gk2_row * this->gk2_col); + psi::memory::synchronize_memory_op()( + this->ctx, ekinetic->get_ctx(), + this->_gk2, ekinetic->get_gk2(), + this->gk2_row * this->gk2_col); + this->gk2 = this->_gk2; + } +#endif if( this->tpiba2 < 1e-10 || this->gk2 == nullptr) { ModuleBase::WARNING_QUIT("EkineticPW", "Copy Constuctor of Operator::EkineticPW is failed, please check your code!"); } @@ -85,9 +98,11 @@ hamilt::Ekinetic>::Ekinetic(const Ekinetic>; +template Ekinetic>::Ekinetic(const Ekinetic> *ekinetic); #if ((defined __CUDA) || (defined __ROCM)) template class Ekinetic>; template Ekinetic>::Ekinetic(const Ekinetic> *ekinetic); template Ekinetic>::Ekinetic(const Ekinetic> *ekinetic); +template Ekinetic>::Ekinetic(const Ekinetic> *ekinetic); #endif } // namespace hamilt \ No newline at end of file diff --git a/source/module_hamilt/ks_pw/ekinetic_pw.h b/source/module_hamilt/ks_pw/ekinetic_pw.h index e83332a0dfa..cc256f124a5 100644 --- a/source/module_hamilt/ks_pw/ekinetic_pw.h +++ b/source/module_hamilt/ks_pw/ekinetic_pw.h @@ -44,7 +44,7 @@ class Ekinetic> : public OperatorPW int get_gk2_row() const {return this->gk2_row;} int get_gk2_col() const {return this->gk2_col;} FPTYPE get_tpiba2() const {return this->tpiba2;} - const FPTYPE* get_gk2() const {return this->gk2;} + const FPTYPE* get_gk2() const {return this->gk2_in;} Device* get_ctx() const {return this->ctx;} private: @@ -54,21 +54,20 @@ class Ekinetic> : public OperatorPW mutable int npol = 0; FPTYPE tpiba2 = 0.0; -#if ((defined __CUDA) || (defined __ROCM)) - FPTYPE* gk2 = nullptr; -#else + FPTYPE* _gk2 = nullptr; const FPTYPE* gk2 = nullptr; -#endif + const FPTYPE* gk2_in = nullptr; int gk2_row = 0; int gk2_col = 0; Device* ctx = {}; psi::DEVICE_CPU* cpu_ctx = {}; + psi::AbacusDevice_t device = {}; using ekinetic_op = ekinetic_pw_op; - using resize_memory_op = psi::memory::resize_memory_op; - using delete_memory_op = psi::memory::delete_memory_op; - using synchronize_memory_op = psi::memory::synchronize_memory_op; + using resmem_var_op = psi::memory::resize_memory_op; + using delmem_var_op = psi::memory::delete_memory_op; + using syncmem_var_h2d_op = psi::memory::synchronize_memory_op; }; } // namespace hamilt diff --git a/source/module_hamilt/ks_pw/meta_pw.cpp b/source/module_hamilt/ks_pw/meta_pw.cpp index 43914f95708..af61a43d04c 100644 --- a/source/module_hamilt/ks_pw/meta_pw.cpp +++ b/source/module_hamilt/ks_pw/meta_pw.cpp @@ -104,9 +104,11 @@ hamilt::Meta>::Meta(const Meta>; +template Meta>::Meta(const Meta> *meta); #if ((defined __CUDA) || (defined __ROCM)) template class Meta>; template Meta>::Meta(const Meta> *meta); template Meta>::Meta(const Meta> *meta); +template Meta>::Meta(const Meta> *meta); #endif } // namespace hamilt \ No newline at end of file diff --git a/source/module_hamilt/ks_pw/nonlocal_pw.cpp b/source/module_hamilt/ks_pw/nonlocal_pw.cpp index 994b0a67047..7e017626d3a 100644 --- a/source/module_hamilt/ks_pw/nonlocal_pw.cpp +++ b/source/module_hamilt/ks_pw/nonlocal_pw.cpp @@ -24,7 +24,7 @@ Nonlocal>::Nonlocal( if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { this->deeq = this->ppcell->d_deeq; this->deeq_nc = this->ppcell->d_deeq_nc; - resize_memory_op()(this->ctx, this->vkb, this->ppcell->vkb.size); + resmem_complex_op()(this->ctx, this->vkb, this->ppcell->vkb.size); } else { this->deeq = this->ppcell->deeq.ptr; @@ -39,16 +39,17 @@ Nonlocal>::Nonlocal( template Nonlocal>::~Nonlocal() { - delete_memory_op()(this->ctx, this->ps); - delete_memory_op()(this->ctx, this->becp); + delmem_complex_op()(this->ctx, this->ps); + delmem_complex_op()(this->ctx, this->becp); if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { - delete_memory_op()(this->ctx, this->vkb); + delmem_complex_op()(this->ctx, this->vkb); } } template void Nonlocal>::init(const int ik_in) { + ModuleBase::timer::tick("Nonlocal", "getvnl"); this->ik = ik_in; // Calculate nonlocal pseudopotential vkb if(this->ppcell->nkb > 0) //xiaohui add 2013-09-02. Attention... @@ -60,6 +61,11 @@ void Nonlocal>::init(const int ik_in) { this->next_op->init(ik_in); } + + if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { + syncmem_complex_h2d_op()(this->ctx, this->cpu_ctx, this->vkb, this->ppcell->vkb.c, this->ppcell->vkb.size); + } + ModuleBase::timer::tick("Nonlocal", "getvnl"); } //-------------------------------------------------------------------------- @@ -75,8 +81,8 @@ void Nonlocal>::add_nonlocal_pp(std::complex // std::complex *ps = new std::complex[nkb * m]; // ModuleBase::GlobalFunc::ZEROS(ps, m * nkb); - resize_memory_op()(this->ctx, this->ps, nkb * m); - set_memory_op()(this->ctx, this->ps, 0, nkb * m); + resmem_complex_op()(this->ctx, this->ps, nkb * m); + setmem_complex_op()(this->ctx, this->ps, 0, nkb * m); int sum = 0; int iat = 0; @@ -170,29 +176,19 @@ void Nonlocal>::add_nonlocal_pp(std::complex int inc = 1; // denghui replace 2022-10-20 // <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - gemv_op()(this->ctx, - transa, - this->npw, - this->ppcell->nkb, - &ModuleBase::ONE, - this->vkb, - this->ppcell->vkb.nc, - this->ps, - inc, - &ModuleBase::ONE, - hpsi_in, - inc); - // zgemv_(&transa, - // &this->npw, - // &(this->ppcell->nkb), - // &ModuleBase::ONE, - // this->ppcell->vkb.c, - // &this->ppcell->vkb.nc, - // ps, - // &inc, - // &ModuleBase::ONE, - // hpsi_in, - // &inc); + gemv_op()( + this->ctx, + transa, + this->npw, + this->ppcell->nkb, + &ModuleBase::ONE, + this->vkb, + this->ppcell->vkb.nc, + this->ps, + inc, + &ModuleBase::ONE, + hpsi_in, + inc); } else { @@ -215,19 +211,6 @@ void Nonlocal>::add_nonlocal_pp(std::complex hpsi_in, this->max_npw ); - // zgemm_(&transa, - // &transb, - // &this->npw, - // &npm, - // &(this->ppcell->nkb), - // &ModuleBase::ONE, - // this->ppcell->vkb.c, - // &this->ppcell->vkb.nc, - // ps, - // &npm, - // &ModuleBase::ONE, - // hpsi_in, - // &this->max_npw); } ModuleBase::timer::tick("Nonlocal", "add_nonlocal_pp"); } @@ -245,15 +228,12 @@ void Nonlocal>::act this->max_npw = psi_in->get_nbasis() / psi_in->npol; this->npol = psi_in->npol; - if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { - syncmem_complex_h2d_op()(this->ctx, this->cpu_ctx, this->vkb, this->ppcell->vkb.c, this->ppcell->vkb.size); - } if (this->ppcell->nkb > 0) { //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // qianrui optimize 2021-3-31 int nkb = this->ppcell->nkb; - resize_memory_op()(this->ctx, this->becp, n_npwx * nkb); + resmem_complex_op()(this->ctx, this->becp, n_npwx * nkb); // ModuleBase::ComplexMatrix becp(n_npwx, nkb, false); char transa = 'C'; char transb = 'N'; @@ -262,30 +242,19 @@ void Nonlocal>::act int inc = 1; // denghui replace 2022-10-20 // <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - gemv_op()(this->ctx, - transa, - this->npw, - nkb, - &ModuleBase::ONE, - this->vkb, - this->ppcell->vkb.nc, - tmpsi_in, - inc, - &ModuleBase::ZERO, - this->becp, - inc); - // zgemv_(&transa, - // &this->npw, - // &nkb, - // &ModuleBase::ONE, - // this->ppcell->vkb.c, - // &this->ppcell->vkb.nc, - // tmpsi_in, - // &inc, - // &ModuleBase::ZERO, - // this->becp, - // &inc); - + gemv_op()( + this->ctx, + transa, + this->npw, + nkb, + &ModuleBase::ONE, + this->vkb, + this->ppcell->vkb.nc, + tmpsi_in, + inc, + &ModuleBase::ZERO, + this->becp, + inc); } else { @@ -308,21 +277,6 @@ void Nonlocal>::act this->becp, nkb ); - //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - // zgemm_(&transa, - // &transb, - // &nkb, - // &npm, - // &this->npw, - // &ModuleBase::ONE, - // this->ppcell->vkb.c, - // &this->ppcell->vkb.nc, - // tmpsi_in, - // &this->max_npw, - // &ModuleBase::ZERO, - // becp, - // &nkb); - //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< } Parallel_Reduce::reduce_complex_double_pool(becp, nkb * n_npwx); @@ -344,7 +298,7 @@ hamilt::Nonlocal>::Nonlocal(const Nonlocalucell = nonlocal->get_ucell(); if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { this->deeq = this->ppcell->d_deeq; - resize_memory_op()(this->ctx, this->vkb, this->ppcell->vkb.size); + resmem_complex_op()(this->ctx, this->vkb, this->ppcell->vkb.size); } else { this->deeq = this->ppcell->deeq.ptr; @@ -358,9 +312,11 @@ hamilt::Nonlocal>::Nonlocal(const Nonlocal>; +template Nonlocal>::Nonlocal(const Nonlocal> *nonlocal); #if ((defined __CUDA) || (defined __ROCM)) template class Nonlocal>; template Nonlocal>::Nonlocal(const Nonlocal> *nonlocal); template Nonlocal>::Nonlocal(const Nonlocal> *nonlocal); +template Nonlocal>::Nonlocal(const Nonlocal> *nonlocal); #endif } // namespace hamilt \ No newline at end of file diff --git a/source/module_hamilt/ks_pw/nonlocal_pw.h b/source/module_hamilt/ks_pw/nonlocal_pw.h index 820cabd6d40..fef9f6bbca7 100644 --- a/source/module_hamilt/ks_pw/nonlocal_pw.h +++ b/source/module_hamilt/ks_pw/nonlocal_pw.h @@ -24,7 +24,7 @@ template class Nonlocal : public T template class Nonlocal> : public OperatorPW { - public: + public: Nonlocal(const int* isk_in,const pseudopot_cell_vnl* ppcell_in,const UnitCell* ucell_in); template @@ -45,7 +45,7 @@ class Nonlocal> : public OperatorPW const pseudopot_cell_vnl *get_ppcell() const {return this->ppcell;} const UnitCell *get_ucell() const {return this->ucell;} - private: + private: void add_nonlocal_pp(std::complex *hpsi_in, const std::complex *becp, const int m) const; mutable int max_npw = 0; @@ -71,9 +71,9 @@ class Nonlocal> : public OperatorPW using gemv_op = hsolver::gemv_op; using gemm_op = hsolver::gemm_op; using nonlocal_op = nonlocal_pw_op; - using set_memory_op = psi::memory::set_memory_op, Device>; - using resize_memory_op = psi::memory::resize_memory_op, Device>; - using delete_memory_op = psi::memory::delete_memory_op, Device>; + using setmem_complex_op = psi::memory::set_memory_op, Device>; + using resmem_complex_op = psi::memory::resize_memory_op, Device>; + using delmem_complex_op = psi::memory::delete_memory_op, Device>; using syncmem_complex_h2d_op = psi::memory::synchronize_memory_op, Device, psi::DEVICE_CPU>; }; diff --git a/source/module_hamilt/ks_pw/veff_pw.cpp b/source/module_hamilt/ks_pw/veff_pw.cpp index 8a2ad03166b..48147bb3264 100644 --- a/source/module_hamilt/ks_pw/veff_pw.cpp +++ b/source/module_hamilt/ks_pw/veff_pw.cpp @@ -10,41 +10,31 @@ using hamilt::OperatorPW; template Veff>::Veff( const int* isk_in, - const ModuleBase::matrix* veff_in, + const FPTYPE* veff_in, + const int veff_row, + const int veff_col, ModulePW::PW_Basis_K* wfcpw_in) { this->classname = "Veff"; this->cal_type = pw_veff; this->isk = isk_in; - // this->veff = veff_in; - // TODO: add an GPU veff array - this->veff = veff_in[0].c; + this->veff = veff_in; //note: "veff = nullptr" means that this core does not treat potential but still treats wf. - this->veff_col = veff_in[0].nc; - this->veff_row = veff_in[0].nr; + this->veff_row = veff_row; + this->veff_col = veff_col; this->wfcpw = wfcpw_in; - this->device = psi::device::get_device_type(this->ctx); - resize_memory_complex_op()(this->ctx, this->porter, this->wfcpw->nmaxgr); - resize_memory_complex_op()(this->ctx, this->porter1, this->wfcpw->nmaxgr); - if (this->device == psi::GpuDevice) { - resize_memory_double_op()(this->ctx, this->d_veff, this->veff_col * this->veff_row); - } + resmem_complex_op()(this->ctx, this->porter, this->wfcpw->nmaxgr); + resmem_complex_op()(this->ctx, this->porter1, this->wfcpw->nmaxgr); if (this->isk == nullptr || this->wfcpw == nullptr) { ModuleBase::WARNING_QUIT("VeffPW", "Constuctor of Operator::VeffPW is failed, please check your code!"); } - if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { - resize_memory_double_op()(this->ctx, this->d_veff, this->veff_col * this->veff_row); - } } template Veff>::~Veff() { - delete_memory_complex_op()(this->ctx, this->porter); - delete_memory_complex_op()(this->ctx, this->porter1); - if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { - delete_memory_double_op()(this->ctx, this->d_veff); - } + delmem_complex_op()(this->ctx, this->porter); + delmem_complex_op()(this->ctx, this->porter1); } template @@ -53,7 +43,7 @@ void Veff>::act( const int n_npwx, const std::complex* tmpsi_in, std::complex* tmhpsi -)const +)const { ModuleBase::timer::tick("Operator", "VeffPW"); @@ -73,18 +63,12 @@ void Veff>::act( // but the 3DFFT can not be skipped, it will cause hanging if(this->veff_col != 0) { + veff_op()(this->ctx, this->veff_col, this->porter, this->veff + current_spin * this->veff_col); // const FPTYPE* current_veff = &(this->veff[0](current_spin, 0)); // for (int ir = 0; ir < this->veff->nc; ++ir) // { // porter[ir] *= current_veff[ir]; // } - if (this->device == psi::GpuDevice) { - syncmem_double_h2d_op()(this->ctx, this->cpu_ctx, this->d_veff, this->veff, this->veff_col * this->veff_row); - veff_op()(this->ctx, this->veff_col, this->porter, this->d_veff + current_spin * this->veff_col); - } - else { - veff_op()(this->ctx, this->veff_col, this->porter, this->veff + current_spin * this->veff_col); - } } // wfcpw->real2recip(porter, tmhpsi, this->ik, true); wfcpw->real_to_recip(this->ctx, this->porter, tmhpsi, this->ik, true); @@ -99,13 +83,8 @@ void Veff>::act( { /// denghui added at 20221109 const FPTYPE* current_veff[4]; - if (this->device == psi::GpuDevice) { - syncmem_double_h2d_op()(this->ctx, this->cpu_ctx, this->d_veff, this->veff, this->veff_col * this->veff_row); - } for(int is = 0; is < 4; is++) { - current_veff[is] = this->device == psi::GpuDevice ? - this->d_veff + is * this->veff_col : // for GPU device - this->veff + is * this->veff_col ; // for CPU device + current_veff[is] = this->veff + is * this->veff_col ; // for CPU device } veff_op()(this->ctx, this->veff_col, this->porter, this->porter1, current_veff); // std::complex sup, sdown; @@ -139,31 +118,23 @@ hamilt::Veff>::Veff(const Veffcal_type = pw_veff; this->ik = veff->get_ik(); this->isk = veff->get_isk(); - // this->veff = veff_in; - // TODO: add an GPU veff array - this->veff = veff->get_veff(); this->veff_col = veff->get_veff_col(); this->veff_row = veff->get_veff_row(); this->wfcpw = veff->get_wfcpw(); this->npol = veff->get_npol(); - this->device = psi::device::get_device_type(this->ctx); - resize_memory_complex_op()(this->ctx, this->porter, this->wfcpw->nmaxgr); - if (this->npol != 1) { - resize_memory_complex_op()(this->ctx, this->porter1, this->wfcpw->nmaxgr); - } + resmem_complex_op()(this->ctx, this->porter, this->wfcpw->nmaxgr); + resmem_complex_op()(this->ctx, this->porter1, this->wfcpw->nmaxgr); + this->veff = veff->get_veff(); if (this->isk == nullptr || this->veff == nullptr || this->wfcpw == nullptr) { ModuleBase::WARNING_QUIT("VeffPW", "Constuctor of Operator::VeffPW is failed, please check your code!"); } - if (psi::device::get_device_type(this->ctx) == psi::GpuDevice) { - resize_memory_double_op()(this->ctx, this->d_veff, this->veff_col * this->veff_row); - } } namespace hamilt{ template class Veff>; +template Veff>::Veff(const Veff> *veff); #if ((defined __CUDA) || (defined __ROCM)) template class Veff>; -template Veff>::Veff(const Veff> *veff); -template Veff>::Veff(const Veff> *veff); +template Veff>::Veff(const Veff> *veff); #endif } // namespace hamilt \ No newline at end of file diff --git a/source/module_hamilt/ks_pw/veff_pw.h b/source/module_hamilt/ks_pw/veff_pw.h index 83fa12b1ce5..1717280311b 100644 --- a/source/module_hamilt/ks_pw/veff_pw.h +++ b/source/module_hamilt/ks_pw/veff_pw.h @@ -22,7 +22,7 @@ template class Veff> : public OperatorPW { public: - Veff(const int* isk_in,const ModuleBase::matrix* veff_in,ModulePW::PW_Basis_K* wfcpw_in); + Veff(const int* isk_in, const FPTYPE* veff_in, const int veff_row, const int veff_col, ModulePW::PW_Basis_K* wfcpw_in); template explicit Veff(const Veff>* veff); @@ -37,7 +37,7 @@ class Veff> : public OperatorPW )const override; // denghui added for copy constructor at 20221105 - FPTYPE *get_veff() const {return this->veff;} + const FPTYPE *get_veff() const {return this->veff;} int get_veff_col() const {return this->veff_col;} int get_veff_row() const {return this->veff_row;} int get_npol() const {return this->npol;} @@ -60,19 +60,15 @@ class Veff> : public OperatorPW int veff_col = 0; int veff_row = 0; - FPTYPE *veff = nullptr; - FPTYPE *d_veff = nullptr; + const FPTYPE *veff = nullptr, *h_veff = nullptr, *d_veff = nullptr; std::complex *porter = nullptr; std::complex *porter1 = nullptr; psi::AbacusDevice_t device = {}; using veff_op = veff_pw_op; - using resize_memory_double_op = psi::memory::resize_memory_op; - using delete_memory_double_op = psi::memory::delete_memory_op; - using resize_memory_complex_op = psi::memory::resize_memory_op, Device>; - using delete_memory_complex_op = psi::memory::delete_memory_op, Device>; - using syncmem_double_h2d_op = psi::memory::synchronize_memory_op; + using resmem_complex_op = psi::memory::resize_memory_op, Device>; + using delmem_complex_op = psi::memory::delete_memory_op, Device>; }; } // namespace hamilt diff --git a/source/module_hamilt/operator.cpp b/source/module_hamilt/operator.cpp index b87cf317624..0ece6061781 100644 --- a/source/module_hamilt/operator.cpp +++ b/source/module_hamilt/operator.cpp @@ -27,14 +27,6 @@ typename Operator::hpsi_info Operator::hPsi(hpsi return hpsi_info(nullptr, 0, nullptr); } -#if ((defined __CUDA) || (defined __ROCM)) -template -typename Operator::hpsi_info_gpu Operator::hPsi_gpu(hpsi_info_gpu&) const { - ModuleBase::WARNING_QUIT("Operator::hPsi_gpu", "GPU's implementation is not supported!"); - return hpsi_info_gpu(nullptr, 0, nullptr); -} -#endif // ((defined __CUDA) || (defined __ROCM)) - template void Operator::init(const int ik_in) { diff --git a/source/module_hamilt/operator.h b/source/module_hamilt/operator.h index e41bca50a4b..2ffe57dd18c 100644 --- a/source/module_hamilt/operator.h +++ b/source/module_hamilt/operator.h @@ -47,11 +47,6 @@ class Operator virtual int get_ik() const {return this->ik;} -#if ((defined __CUDA) || (defined __ROCM)) - typedef std::tuple*, const psi::Range, FPTYPE*> hpsi_info_gpu; - virtual hpsi_info_gpu hPsi_gpu(hpsi_info_gpu& input) const; - #endif // ((defined __CUDA) || (defined __ROCM)) - Operator* next_op = nullptr; protected: diff --git a/source/module_hsolver/diago_cg.cpp b/source/module_hsolver/diago_cg.cpp index cc8a96f0c02..a3afc5182aa 100644 --- a/source/module_hsolver/diago_cg.cpp +++ b/source/module_hsolver/diago_cg.cpp @@ -19,25 +19,19 @@ DiagoCG::DiagoCG(const FPTYPE* precondition_in) test_cg = 0; reorder = false; this->device = psi::device::get_device_type(this->ctx); -#if defined(__CUDA) || defined(__ROCM) - hsolver::createBLAShandle(); -#endif } template DiagoCG::~DiagoCG() { // delete this->cg; // delete this->phi_m; - delete_memory_complex_op()(this->ctx, this->sphi); - delete_memory_complex_op()(this->ctx, this->hphi); - delete_memory_complex_op()(this->ctx, this->scg); - delete_memory_complex_op()(this->ctx, this->pphi); - delete_memory_complex_op()(this->ctx, this->gradient); - delete_memory_complex_op()(this->ctx, this->g0); - delete_memory_complex_op()(this->ctx, this->lagrange); -#if defined(__CUDA) || defined(__ROCM) - hsolver::destoryBLAShandle(); -#endif + delmem_complex_op()(this->ctx, this->sphi); + delmem_complex_op()(this->ctx, this->hphi); + delmem_complex_op()(this->ctx, this->scg); + delmem_complex_op()(this->ctx, this->pphi); + delmem_complex_op()(this->ctx, this->gradient); + delmem_complex_op()(this->ctx, this->g0); + delmem_complex_op()(this->ctx, this->lagrange); } template @@ -55,7 +49,7 @@ void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, this->n_band = phi.get_nbands(); this->eigenvalue = eigenvalue_in; - set_memory_var_cpu_op()(this->cpu_ctx, this->eigenvalue, 0, this->n_band); + setmem_var_h_op()(this->cpu_ctx, this->eigenvalue, 0, this->n_band); // ModuleBase::GlobalFunc::ZEROS(this->eigenvalue, this->n_band); /// record for how many loops in cg convergence @@ -70,32 +64,32 @@ void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, //------------------------------------------------------------------- this->phi_m = new psi::Psi, Device>(phi, 1, 1); // this->hphi.resize(this->dmx, ModuleBase::ZERO); - resize_memory_complex_op()(this->ctx, this->hphi, this->dmx); - set_memory_complex_op()(this->ctx, this->hphi, 0, this->dmx); + resmem_complex_op()(this->ctx, this->hphi, this->dmx); + setmem_complex_op()(this->ctx, this->hphi, 0, this->dmx); // this->sphi.resize(this->dmx, ModuleBase::ZERO); - resize_memory_complex_op()(this->ctx, this->sphi, this->dmx); - set_memory_complex_op()(this->ctx, this->sphi, 0, this->dmx); + resmem_complex_op()(this->ctx, this->sphi, this->dmx); + setmem_complex_op()(this->ctx, this->sphi, 0, this->dmx); this->cg = new psi::Psi, Device>(phi, 1, 1); // this->scg.resize(this->dmx, ModuleBase::ZERO); - resize_memory_complex_op()(this->ctx, this->scg, this->dmx); - set_memory_complex_op()(this->ctx, this->scg, 0, this->dmx); + resmem_complex_op()(this->ctx, this->scg, this->dmx); + setmem_complex_op()(this->ctx, this->scg, 0, this->dmx); // this->pphi.resize(this->dmx, ModuleBase::ZERO); - resize_memory_complex_op()(this->ctx, this->pphi, this->dmx); - set_memory_complex_op()(this->ctx, this->pphi, 0, this->dmx); + resmem_complex_op()(this->ctx, this->pphi, this->dmx); + setmem_complex_op()(this->ctx, this->pphi, 0, this->dmx); //in band_by_band CG method, only the first band in phi_m would be calculated psi::Range cg_hpsi_range(0); // this->gradient.resize(this->dmx, ModuleBase::ZERO); - resize_memory_complex_op()(this->ctx, this->gradient, this->dmx); - set_memory_complex_op()(this->ctx, this->gradient, 0, this->dmx); + resmem_complex_op()(this->ctx, this->gradient, this->dmx); + setmem_complex_op()(this->ctx, this->gradient, 0, this->dmx); // this->g0.resize(this->dmx, ModuleBase::ZERO); - resize_memory_complex_op()(this->ctx, this->g0, this->dmx); - set_memory_complex_op()(this->ctx, this->g0, 0, this->dmx); + resmem_complex_op()(this->ctx, this->g0, this->dmx); + setmem_complex_op()(this->ctx, this->g0, 0, this->dmx); // this->lagrange.resize(this->n_band, ModuleBase::ZERO); - resize_memory_complex_op()(this->ctx, this->lagrange, this->n_band); - set_memory_complex_op()(this->ctx, this->lagrange, 0, this->n_band); + resmem_complex_op()(this->ctx, this->lagrange, this->n_band); + setmem_complex_op()(this->ctx, this->lagrange, 0, this->n_band); for (int m = 0; m < this->n_band; m++) { @@ -107,12 +101,12 @@ void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, const std::complex* psi_m_in = &(phi(m, 0)); auto pphi_m = this->phi_m->get_pointer(); // haozhihan replace COPY_ARRAY - psi::memory::synchronize_memory_op, Device, Device>()(this->ctx, this->ctx, pphi_m, psi_m_in, this->dim); + syncmem_complex_op()(this->ctx, this->ctx, pphi_m, psi_m_in, this->dim); // ModuleBase::GlobalFunc::COPYARRAY(psi_m_in, pphi_m, this->dim); } - phm_in->sPsi(this->phi_m->get_pointer(), this->sphi, (size_t)this->dim); // sphi = S|psi(m)> + phm_in->sPsi(this->phi_m->get_pointer(), this->sphi, static_cast(this->dim)); // sphi = S|psi(m)> this->schmit_orth(m, phi); - phm_in->sPsi(this->phi_m->get_pointer(), this->sphi, (size_t)this->dim); // sphi = S|psi(m)> + phm_in->sPsi(this->phi_m->get_pointer(), this->sphi, static_cast(this->dim)); // sphi = S|psi(m)> //do hPsi, actually the result of hpsi stored in Operator, //the necessary of copying operation should be checked later @@ -149,9 +143,8 @@ void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, } // end iter std::complex* psi_temp = &(phi(m, 0)); - // haozhihan replace COPY_ARRAY - psi::memory::synchronize_memory_op, Device, Device>()(this->ctx, this->ctx, psi_temp, this->phi_m->get_pointer(), this->dim); // ModuleBase::GlobalFunc::COPYARRAY(this->phi_m->get_pointer(), psi_temp, this->dim); + syncmem_complex_op()(this->ctx, this->ctx, psi_temp, this->phi_m->get_pointer(), this->dim); if (!converged) { @@ -178,33 +171,25 @@ void DiagoCG::diag_mock(hamilt::Hamilt *phm_in, // last calculated eigenvalue should be in the i-th position: reorder FPTYPE e0 = eigenvalue[m]; - // haozhihan replace COPY_ARRAY - psi::memory::synchronize_memory_op, Device, Device>()(this->ctx, this->ctx, pphi, psi_temp, this->dim); // ModuleBase::GlobalFunc::COPYARRAY(psi_temp, pphi, this->dim); + syncmem_complex_op()(this->ctx, this->ctx, pphi, psi_temp, this->dim); for (int j = m; j >= i + 1; j--) { eigenvalue[j] = eigenvalue[j - 1]; std::complex* phi_j = &phi(j, 0); std::complex* phi_j1 = &phi(j-1, 0); - // haozhihan replace COPY_ARRAY - psi::memory::synchronize_memory_op, Device, Device>()(this->ctx, this->ctx, phi_j, phi_j1, this->dim); // ModuleBase::GlobalFunc::COPYARRAY(phi_j1, phi_j, this->dim); + syncmem_complex_op()(this->ctx, this->ctx, phi_j, phi_j1, this->dim); } eigenvalue[i] = e0; // dcopy(pphi, phi, i); std::complex* phi_pointer = &phi(i, 0); - // haozhihan replace COPY_ARRAY - psi::memory::synchronize_memory_op, Device, Device>()(this->ctx, this->ctx, phi_pointer, pphi, this->dim); // ModuleBase::GlobalFunc::COPYARRAY(pphi, phi_pointer, this->dim); - - // this procedure should be good if only a few inversions occur, - // extremely inefficient if eigenvectors are often in bad order - // (but this should not happen) + syncmem_complex_op()(this->ctx, this->ctx, phi_pointer, pphi, this->dim); } // endif } // end reorder - } // end m avg /= this->n_band; @@ -273,108 +258,52 @@ void DiagoCG::orthogonal_gradient(hamilt::Hamilt // int inc = 1; //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // haozhihan replace 2022-10-07 - gemv_op()(this->ctx, - 'C', - this->dim, - m, - &ModuleBase::ONE, - eigenfunction.get_pointer(), - this->dmx, - this->scg, - 1, - &ModuleBase::ZERO, - this->lagrange, - 1); - //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - // qianrui replace 2021-3-15 - // char trans = 'C'; - // zgemv_(&trans, - // &(this->dim), - // &m, - // &ModuleBase::ONE, - // eigenfunction.get_pointer(), - // &(this->dmx), - // this->scg, - // &inc, - // &ModuleBase::ZERO, - // this->lagrange, - // &inc); - //====================================================================== - /*for (int i=0; i>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + gemv_op()( + this->ctx, + 'C', + this->dim, + m, + &ModuleBase::ONE, + eigenfunction.get_pointer(), + this->dmx, + this->scg, + 1, + &ModuleBase::ZERO, + this->lagrange, + 1); Parallel_Reduce::reduce_complex_double_pool(this->lagrange, m); // (3) orthogonal |g> and |scg> to all states (0~m-1) //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // haozhihan replace 2022-10-07 - gemv_op()(this->ctx, - 'N', - this->dim, - m, - &ModuleBase::NEG_ONE, - eigenfunction.get_pointer(), - this->dmx, - this->lagrange, - 1, - &ModuleBase::ONE, - this->gradient, - 1); - gemv_op()(this->ctx, - 'N', - this->dim, - m, - &ModuleBase::NEG_ONE, - eigenfunction.get_pointer(), - this->dmx, - this->lagrange, - 1, - &ModuleBase::ONE, - this->scg, - 1); - //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - // qianrui replace 2021-3-15 - // char trans2 = 'N'; - // zgemv_(&trans2, - // &(this->dim), - // &m, - // &ModuleBase::NEG_ONE, - // eigenfunction.get_pointer(), - // &(this->dmx), - // this->lagrange, - // &inc, - // &ModuleBase::ONE, - // this->gradient, - // &inc); - // zgemv_(&trans2, - // &(this->dim), - // &m, - // &ModuleBase::NEG_ONE, - // eigenfunction.get_pointer(), - // &(this->dmx), - // this->lagrange, - // &inc, - // &ModuleBase::ONE, - // this->scg, - // &inc); - //====================================================================== - /*for (int i=0; i oo = lagrange[i] * eigenfunction(i, j); - g[j] -= oo; - scg[j] -= oo; - } - }*/ - //>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + gemv_op()( + this->ctx, + 'N', + this->dim, + m, + &ModuleBase::NEG_ONE, + eigenfunction.get_pointer(), + this->dmx, + this->lagrange, + 1, + &ModuleBase::ONE, + this->gradient, + 1); + + gemv_op()( + this->ctx, + 'N', + this->dim, + m, + &ModuleBase::NEG_ONE, + eigenfunction.get_pointer(), + this->dmx, + this->lagrange, + 1, + &ModuleBase::ONE, + this->scg, + 1); } template @@ -423,9 +352,8 @@ void DiagoCG::calculate_gamma_cg(const int iter, FPTYPE &gg_last gg_last = gg_now; // (50) cg direction first value : |g> // |cg> = |g> - // haozhihan replace COPY_ARRAY - psi::memory::synchronize_memory_op, Device, Device>()(this->ctx, this->ctx, pcg, this->gradient, this->dim); // ModuleBase::GlobalFunc::COPYARRAY(this->gradient, pcg, this->dim); + syncmem_complex_op ()(this->ctx, this->ctx, pcg, this->gradient, this->dim); } else { @@ -448,13 +376,13 @@ void DiagoCG::calculate_gamma_cg(const int iter, FPTYPE &gg_last std::complex znorma(norma * -1, 0.0); // haozhihan replace this 2022-10-6 - axpy_op()(this->ctx, this->dim, &znorma, pphi_m, 1, pcg, 1); // const int one = 1; // zaxpy_(&this->dim, &znorma, pphi_m, &one, pcg, &one); /*for (int i = 0; i < this->dim; i++) { pcg[i] -= norma * pphi_m[i]; }*/ + axpy_op()(this->ctx, this->dim, &znorma, pphi_m, 1, pcg, 1); } } @@ -494,10 +422,6 @@ bool DiagoCG::update_psi(FPTYPE &cg_norm, FPTYPE &theta, FPTYPE const FPTYPE cost = cos(theta); const FPTYPE sint_norm = sin(theta) / cg_norm; - // std::cout << "\n cg_norm = " << this->ddot(dim, cg, cg); - // std::cout << "\n cg_norm_fac = "<< cg_norm * cg_norm; - // std::cout << "\n overlap = " << this->ddot(dim, phi_m, phi_m); - auto pcg = this->cg->get_pointer(); // for (int i = 0; i < this->dim; i++) // { @@ -526,8 +450,6 @@ bool DiagoCG::update_psi(FPTYPE &cg_norm, FPTYPE &theta, FPTYPE // haozhihan replace this 2022-10-6 constantvector_addORsub_constantVector_op()(this->ctx, this->dim, this->sphi, this->sphi, cost, this->scg, sint_norm); constantvector_addORsub_constantVector_op()(this->ctx, this->dim, this->hphi, this->hphi, cost, this->pphi, sint_norm); - - // ModuleBase::timer::tick("DiagoCG","update"); return 0; } } @@ -548,48 +470,24 @@ void DiagoCG::schmit_orth( assert(psi.get_nbands() >= m); std::complex * lagrange_so = nullptr; - resize_memory_complex_op()(this->ctx, lagrange_so, m + 1); + resmem_complex_op()(this->ctx, lagrange_so, m + 1); //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // haozhihan replace 2022-10-6 int inc = 1; - gemv_op()(this->ctx, - 'C', - this->dim, - (m + 1), - &ModuleBase::ONE, - psi.get_pointer(), - this->dmx, - this->sphi, - inc, - &ModuleBase::ZERO, - lagrange_so, - inc); - //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - // qianrui replace 2021-3-15 - // int inc = 1; - // int mp1 = m + 1; - // char trans = 'C'; - // zgemv_(&trans, - // &(this->dim), - // &mp1, - // &ModuleBase::ONE, - // psi.get_pointer(), - // &(this->dmx), - // this->sphi, - // &inc, - // &ModuleBase::ZERO, - // lagrange_so.data(), - // &inc); - //====================================================================== - /*for (int j = 0; j <= m; j++) - { - for (int ig=0; ig < dim; ig++) - { - lagrange_so[j] += conj(psi( j, ig)) * sphi[ig] ; - } - }*/ - //>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + gemv_op()( + this->ctx, + 'C', + this->dim, + (m + 1), + &ModuleBase::ONE, + psi.get_pointer(), + this->dmx, + this->sphi, + inc, + &ModuleBase::ZERO, + lagrange_so, + inc); // be careful , here reduce m+1 Parallel_Reduce::reduce_complex_double_pool(lagrange_so, m + 1); @@ -600,34 +498,20 @@ void DiagoCG::schmit_orth( //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // haozhihan replace 2022-10-6 - gemv_op()(this->ctx, - 'N', - this->dim, - m, - &ModuleBase::NEG_ONE, - psi.get_pointer(), - this->dmx, - lagrange_so, - inc, - &ModuleBase::ONE, - this->phi_m->get_pointer(), - inc); - //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< - // qianrui replace 2021-3-15 - // char trans2 = 'N'; - // zgemv_(&trans2, - // &(this->dim), - // &m, - // &ModuleBase::NEG_ONE, - // psi.get_pointer(), - // &(this->dmx), - // lagrange_so.data(), - // &inc, - // &ModuleBase::ONE, - // this->phi_m->get_pointer(), - // &inc); + gemv_op()( + this->ctx, + 'N', + this->dim, + m, + &ModuleBase::NEG_ONE, + psi.get_pointer(), + this->dmx, + lagrange_so, + inc, + &ModuleBase::ONE, + this->phi_m->get_pointer(), + inc); - psi_norm -= hsolver::zdot_real_op()(this->ctx, m, lagrange_so, lagrange_so, false); //====================================================================== /*for (int j = 0; j < m; j++) { @@ -638,6 +522,7 @@ void DiagoCG::schmit_orth( psi_norm -= ( conj(lagrange[j]) * lagrange[j] ).real(); }*/ //>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> + psi_norm -= hsolver::zdot_real_op()(this->ctx, m, lagrange_so, lagrange_so, false); if (psi_norm <= 0.0) { @@ -657,16 +542,16 @@ void DiagoCG::schmit_orth( auto pphi_m = this->phi_m->get_pointer(); //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< // haozhihan replace 2022-10-6 - vector_div_constant_op()(this->ctx, this->dim, pphi_m, pphi_m, psi_norm); // scal_op()(this->ctx, this->dim, &psi_norm, pphi_m, 1); //====================================================================== // for (int ig = 0; ig < this->dim; ig++) // { // pphi_m[ig] /= psi_norm; // } + vector_div_constant_op()(this->ctx, this->dim, pphi_m, pphi_m, psi_norm); // ModuleBase::timer::tick("DiagoCG","schmit_orth"); - delete_memory_complex_op()(this->ctx, lagrange_so); + delmem_complex_op()(this->ctx, lagrange_so); } template @@ -676,7 +561,7 @@ void DiagoCG::diag(hamilt::Hamilt *phm_in, psi:: int ntry = 0; this->notconv = 0; if (this->device == psi::GpuDevice) { - resize_memory_var_op()(this->ctx, this->d_precondition, psi.get_nbasis()); + resmem_var_op()(this->ctx, this->d_precondition, psi.get_nbasis()); syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, this->d_precondition, this->precondition, psi.get_nbasis()); } do diff --git a/source/module_hsolver/diago_cg.h b/source/module_hsolver/diago_cg.h index 4a04a48a03b..3690595d47d 100644 --- a/source/module_hsolver/diago_cg.h +++ b/source/module_hsolver/diago_cg.h @@ -94,13 +94,14 @@ class DiagoCG : public DiagH using hpsi_info = typename hamilt::Operator, Device>::hpsi_info; using zdot_real_op = hsolver::zdot_real_op; - using set_memory_complex_op = psi::memory::set_memory_op, Device>; - using delete_memory_complex_op = psi::memory::delete_memory_op, Device>; - using resize_memory_complex_op = psi::memory::resize_memory_op, Device>; + using setmem_complex_op = psi::memory::set_memory_op, Device>; + using delmem_complex_op = psi::memory::delete_memory_op, Device>; + using resmem_complex_op = psi::memory::resize_memory_op, Device>; + using syncmem_complex_op = psi::memory::synchronize_memory_op, Device, Device>; using syncmem_complex_d2h_op = psi::memory::synchronize_memory_op, psi::DEVICE_CPU, Device>; - using set_memory_var_cpu_op = psi::memory::set_memory_op; - using resize_memory_var_op = psi::memory::resize_memory_op; + using resmem_var_op = psi::memory::resize_memory_op; + using setmem_var_h_op = psi::memory::set_memory_op; using syncmem_var_h2d_op = psi::memory::synchronize_memory_op; }; diff --git a/source/module_hsolver/diago_david.cpp b/source/module_hsolver/diago_david.cpp index 5ef39a43c6f..dfda95ab99d 100644 --- a/source/module_hsolver/diago_david.cpp +++ b/source/module_hsolver/diago_david.cpp @@ -195,373 +195,6 @@ void DiagoDavid::cal_grad(hamilt::Hamilt* phm_in return; } -#if ((defined __CUDA) || (defined __ROCM)) -template<> -void DiagoDavid::cal_grad(hamilt::Hamilt* phm_in, - const int &npw, - const int &nbase, // current dimension of the reduced basis - const int ¬conv, - psi::Psi, psi::DEVICE_GPU> &basis, - ModuleBase::ComplexMatrix &hp, - ModuleBase::ComplexMatrix &sp, - const ModuleBase::ComplexMatrix &vc, - const int *unconv, - const double *eigenvalue) -{ - if (test_david == 1) - ModuleBase::TITLE("DiagoDavid", "cal_grad"); - if(notconv == 0) return; - ModuleBase::timer::tick("DiagoDavid", "cal_grad"); - - //use template pointer for accelerate - std::complex *hpsi; - std::complex *spsi; - std::complex *ppsi; - - // expand the reduced basis set with the new basis vectors P|R(psi)>... - // in which psi are the last eigenvectors - // we define |R(psi)> as (H-ES)*|Psi>, E = / - ModuleBase::ComplexMatrix vc_ev_vector(notconv, nbase); - for (int m = 0; m < notconv; m++) - { - for(int i = 0; i < nbase; i++) - { - vc_ev_vector(m, i) = vc(i, unconv[m]); - } - } - ppsi = &basis(nbase, 0); - int inc = 1; - char trans = 'N'; - char transb = 'N'; - zgemm_(&trans, - &transb, - &npw, // m: row of A,C - ¬conv, // n: col of B,C - &nbase, // k: col of A, row of B - &ModuleBase::ONE, // alpha - hp.c, // A - &hp.nc, // LDA: if(N) max(1,m) if(T) max(1,k) - vc_ev_vector.c, // B - &vc_ev_vector.nc, // LDB: if(N) max(1,k) if(T) max(1,n) - &ModuleBase::ZERO, // belta - ppsi, // C - &basis.get_nbasis()); // LDC: if(N) max(1, m) - /*zgemv_(&trans, - &npw, - &nbase, - &ModuleBase::ONE, - hp.c, - &hp.nc, - vc_ev_vector.data(), - &inc, - &ModuleBase::ZERO, - respsi, - &inc);*/ - for (int m = 0; m < notconv; m++) - { - for(int i = 0; i < nbase; i++) - { - vc_ev_vector(m, i) *= -1 * eigenvalue[unconv[m]]; - } - } - zgemm_(&trans, - &transb, - &npw, // m: row of A,C - ¬conv, // n: col of B,C - &nbase, // k: col of A, row of B - &ModuleBase::ONE, // alpha - sp.c, // A - &sp.nc, // LDA: if(N) max(1,m) if(T) max(1,k) - vc_ev_vector.c, // B - &vc_ev_vector.nc, // LDB: if(N) max(1,k) if(T) max(1,n) - &ModuleBase::ONE, // belta - ppsi, // C - &basis.get_nbasis()); // LDC: if(N) max(1, m) - /*zgemv_(&trans, - &npw, - &nbase, - &ModuleBase::ONE, - sp.c, - &sp.nc, - vc_ev_vector.data(), - &inc, - &ModuleBase::ONE, - respsi, - &inc);*/ - - /*ModuleBase::GlobalFunc::ZEROS(respsi, npw); - for (int i = 0; i < nbase; i++) - { - hpsi = &(hp(i, 0)); - spsi = &(sp(i, 0)); - auto vc_value = vc(i, unconv[m]); - auto ev_value = eigenvalue[unconv[m]]; - for (int ig = 0; ig < npw; ig++) - { - respsi[ig] += vc_value * (hpsi[ig] - ev_value * spsi[ig]); - } - }*/ - - for (int m = 0; m < notconv; m++) - { - ppsi = &basis(nbase + m, 0); - for (int ig = 0; ig < npw; ig++) - { - ppsi[ig] /= this->precondition[ig]; - } - } - - //there is a nbase to nbase + notconv band orthogonalise - //plan for SchmitOrth - ModuleBase::ComplexMatrix lagrange_matrix(notconv, nbase + notconv); - std::vector pre_matrix_mm_m(notconv, 0); - std::vector pre_matrix_mv_m(notconv, 1); - this->planSchmitOrth(notconv, pre_matrix_mm_m.data(), pre_matrix_mv_m.data()); - for( int m = 0; m < notconv; m++) - { - phm_in->sPsi(&basis(nbase + m, 0), &sp(nbase + m, 0), (size_t)npw); - } - //first nbase bands psi* dot notconv bands spsi to prepare lagrange_matrix - trans = 'C'; - transb = 'N'; - //calculate the square matrix for future lagranges - zgemm_(&trans, - &transb, - &nbase, // m: row of A,C - ¬conv, // n: col of B,C - &npw, // k: col of A, row of B - &ModuleBase::ONE, // alpha - &basis(0, 0), // A - &basis.get_nbasis(), // LDA: if(N) max(1,m) if(T) max(1,k) - &sp(nbase, 0), // B - &sp.nc, // LDB: if(N) max(1,k) if(T) max(1,n) - &ModuleBase::ZERO, // belta - &lagrange_matrix(0, 0), // C - &lagrange_matrix.nc); // LDC: if(N) max(1, m) - - for (int m = 0; m < notconv; m++) - { - ppsi = &basis(nbase + m, 0); - spsi = &sp(nbase + m, 0); - - this->SchmitOrth( - npw, - nbase + notconv, - nbase + m, - basis, - sp, - &lagrange_matrix(m, 0), - pre_matrix_mm_m[m], - pre_matrix_mv_m[m] - ); - phm_in->sPsi(ppsi, spsi, (size_t)npw); - - } - //calculate H|psi> for not convergence bands - hpsi_info_gpu dav_hpsi_in(&basis, psi::Range(1, 0, nbase, nbase + notconv-1), &hp(nbase, 0)); - phm_in->ops->hPsi_gpu(dav_hpsi_in); - - ModuleBase::timer::tick("DiagoDavid", "cal_grad"); -} - - -template<> -void DiagoDavid::diag_mock(hamilt::Hamilt* phm_in, psi::Psi, psi::DEVICE_GPU> &psi, double* eigenvalue_in) -{ - if (test_david == 1) - ModuleBase::TITLE("DiagoDavid", "diag_mock"); - ModuleBase::timer::tick("DiagoDavid", "diag_mock"); - - const int dim = psi.get_current_nbas(); - const int nband = psi.get_nbands(); - - assert(DiagoDavid::PW_DIAG_NDIM > 1); - assert(DiagoDavid::PW_DIAG_NDIM * nband < dim * GlobalV::NPROC_IN_POOL); - // qianrui change it 2021-7-25. - // In strictly speaking, it shoule be PW_DIAG_NDIM*nband < npw sum of all pools. We roughly estimate it here. - // However, in most cases, total number of plane waves should be much larger than nband*PW_DIAG_NDIM - - int nbase_x = DiagoDavid::PW_DIAG_NDIM * nband; // maximum dimension of the reduced basis set - - psi::Psi, psi::DEVICE_GPU> basis(1, nbase_x, dim, &(psi.get_ngk(0))); // the reduced basis set - ModuleBase::ComplexMatrix hp(nbase_x, dim); // the product of H and psi in the reduced basis set - ModuleBase::ComplexMatrix sp(nbase_x, dim); // the Product of S and psi in the reduced basis set - - ModuleBase::ComplexMatrix hc(nbase_x, nbase_x); // Hamiltonian on the reduced basis - ModuleBase::ComplexMatrix sc(nbase_x, nbase_x); // Overlap on the reduced basis - ModuleBase::ComplexMatrix vc(nbase_x, nbase_x); // Eigenvectors of hc - std::vector eigenvalue(nbase_x); // the lowest N eigenvalues of hc - - std::vector convflag(nband, false); // convflag[m] = true if the m th band is convergent - std::vector unconv(nband); // unconv[m] store the number of the m th unconvergent band - - int nbase = 0; // the dimension of the reduced basis set - this->notconv = nband; // the number of the unconvergent bands - for (int m = 0; m < nband; m++) - unconv[m] = m; - - ModuleBase::timer::tick("DiagoDavid", "first"); - // orthogonalise the initial trial psi(0~nband-1) - //plan for SchmitOrth - ModuleBase::ComplexMatrix lagrange_matrix(nband, nband); - std::vector pre_matrix_mm_m(nband, 0); - std::vector pre_matrix_mv_m(nband, 1); - this->planSchmitOrth(nband, pre_matrix_mm_m.data(), pre_matrix_mv_m.data()); - for( int m = 0; m < nband; m++) - { - phm_in->sPsi(&psi(m, 0), &sp(m, 0), (size_t)dim); - } - //begin SchmitOrth - for (int m = 0; m < nband; m++) - { - // psi_m = psi(m) - ModuleBase::GlobalFunc::COPYARRAY(&psi(m, 0), &basis(m, 0), dim); - /*for (int ig = 0; ig < dim; ig++) - { - psi_m[ig] = psi(m, ig); - }*/ - - //phm_in->sPsi(psi_m.data(), spsi.data(), (size_t)dim); - this->SchmitOrth( - dim, - nband, - m, - basis, - sp, - &lagrange_matrix(m, 0), - pre_matrix_mm_m[m], - pre_matrix_mv_m[m] - ); - phm_in->sPsi(&basis(m, 0), &sp(m, 0), (size_t)dim); - - // basis(m) = psi_m, hp(m) = H |psi_m>, sp(m) = S |psi_m> - //ModuleBase::GlobalFunc::COPYARRAY(psi_m.data(), &basis(m, 0), dim); - //ModuleBase::GlobalFunc::COPYARRAY(spsi.data(), &sp(m, 0), dim); - /*std::complex* sp_p = &sp(m, 0); - std::complex* basis_p = &basis(m, 0); - for (int ig = 0; ig < dim; ig++) - { - basis_p[ig] = psi_m[ig]; - //hp(m, ig) = hpsi[ig]; - sp_p[ig] = spsi[ig]; - }*/ - } - //end of SchmitOrth and calculate H|psi> - hpsi_info_gpu dav_hpsi_in(&basis, psi::Range(1, 0, 0, nband-1), &hp(0, 0)); - phm_in->ops->hPsi_gpu(dav_hpsi_in); - - hc.zero_out(); - sc.zero_out(); - - this->cal_elem(dim, nbase, this->notconv, basis, hp, sp, hc, sc); - // TODO diag_zhegvx - this->diag_zhegvx(nbase, nband, hc, sc, nbase_x, eigenvalue.data(), vc); - - for (int m = 0; m < nband; m++) - { - eigenvalue_in[m] = eigenvalue[m]; - } - - ModuleBase::timer::tick("DiagoDavid", "first"); - - int dav_iter = 0; - do - { - dav_iter++; - - this->cal_grad(phm_in, - dim, - nbase, - this->notconv, - basis, - hp, - sp, - vc, - unconv.data(), - eigenvalue.data()); - - this->cal_elem(dim, nbase, this->notconv, basis, hp, sp, hc, sc); - - this->diag_zhegvx(nbase, nband, hc, sc, nbase_x, eigenvalue.data(), vc); - - // check convergence and update eigenvalues - ModuleBase::timer::tick("DiagoDavid", "check_update"); - - this->notconv = 0; - for (int m = 0; m < nband; m++) - { - convflag[m] = (abs(eigenvalue[m] - eigenvalue_in[m]) < DiagoIterAssist::PW_DIAG_THR); - - if (!convflag[m]) - { - unconv[this->notconv] = m; - this->notconv++; - } - - eigenvalue_in[m] = eigenvalue[m]; - } - - ModuleBase::timer::tick("DiagoDavid", "check_update"); - if (!this->notconv || (nbase + this->notconv > nbase_x) || (dav_iter == DiagoIterAssist::PW_DIAG_NMAX)) - { - ModuleBase::timer::tick("DiagoDavid", "last"); - - // updata eigenvectors of Hamiltonian - ModuleBase::GlobalFunc::ZEROS(psi.get_pointer(), psi.get_nbands() * psi.get_nbasis()); - char transa = 'N'; - char transb = 'T'; - zgemm_(&transa, - &transb, - &dim, // m: row of A,C - &nband, // n: col of B,C - &nbase, // k: col of A, row of B - &ModuleBase::ONE, // alpha - basis.get_pointer(), // A - &basis.get_nbasis(), // LDA: if(N) max(1,m) if(T) max(1,k) - vc.c, // B - &nbase_x, // LDB: if(N) max(1,k) if(T) max(1,n) - &ModuleBase::ZERO, // belta - psi.get_pointer(), // C - &psi.get_nbasis()); // LDC: if(N) max(1, m) - /*for (int m = 0; m < nband; m++) - { - for (int j = 0; j < nbase; j++) - { - for (int ig = 0; ig < dim; ig++) - { - psi(m, ig) += vc(j, m) * basis(j, ig); - } - } - }*/ - - if (!this->notconv || (dav_iter == DiagoIterAssist::PW_DIAG_NMAX)) - { - // overall convergence or last iteration: exit the iteration - - ModuleBase::timer::tick("DiagoDavid", "last"); - break; - } - else - { - // if the dimension of the reduced basis set is becoming too large, - // then replace the first N (=nband) basis vectors with the current - // estimate of the eigenvectors and set the basis dimension to N; - - this->refresh(dim, nband, nbase, eigenvalue_in, psi, basis, hp, sp, hc, sc, vc); - ModuleBase::timer::tick("DiagoDavid", "last"); - } - - } // end of if - - } while (1); - - DiagoIterAssist::avg_iter += static_cast(dav_iter); - - ModuleBase::timer::tick("DiagoDavid", "diag_mock"); - return; -} -#endif // ((defined __CUDA) || (defined __ROCM)) - template void DiagoDavid::diag_mock(hamilt::Hamilt* phm_in, psi::Psi, Device> &psi, FPTYPE *eigenvalue_in) { diff --git a/source/module_hsolver/diago_david.h b/source/module_hsolver/diago_david.h index 892fadbb1ef..36e0531d533 100644 --- a/source/module_hsolver/diago_david.h +++ b/source/module_hsolver/diago_david.h @@ -103,10 +103,7 @@ class DiagoDavid : public DiagH const FPTYPE* precondition = nullptr; - using hpsi_info = typename hamilt::Operator, psi::DEVICE_CPU>::hpsi_info; -#if ((defined __CUDA) || (defined __ROCM)) - using hpsi_info_gpu = typename hamilt::Operator, psi::DEVICE_GPU>::hpsi_info_gpu; -#endif + using hpsi_info = typename hamilt::Operator, Device>::hpsi_info; }; } // namespace hsolver diff --git a/source/module_hsolver/diago_iter_assist.cpp b/source/module_hsolver/diago_iter_assist.cpp index d122386bf22..1633aeef335 100644 --- a/source/module_hsolver/diago_iter_assist.cpp +++ b/source/module_hsolver/diago_iter_assist.cpp @@ -9,6 +9,7 @@ #include "src_parallel/parallel_reduce.h" #include "module_hsolver/include/math_kernel.h" #include "module_hsolver/include/dngvd_op.h" +#include "module_psi/include/device.h" using namespace hsolver; @@ -48,17 +49,13 @@ void DiagoIterAssist::diagH_subspace( n_band = nstart; assert(n_band <= nstart); - Device* ctx = {}; - - std::complex* hcc = nullptr; - std::complex* scc = nullptr; - std::complex* vcc = nullptr; - psi::memory::resize_memory_op, Device>()(ctx, hcc, nstart * nstart); - psi::memory::resize_memory_op, Device>()(ctx, scc, nstart * nstart); - psi::memory::resize_memory_op, Device>()(ctx, vcc, nstart * nstart); - psi::memory::set_memory_op, Device>()(ctx, hcc, 0, nstart * nstart); - psi::memory::set_memory_op, Device>()(ctx, scc, 0, nstart * nstart); - psi::memory::set_memory_op, Device>()(ctx, vcc, 0, nstart * nstart); + std::complex* hcc = nullptr, * scc = nullptr, * vcc = nullptr; + resmem_complex_op()(ctx, hcc, nstart * nstart); + resmem_complex_op()(ctx, scc, nstart * nstart); + resmem_complex_op()(ctx, vcc, nstart * nstart); + setmem_complex_op()(ctx, hcc, 0, nstart * nstart); + setmem_complex_op()(ctx, scc, 0, nstart * nstart); + setmem_complex_op()(ctx, vcc, 0, nstart * nstart); const int dmin = psi.get_current_nbas(); const int dmax = psi.get_nbasis(); @@ -69,8 +66,8 @@ void DiagoIterAssist::diagH_subspace( // allocated hpsi // std::vector> hpsi(psi.get_nbands() * psi.get_nbasis()); std::complex* hphi = nullptr; - psi::memory::resize_memory_op, Device>()(ctx, hphi, psi.get_nbands() * psi.get_nbasis()); - psi::memory::set_memory_op, Device>()(ctx, hphi, 0, psi.get_nbands() * psi.get_nbasis()); + resmem_complex_op()(ctx, hphi, psi.get_nbands() * psi.get_nbasis()); + setmem_complex_op()(ctx, hphi, 0, psi.get_nbands() * psi.get_nbasis()); // do hPsi for all bands psi::Range all_bands_range(1, psi.get_current_k(), 0, psi.get_nbands()-1); hpsi_info hpsi_in(&psi, all_bands_range, hphi); @@ -133,10 +130,6 @@ void DiagoIterAssist::diagH_subspace( // after generation of H and S matrix, diag them DiagoIterAssist::diagH_LAPACK(nstart, n_band, hcc, scc, nstart, en, vcc); -#if defined(__CUDA) || defined(__ROCM) - hsolver::createBLAShandle(); -#endif - //======================= // diagonize the H-matrix //======================= @@ -174,8 +167,8 @@ void DiagoIterAssist::diagH_subspace( // create a temporary matrix to store the result. (by wangjp) // qianrui improve this part 2021-3-13 std::complex* evctemp = nullptr; - psi::memory::resize_memory_op, Device>()(ctx, evctemp, n_band * dmin); - psi::memory::set_memory_op, Device>()(ctx, evctemp, 0, n_band * dmin); + resmem_complex_op()(ctx, evctemp, n_band * dmin); + setmem_complex_op()(ctx, evctemp, 0, n_band * dmin); gemm_op()( ctx, @@ -202,22 +195,23 @@ void DiagoIterAssist::diagH_subspace( // evc(ib, ig) = evctmp(ib, ig); // } // } - psi::memory::delete_memory_op, Device>()(ctx, evctemp); + delmem_complex_op()(ctx, evctemp); } - psi::memory::delete_memory_op, Device>()(ctx, hcc); - psi::memory::delete_memory_op, Device>()(ctx, scc); - psi::memory::delete_memory_op, Device>()(ctx, vcc); - psi::memory::delete_memory_op, Device>()(ctx, hphi); + delmem_complex_op()(ctx, hcc); + delmem_complex_op()(ctx, scc); + delmem_complex_op()(ctx, vcc); + delmem_complex_op()(ctx, hphi); ModuleBase::timer::tick("DiagoIterAssist", "diagH_subspace"); - return; } template void DiagoIterAssist::diagH_subspace_init( hamilt::Hamilt* pHamilt, - const ModuleBase::ComplexMatrix &psi, + const std::complex* psi, + int psi_nr, + int psi_nc, psi::Psi, Device> &evc, FPTYPE *en) { @@ -227,68 +221,100 @@ void DiagoIterAssist::diagH_subspace_init( // two case: // 1. pw base: nstart = n_band, psi(nbands * npwx) // 2. lcao_in_pw base: nstart >= n_band, psi(NLOCAL * npwx) - const int nstart = psi.nr; + + const int nstart = psi_nr; const int n_band = evc.get_nbands(); - ModuleBase::ComplexMatrix hc(nstart, nstart); - ModuleBase::ComplexMatrix sc(nstart, nstart); - ModuleBase::ComplexMatrix hvec(nstart, n_band); + Device* ctx = {}; + + // ModuleBase::ComplexMatrix hc(nstart, nstart); + // ModuleBase::ComplexMatrix sc(nstart, nstart); + // ModuleBase::ComplexMatrix hvec(nstart, n_band); + std::complex* hcc = nullptr, * scc = nullptr, * vcc = nullptr; + resmem_complex_op()(ctx, hcc, nstart * nstart); + resmem_complex_op()(ctx, scc, nstart * nstart); + resmem_complex_op()(ctx, vcc, nstart * nstart); + setmem_complex_op()(ctx, hcc, 0, nstart * nstart); + setmem_complex_op()(ctx, scc, 0, nstart * nstart); + setmem_complex_op()(ctx, vcc, 0, nstart * nstart); const int dmin = evc.get_current_nbas(); const int dmax = evc.get_nbasis(); // qianrui improve this part 2021-3-14 - //std::complex *aux = new std::complex[dmax * nstart]; - //const std::complex *paux = aux; - psi::Psi, Device> psi_temp(1, nstart, psi.nc, &evc.get_ngk(0)); - ModuleBase::GlobalFunc::COPYARRAY(psi.c, psi_temp.get_pointer(), psi_temp.size()); + // std::complex *aux = new std::complex[dmax * nstart]; + // const std::complex *paux = aux; + psi::Psi, Device> psi_temp(1, nstart, psi_nc, &evc.get_ngk(0)); + syncmem_complex_op ()(ctx, ctx, psi_temp.get_pointer(), psi, psi_temp.size()); + // ModuleBase::GlobalFunc::COPYARRAY(psi, psi_temp.get_pointer(), psi_temp.size()); + const std::complex *ppsi = psi_temp.get_pointer(); - //allocated hpsi - std::vector> hpsi(psi_temp.get_nbands() * psi_temp.get_nbasis()); - //do hPsi for all bands + // allocated hpsi + std::complex* hpsi = nullptr; + resmem_complex_op()(ctx, hpsi, psi_temp.get_nbands() * psi_temp.get_nbasis()); + setmem_complex_op()(ctx, hpsi, 0, psi_temp.get_nbands() * psi_temp.get_nbasis()); + // ================================================ + // std::vector> hpsi(psi_temp.get_nbands() * psi_temp.get_nbasis()); + + + // do hPsi for all bands psi::Range all_bands_range(1, psi_temp.get_current_k(), 0, psi_temp.get_nbands()-1); - hpsi_info hpsi_in(&psi_temp, all_bands_range, hpsi.data()); + hpsi_info hpsi_in(&psi_temp, all_bands_range, hpsi); pHamilt->ops->hPsi(hpsi_in); - //use aux as a data pointer for hpsi - const std::complex *aux = hpsi.data(); - - char trans1 = 'C'; - char trans2 = 'N'; - zgemm_(&trans1, - &trans2, - &nstart, - &nstart, - &dmin, - &ModuleBase::ONE, - ppsi, - &dmax, - aux, - &dmax, - &ModuleBase::ZERO, - hc.c, - &nstart); - hc = transpose(hc, false); - - zgemm_(&trans1, - &trans2, - &nstart, - &nstart, - &dmin, - &ModuleBase::ONE, - ppsi, - &dmax, - ppsi, - &dmax, - &ModuleBase::ZERO, - sc.c, - &nstart); - sc = transpose(sc, false); + + gemm_op()( + ctx, + 'C', + 'N', + nstart, + nstart, + dmin, + &ModuleBase::ONE, + ppsi, + dmax, + hpsi, + dmax, + &ModuleBase::ZERO, + hcc, + nstart + ); + matrixTranspose_op()( + ctx, + nstart, + nstart, + hcc, + hcc + ); + + gemm_op()( + ctx, + 'C', + 'N', + nstart, + nstart, + dmin, + &ModuleBase::ONE, + ppsi, + dmax, + ppsi, + dmax, + &ModuleBase::ZERO, + scc, + nstart + ); + matrixTranspose_op()( + ctx, + nstart, + nstart, + scc, + scc + ); if (GlobalV::NPROC_IN_POOL > 1) { - Parallel_Reduce::reduce_complex_double_pool(hc.c, nstart * nstart); - Parallel_Reduce::reduce_complex_double_pool(sc.c, nstart * nstart); + Parallel_Reduce::reduce_complex_double_pool(hcc, nstart * nstart); + Parallel_Reduce::reduce_complex_double_pool(scc, nstart * nstart); } // after generation of H and S matrix, diag them @@ -304,7 +330,8 @@ void DiagoIterAssist::diagH_subspace_init( else sc(i,j) = std::complex(double(int(sc(i,j).real()*100000000))/100000000, 0); } }*/ - DiagoIterAssist::diagH_LAPACK(nstart, n_band, hc, sc, nstart, en, hvec); + + DiagoIterAssist::diagH_LAPACK(nstart, n_band, hcc, scc, nstart, en, vcc); //======================= // diagonize the H-matrix @@ -320,198 +347,61 @@ void DiagoIterAssist::diagH_subspace_init( // because psi and evc are different here, // I think if psi and evc are the same, // there may be problems, mohan 2011-01-01 - char transa = 'N'; - char transb = 'T'; - zgemm_(&transa, - &transb, - &dmax, // m: row of A,C - &n_band, // n: col of B,C - &nstart, // k: col of A, row of B - &ModuleBase::ONE, // alpha - ppsi, // A - &dmax, // LDA: if(N) max(1,m) if(T) max(1,k) - hvec.c, // B - &n_band, // LDB: if(N) max(1,k) if(T) max(1,n) - &ModuleBase::ZERO, // belta - evc.get_pointer(), // C - &dmax); // LDC: if(N) max(1, m) + gemm_op()( + ctx, + 'N', + 'T', + dmax, + n_band, + nstart, + &ModuleBase::ONE, + ppsi, + dmax, + vcc, + n_band, + &ModuleBase::ZERO, + evc.get_pointer(), + dmax + ); } else { // As the evc and psi may refer to the same matrix, we first // create a temporary matrix to store the result. (by wangjp) // qianrui improve this part 2021-3-13 - char transa = 'N'; - char transb = 'T'; - ModuleBase::ComplexMatrix evctmp(n_band, dmin, false); - zgemm_(&transa, - &transb, - &dmin, - &n_band, - &nstart, - &ModuleBase::ONE, - ppsi, - &dmax, - hvec.c, - &n_band, - &ModuleBase::ZERO, - evctmp.c, - &dmin); - for (int ib = 0; ib < n_band; ib++) - { - for (int ig = 0; ig < dmin; ig++) - { - evc(ib, ig) = evctmp(ib, ig); - } - } - } - - ModuleBase::timer::tick("DiagoIterAssist", "diagH_subspace"); - return; -} - -#if ((defined __CUDA) || (defined __ROCM)) -//---------------------------------------------------------------------- -// Hamiltonian diagonalization in the subspace spanned -// by nstart states psi (atomic or random wavefunctions). -// Produces on output n_band eigenvectors (n_band <= nstart) in evc. -//---------------------------------------------------------------------- -// template<> -// void DiagoIterAssist::diagH_subspace(hamilt::Hamilt* pHamilt, const psi::Psi, psi::DEVICE_GPU> &psi, psi::Psi, psi::DEVICE_GPU> &evc, double *en, int n_band) { -// ModuleBase::WARNING_QUIT("DiagoIterAssist::diagH_subspace","GPU's implementation is not supported currently!"); -// } - -template<> -void DiagoIterAssist::diagH_subspace_init(hamilt::Hamilt* pHamilt, const ModuleBase::ComplexMatrix &psi, psi::Psi, psi::DEVICE_GPU> &evc, double *en) { - ModuleBase::WARNING_QUIT("DiagoIterAssist::diagH_subspace_init","GPU's implementation is not supported currently!"); -} -#endif - -template -void DiagoIterAssist::diagH_LAPACK( - const int nstart, - const int nbands, - const ModuleBase::ComplexMatrix &hc, - const ModuleBase::ComplexMatrix &sc, - const int ldh, // nstart - FPTYPE *e, - ModuleBase::ComplexMatrix &hvec) -{ - ModuleBase::TITLE("DiagoIterAssist", "diagH_LAPACK"); - ModuleBase::timer::tick("DiagoIterAssist", "diagH_LAPACK"); - - int lwork = 0; - - ModuleBase::ComplexMatrix sdum(nstart, ldh); - ModuleBase::ComplexMatrix hdum; - - sdum = sc; - - const bool all_eigenvalues = (nstart == nbands); - - // workspace query - int nb = LapackConnector::ilaenv(1, "ZHETRD", "U", nstart, -1, -1, -1); - - if (nb < 1) - { - nb = std::max(1, nstart); - } - - if (nb == 1 || nb >= nstart) - { - lwork = 2 * nstart; // mohan modify 2009-08-02 - } - else - { - lwork = (nb + 1) * nstart; - } - - std::complex *work = new std::complex[lwork]; - ModuleBase::GlobalFunc::ZEROS(work, lwork); - - //===================================================================== - // input s and (see below) h are copied so that they are not destroyed - //===================================================================== - - int info = 0; - int rwork_dim; - if (all_eigenvalues) - { - rwork_dim = 3 * nstart - 2; - } - else - { - rwork_dim = 7 * nstart; - } - - FPTYPE *rwork = new FPTYPE[rwork_dim]; - ModuleBase::GlobalFunc::ZEROS(rwork, rwork_dim); - - if (all_eigenvalues) - { - //=========================== - // calculate all eigenvalues - //=========================== - hvec = hc; - LapackConnector::zhegv(1, 'V', 'U', nstart, hvec, ldh, sdum, ldh, e, work, lwork, rwork, info); - } - else - { - //===================================== - // calculate only m lowest eigenvalues - //===================================== - int *iwork = new int[5 * nstart]; - int *ifail = new int[nstart]; - - ModuleBase::GlobalFunc::ZEROS(rwork, 7 * nstart); - ModuleBase::GlobalFunc::ZEROS(iwork, 5 * nstart); - ModuleBase::GlobalFunc::ZEROS(ifail, nstart); - - hdum.create(nstart, ldh); - hdum = hc; - - //============================= - // Number of calculated bands - //============================= - int mm = nbands; - - LapackConnector::zhegvx(1, // INTEGER - 'V', // CHARACTER*1 - 'I', // CHARACTER*1 - 'U', // CHARACTER*1 - nstart, // INTEGER - hdum, // COMPLEX*16 array - ldh, // INTEGER - sdum, // COMPLEX*16 array - ldh, // INTEGER - 0.0, // DOUBLE PRECISION - 0.0, // DOUBLE PRECISION - 1, // INTEGER - nbands, // INTEGER - 0.0, // DOUBLE PRECISION - mm, // INTEGER - e, // DOUBLE PRECISION array - hvec, // COMPLEX*16 array - ldh, // INTEGER - work, // DOUBLE array, dimension (MAX(1,LWORK)) - lwork, // INTEGER - rwork, // DOUBLE PRECISION array, dimension (7*N) - iwork, // INTEGER array, dimension (5*N) - ifail, // INTEGER array, dimension (N) - info // INTEGER + std::complex* evctemp = nullptr; + resmem_complex_op()(ctx, evctemp, n_band * dmin); + setmem_complex_op()(ctx, evctemp, 0, n_band * dmin); + + gemm_op()( + ctx, + 'N', + 'T', + dmin, + n_band, + nstart, + &ModuleBase::ONE, + ppsi, + dmax, + vcc, + n_band, + &ModuleBase::ZERO, + evctemp, + dmin ); - delete[] iwork; - delete[] ifail; + matrixSetToAnother()(ctx, n_band, evctemp, dmin, evc.get_pointer(), dmax); + + delmem_complex_op()(ctx, evctemp); } - delete[] rwork; - delete[] work; - ModuleBase::timer::tick("DiagoIterAssist", "diagH_LAPACK"); - return; + delmem_complex_op()(ctx, hcc); + delmem_complex_op()(ctx, scc); + delmem_complex_op()(ctx, vcc); + delmem_complex_op()(ctx, hpsi); + ModuleBase::timer::tick("DiagoIterAssist", "diagH_subspace"); } - template void DiagoIterAssist::diagH_LAPACK( const int nstart, @@ -525,70 +415,40 @@ void DiagoIterAssist::diagH_LAPACK( ModuleBase::TITLE("DiagoIterAssist", "LAPACK_subspace"); ModuleBase::timer::tick("DiagoIterAssist", "LAPACK_subspace"); - Device* ctx = {}; - const bool all_eigenvalues = (nstart == nbands); -#if defined(__CUDA) || defined(__ROCM) - const psi::DEVICE_CPU * cpu_ctx = {}; - const psi::DEVICE_GPU * gpu_ctx = {}; - - FPTYPE* e_gpu = nullptr; - psi::memory::resize_memory_op()(gpu_ctx, e_gpu, nbands); - // set e in CPU value to e_gpu - psi::memory::synchronize_memory_op()( - gpu_ctx, - cpu_ctx, - e_gpu, - e, - nbands - ); - - if (all_eigenvalues) - { - //=========================== - // calculate all eigenvalues - //=========================== - dngv_op()(ctx, nstart, ldh, hcc, scc, e_gpu, vcc); - } - else - { - //===================================== - // calculate only m lowest eigenvalues - //===================================== - dngvx_op()(ctx, nstart, ldh, hcc, scc, nbands, e_gpu, vcc); + FPTYPE * res = e, *e_gpu = nullptr; +#if ((defined __CUDA) || (defined __ROCM)) + if (psi::device::get_device_type(ctx) == psi::GpuDevice) { + psi::memory::resize_memory_op()(gpu_ctx, e_gpu, nbands); + // set e in CPU value to e_gpu + syncmem_var_h2d_op()(gpu_ctx, cpu_ctx, e_gpu, e, nbands); + res = e_gpu; } +#endif - // set e_gpu value to e in CPU - psi::memory::synchronize_memory_op()( - cpu_ctx, - gpu_ctx, - e, - e_gpu, - nbands - ); - psi::memory::delete_memory_op()(gpu_ctx, e_gpu); -#else - - if (all_eigenvalues) - { + if (all_eigenvalues) { //=========================== // calculate all eigenvalues //=========================== - dngv_op()(ctx, nstart, ldh, hcc, scc, e, vcc); + dngv_op()(ctx, nstart, ldh, hcc, scc, res, vcc); } - else - { + else { //===================================== // calculate only m lowest eigenvalues //===================================== - dngvx_op()(ctx, nstart, ldh, hcc, scc, nbands, e, vcc); + dngvx_op()(ctx, nstart, ldh, hcc, scc, nbands, res, vcc); } +#if ((defined __CUDA) || (defined __ROCM)) + if (psi::device::get_device_type(ctx) == psi::GpuDevice) { + // set e_gpu value to e in CPU + syncmem_var_d2h_op()(cpu_ctx, gpu_ctx, e, res, nbands); + delmem_var_op()(gpu_ctx, e_gpu); + } #endif ModuleBase::timer::tick("DiagoIterAssist", "LAPACK_subspace"); - return; } template diff --git a/source/module_hsolver/diago_iter_assist.h b/source/module_hsolver/diago_iter_assist.h index 09692974f03..89ec8854c69 100644 --- a/source/module_hsolver/diago_iter_assist.h +++ b/source/module_hsolver/diago_iter_assist.h @@ -26,22 +26,15 @@ class DiagoIterAssist psi::Psi, Device> &evc, FPTYPE *en, int n_band = 0); - // for initializing wave function , this is a template function + static void diagH_subspace_init( hamilt::Hamilt* pHamilt, - const ModuleBase::ComplexMatrix &psi, + const std::complex* psi, + int psi_nr, + int psi_nc, psi::Psi, Device> &evc, FPTYPE *en); - static void diagH_LAPACK( - const int nstart, - const int nbands, - const ModuleBase::ComplexMatrix &hc, - const ModuleBase::ComplexMatrix &sc, - const int ldh, // nstart - FPTYPE *e, - ModuleBase::ComplexMatrix &hvec); - static void diagH_LAPACK( const int nstart, const int nbands, @@ -53,7 +46,23 @@ class DiagoIterAssist static bool test_exit_cond(const int &ntry, const int ¬conv); + private: + constexpr static const Device * ctx = {}; + constexpr static const psi::DEVICE_CPU * cpu_ctx = {}; + constexpr static const psi::DEVICE_GPU * gpu_ctx = {}; + using hpsi_info = typename hamilt::Operator, Device>::hpsi_info; + using resmem_var_op = psi::memory::resize_memory_op; + using delmem_var_op = psi::memory::delete_memory_op; + using syncmem_var_h2d_op = psi::memory::synchronize_memory_op; + using syncmem_var_d2h_op = psi::memory::synchronize_memory_op; + + using setmem_complex_op = psi::memory::set_memory_op, Device>; + using resmem_complex_op = psi::memory::resize_memory_op, Device>; + using delmem_complex_op = psi::memory::delete_memory_op, Device>; + using syncmem_complex_op = psi::memory::synchronize_memory_op, Device, Device>; + using syncmem_complex_h2d_op = psi::memory::synchronize_memory_op, Device, psi::DEVICE_CPU>; + using syncmem_complex_d2h_op = psi::memory::synchronize_memory_op, psi::DEVICE_CPU, Device>; }; } // namespace hsolver diff --git a/source/module_hsolver/hsolver.h b/source/module_hsolver/hsolver.h index 9ba048c7cae..2c1152641f1 100644 --- a/source/module_hsolver/hsolver.h +++ b/source/module_hsolver/hsolver.h @@ -12,6 +12,7 @@ namespace hsolver { +template class HSolver { public: @@ -32,8 +33,8 @@ class HSolver // solve Hamiltonian to electronic density in ElecState virtual void solve ( - hamilt::Hamilt* phm, - psi::Psi>& ppsi, + hamilt::Hamilt* phm, + psi::Psi, Device>& ppsi, elecstate::ElecState* pes, const std::string method, const bool skip_charge=false @@ -43,8 +44,8 @@ class HSolver } virtual void solve ( - hamilt::Hamilt* phm, - psi::Psi& ppsi, + hamilt::Hamilt* phm, + psi::Psi& ppsi, elecstate::ElecState* pes, const std::string method, const bool skip_charge=false @@ -55,8 +56,8 @@ class HSolver virtual void solve ( - hamilt::Hamilt* phm, - psi::Psi>& ppsi, + hamilt::Hamilt* phm, + psi::Psi, Device>& ppsi, elecstate::ElecState* pes, Stochastic_WF& stowf, const int istep, @@ -73,29 +74,28 @@ class HSolver // cg, dav, elpa, scalapack, hpseps, cusolver std::string method = "none"; public: - double diag_ethr=0.0; //threshold for diagonalization + FPTYPE diag_ethr=0.0; //threshold for diagonalization //set diag_ethr according to drho //for lcao, we suppose the error is zero and we set diag_ethr to 0 - virtual double set_diagethr(const int istep, const int iter, const double drho) + virtual FPTYPE set_diagethr(const int istep, const int iter, const FPTYPE drho) { return 0.0; } //reset diag_ethr according to drho and hsolver_error - virtual double reset_diagethr(std::ofstream& ofs_running, const double hsover_error, const double drho) + virtual FPTYPE reset_diagethr(std::ofstream& ofs_running, const FPTYPE hsover_error, const FPTYPE drho) { return 0.0; } // calculate hsolver_error // for sdft and lcao, we suppose the error is zero - virtual double cal_hsolerror() + virtual FPTYPE cal_hsolerror() { return 0.0; }; protected: - DiagH* pdiagh = nullptr; // for single Hamiltonian matrix diagonal solver - DiagH* gpu_diagh = nullptr; + DiagH* pdiagh = nullptr; // for single Hamiltonian matrix diagonal solver }; diff --git a/source/module_hsolver/hsolver_lcao.h b/source/module_hsolver/hsolver_lcao.h index e88fad62a1f..d7d27e95309 100644 --- a/source/module_hsolver/hsolver_lcao.h +++ b/source/module_hsolver/hsolver_lcao.h @@ -7,7 +7,7 @@ namespace hsolver { -class HSolverLCAO : public HSolver +class HSolverLCAO : public HSolver { public: HSolverLCAO(const Parallel_Orbitals* ParaV_in) diff --git a/source/module_hsolver/hsolver_pw.cpp b/source/module_hsolver/hsolver_pw.cpp index 14f2f7b5461..2e639f83f3c 100644 --- a/source/module_hsolver/hsolver_pw.cpp +++ b/source/module_hsolver/hsolver_pw.cpp @@ -10,9 +10,10 @@ #include "src_pw/global.h" #include -namespace hsolver -{ -HSolverPW::HSolverPW(ModulePW::PW_Basis_K* wfc_basis_in) +namespace hsolver { + +template +HSolverPW::HSolverPW(ModulePW::PW_Basis_K* wfc_basis_in) { this->wfc_basis = wfc_basis_in; this->classname = "HSolverPW"; @@ -29,46 +30,42 @@ void HSolverPW::update() { return; }*/ -void HSolverPW::initDiagh() +template +void HSolverPW::initDiagh() { if (this->method == "cg") { - if(pdiagh!=nullptr) + if(this->pdiagh!=nullptr) { - if(pdiagh->method != this->method) + if(this->pdiagh->method != this->method) { - delete[] pdiagh; - pdiagh = new DiagoCG(precondition.data()); - pdiagh->method = this->method; + delete (DiagoCG*)this->pdiagh; + this->pdiagh = new DiagoCG(precondition.data()); + this->pdiagh->method = this->method; } } else { - pdiagh = new DiagoCG(precondition.data()); - pdiagh->method = this->method; - // temperary added for debugging! - #if defined(__CUDA) || defined(__ROCM) - gpu_diagh = new DiagoCG(precondition.data()); - gpu_diagh->method = this->method; - #endif + this->pdiagh = new DiagoCG(precondition.data()); + this->pdiagh->method = this->method; } } else if (this->method == "dav") { DiagoDavid::PW_DIAG_NDIM = GlobalV::PW_DIAG_NDIM; - if (pdiagh != nullptr) + if (this->pdiagh != nullptr) { - if (pdiagh->method != this->method) + if (this->pdiagh->method != this->method) { - delete[] pdiagh; - pdiagh = new DiagoDavid( precondition.data()); - pdiagh->method = this->method; + delete (DiagoDavid*)this->pdiagh; + this->pdiagh = new DiagoDavid(precondition.data()); + this->pdiagh->method = this->method; } } else { - pdiagh = new DiagoDavid( precondition.data()); - pdiagh->method = this->method; + this->pdiagh = new DiagoDavid( precondition.data()); + this->pdiagh->method = this->method; } } else @@ -77,7 +74,8 @@ void HSolverPW::initDiagh() } } -void HSolverPW::solve(hamilt::Hamilt* pHamilt, psi::Psi>& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) +template +void HSolverPW::solve(hamilt::Hamilt* pHamilt, psi::Psi, Device>& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) { ModuleBase::TITLE("HSolverPW", "solve"); ModuleBase::timer::tick("HSolverPW", "solve"); @@ -112,29 +110,26 @@ void HSolverPW::solve(hamilt::Hamilt* pHamilt, psi::PsipsiToRho(psi); + reinterpret_cast*>(pes)->psiToRho(psi); ModuleBase::timer::tick("HSolverPW", "solve"); return; } -void HSolverPW::endDiagh() +template +void HSolverPW::endDiagh() { // DiagoCG would keep 9*nbasis memory in cache during loop-k // it should be deleted before calculating charge if(this->method == "cg") { - delete (DiagoCG*)pdiagh; - pdiagh = nullptr; - #if defined(__CUDA) || defined(__ROCM) - delete (DiagoCG*)gpu_diagh; - gpu_diagh = nullptr; - #endif + delete (DiagoCG*)this->pdiagh; + this->pdiagh = nullptr; } if(this->method == "dav") { - delete (DiagoDavid*)pdiagh; - pdiagh = nullptr; + delete (DiagoDavid*)this->pdiagh; + this->pdiagh = nullptr; } //in PW base, average iteration steps for each band and k-point should be printing @@ -150,7 +145,8 @@ void HSolverPW::endDiagh() } } -void HSolverPW::updatePsiK(hamilt::Hamilt* pHamilt, psi::Psi>& psi, const int ik) +template +void HSolverPW::updatePsiK(hamilt::Hamilt* pHamilt, psi::Psi, Device>& psi, const int ik) { psi.fix_k(ik); if(!this->initialed_psi) @@ -159,62 +155,28 @@ void HSolverPW::updatePsiK(hamilt::Hamilt* pHamilt, psi::Psictx, ik, psi, pHamilt); } else { ModuleBase::WARNING_QUIT("HSolverPW::updatePsiK", "lcao_in_pw is not supported now."); } - return; } } -void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi>& psi, double* eigenvalue) +template +void HSolverPW::hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi, Device>& psi, FPTYPE* eigenvalue) { - /// a huge victory here! - /// psi::Psi, psi::DEVICE_GPU> gpu_psi = psi; - /// psi::Psi, psi::DEVICE_CPU> cpu_psi = gpu_psi; - /// pdiagh->diag(hm, psi, eigenvalue); - /// psi::memory::synchronize_memory_op, psi::DEVICE_CPU, 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()); - - - /// hamilt::Hamilt* h_phm_in = - /// new hamilt::HamiltPW( - /// reinterpret_cast*>(d_phm_in)); - /// - /// pdiagh->diag(h_phm_in, psi, eigenvalue); - /// new era - /// if this works, then everything done! - -#if defined(__CUDA) || defined(__ROCM) - psi::Psi, psi::DEVICE_GPU> gpu_psi = psi; - hamilt::Hamilt* d_phm_in = - new hamilt::HamiltPW( - reinterpret_cast*>(hm)); - gpu_diagh->diag(d_phm_in, gpu_psi, eigenvalue); - psi::memory::synchronize_memory_op, psi::DEVICE_CPU, psi::DEVICE_GPU>()( - psi.get_device(), - gpu_psi.get_device(), - psi.get_pointer() - psi.get_psi_bias(), - gpu_psi.get_pointer() - gpu_psi.get_psi_bias(), - psi.size()); - delete reinterpret_cast*>(d_phm_in); -#else - pdiagh->diag(hm, psi, eigenvalue); -#endif + this->pdiagh->diag(hm, psi, eigenvalue); } -void HSolverPW::update_precondition(std::vector &h_diag, const int ik, const int npw) +template +void HSolverPW::update_precondition(std::vector &h_diag, const int ik, const int npw) { h_diag.assign(h_diag.size(), 1.0); int precondition_type = 2; - const double tpiba2 = this->wfc_basis->tpiba2; - + const FPTYPE tpiba2 = this->wfc_basis->tpiba2; + //=========================================== // Conjugate-Gradient diagonalization // h_diag is the precondition matrix @@ -224,7 +186,7 @@ void HSolverPW::update_precondition(std::vector &h_diag, const int ik, c { for (int ig = 0; ig < npw; ig++) { - double g2kin = this->wfc_basis->getgk2(ik,ig) * tpiba2; + FPTYPE g2kin = this->wfc_basis->getgk2(ik,ig) * tpiba2; h_diag[ig] = std::max(1.0, g2kin); } } @@ -232,7 +194,7 @@ void HSolverPW::update_precondition(std::vector &h_diag, const int ik, c { for (int ig = 0; ig < npw; ig++) { - double g2kin = this->wfc_basis->getgk2(ik,ig) * tpiba2; + FPTYPE g2kin = this->wfc_basis->getgk2(ik,ig) * tpiba2; h_diag[ig] = 1 + g2kin + sqrt(1 + (g2kin - 1) * (g2kin - 1)); } } @@ -246,12 +208,14 @@ void HSolverPW::update_precondition(std::vector &h_diag, const int ik, c } } -double HSolverPW::cal_hsolerror() +template +FPTYPE HSolverPW::cal_hsolerror() { return this->diag_ethr * std::max(1.0, GlobalV::nelec); } -double HSolverPW::set_diagethr(const int istep, const int iter, const double drho) +template +FPTYPE HSolverPW::set_diagethr(const int istep, const int iter, const FPTYPE drho) { //It is too complex now and should be modified. if (iter == 1) @@ -293,7 +257,8 @@ double HSolverPW::set_diagethr(const int istep, const int iter, const double drh return this->diag_ethr; } -double HSolverPW::reset_diagethr(std::ofstream& ofs_running, const double hsover_error, const double drho) +template +FPTYPE HSolverPW::reset_diagethr(std::ofstream& ofs_running, const FPTYPE hsover_error, const FPTYPE drho) { ofs_running << " Notice: Threshold on eigenvalues was too large.\n"; ModuleBase::WARNING("scf", "Threshold on eigenvalues was too large."); @@ -304,5 +269,9 @@ double HSolverPW::reset_diagethr(std::ofstream& ofs_running, const double hsover return this->diag_ethr; } +template class HSolverPW; +#if ((defined __CUDA) || (defined __ROCM)) +template class HSolverPW; +#endif } // namespace hsolver \ No newline at end of file diff --git a/source/module_hsolver/hsolver_pw.h b/source/module_hsolver/hsolver_pw.h index 61e250518c7..4dde0c93cf9 100644 --- a/source/module_hsolver/hsolver_pw.h +++ b/source/module_hsolver/hsolver_pw.h @@ -4,10 +4,10 @@ #include "hsolver.h" #include "module_pw/pw_basis_k.h" -namespace hsolver -{ +namespace hsolver { -class HSolverPW : public HSolver +template +class HSolverPW: public HSolver { public: HSolverPW(ModulePW::PW_Basis_K* wfc_basis_in); @@ -19,26 +19,28 @@ class HSolverPW : public HSolver void update(//Input &in ) override;*/ - void solve(hamilt::Hamilt* pHamilt, psi::Psi>& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) override; + void solve(hamilt::Hamilt* pHamilt, psi::Psi, Device>& psi, elecstate::ElecState* pes, const std::string method_in, const bool skip_charge) override; - virtual double cal_hsolerror() override; - virtual double set_diagethr(const int istep, const int iter, const double drho) override; - virtual double reset_diagethr(std::ofstream& ofs_running, const double hsover_error, const double drho) override; + virtual FPTYPE cal_hsolerror() override; + virtual FPTYPE set_diagethr(const int istep, const int iter, const FPTYPE drho) override; + virtual FPTYPE reset_diagethr(std::ofstream& ofs_running, const FPTYPE hsover_error, const FPTYPE drho) override; protected: void initDiagh(); void endDiagh(); - void hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi>& psi, double* eigenvalue); + void hamiltSolvePsiK(hamilt::Hamilt* hm, psi::Psi, Device>& psi, FPTYPE* eigenvalue); - void updatePsiK(hamilt::Hamilt* pHamilt, psi::Psi>& psi, const int ik); + void updatePsiK(hamilt::Hamilt* pHamilt, psi::Psi, Device>& psi, const int ik); ModulePW::PW_Basis_K* wfc_basis = nullptr; // calculate the precondition array for diagonalization in PW base - void update_precondition(std::vector &h_diag, const int ik, const int npw); + void update_precondition(std::vector &h_diag, const int ik, const int npw); - std::vector precondition; + std::vector precondition; bool initialed_psi = false; + + Device *ctx = {}; }; } // namespace hsolver diff --git a/source/module_hsolver/hsolver_pw_sdft.h b/source/module_hsolver/hsolver_pw_sdft.h index 333d51e47f1..371124aa49d 100644 --- a/source/module_hsolver/hsolver_pw_sdft.h +++ b/source/module_hsolver/hsolver_pw_sdft.h @@ -4,7 +4,7 @@ #include "src_pw/sto_iter.h" namespace hsolver { - class HSolverPW_SDFT : public HSolverPW + class HSolverPW_SDFT : public HSolverPW { public: HSolverPW_SDFT(ModulePW::PW_Basis_K* wfc_basis_in, Stochastic_WF& stowf, const int method_sto):HSolverPW(wfc_basis_in) diff --git a/source/module_hsolver/src/cuda/dngvd_op.cu b/source/module_hsolver/src/cuda/dngvd_op.cu index e67cb4cef44..28854af019b 100644 --- a/source/module_hsolver/src/cuda/dngvd_op.cu +++ b/source/module_hsolver/src/cuda/dngvd_op.cu @@ -155,8 +155,6 @@ void dngv_op::operator()(const psi::DEVICE_GPU* d, double* W, std::complex* V) { - createBLAShandle(); - // init A_eigenvectors & transpose_B double2 *A_eigenvectors, *transpose_B; checkCudaErrors(cudaMalloc((void**)&A_eigenvectors, sizeof(double2) * row * col)); @@ -228,8 +226,6 @@ void dngv_op::operator()(const psi::DEVICE_GPU* d, checkCudaErrors(cudaFree(all_W)); checkCudaErrors(cudaFree(devInfo)); cusolverErrcheck(cusolverDnDestroy(cusolverH)); - - destoryBLAShandle(); } } // namespace hsolver \ No newline at end of file diff --git a/source/module_hsolver/src/cuda/math_kernel.cu b/source/module_hsolver/src/cuda/math_kernel.cu index bb39ac62428..cb68dcab85d 100644 --- a/source/module_hsolver/src/cuda/math_kernel.cu +++ b/source/module_hsolver/src/cuda/math_kernel.cu @@ -9,19 +9,19 @@ namespace hsolver { -static cublasHandle_t diag_handle = nullptr; +static cublasHandle_t cublas_handle = nullptr; void createBLAShandle(){ - if (diag_handle == nullptr) { - cublasErrcheck(cublasCreate(&diag_handle)); + if (cublas_handle == nullptr) { + cublasErrcheck(cublasCreate(&cublas_handle)); } } void destoryBLAShandle(){ - if (diag_handle != nullptr) { - cublasErrcheck(cublasDestroy(diag_handle)); - diag_handle = nullptr; + if (cublas_handle != nullptr) { + cublasErrcheck(cublasDestroy(cublas_handle)); + cublas_handle = nullptr; } } @@ -101,7 +101,8 @@ FPTYPE zdot_real_op::operator()( // Note that ddot_(2*dim,a,1,b,1) = REAL( zdotc_(dim,a,1,b,1) ) const FPTYPE* pL = reinterpret_cast(psi_L); const FPTYPE* pR = reinterpret_cast(psi_R); - FPTYPE result = thrust::inner_product(thrust::device, pL, pL + dim * 2, pR, FPTYPE(0.0)); + FPTYPE result = 0.0; + cublasErrcheck(cublasDdot(cublas_handle, dim * 2, pL, 1, pR, 1, &result)); if (reduce) { Parallel_Reduce::reduce_double_pool(result); } @@ -190,7 +191,7 @@ void axpy_op::operator()( std::complex *Y, const int& incY) { - cublasErrcheck(cublasCaxpy(diag_handle, N, (float2*)alpha, (float2*)X, incX, (float2*)Y, incY)); + cublasErrcheck(cublasCaxpy(cublas_handle, N, (float2*)alpha, (float2*)X, incX, (float2*)Y, incY)); } template <> @@ -203,7 +204,7 @@ void axpy_op::operator()( std::complex *Y, const int& incY) { - cublasErrcheck(cublasZaxpy(diag_handle, N, (double2*)alpha, (double2*)X, incX, (double2*)Y, incY)); + cublasErrcheck(cublasZaxpy(cublas_handle, N, (double2*)alpha, (double2*)X, incX, (double2*)Y, incY)); } template <> @@ -231,7 +232,7 @@ void gemv_op::operator()( else if (trans == 'C'){ cutrans = CUBLAS_OP_C; } - cublasErrcheck(cublasCgemv(diag_handle, cutrans, m, n, (float2*)alpha, (float2*)A, lda, (float2*)X, incx, (float2*)beta, (float2*)Y, incx)); + cublasErrcheck(cublasCgemv(cublas_handle, cutrans, m, n, (float2*)alpha, (float2*)A, lda, (float2*)X, incx, (float2*)beta, (float2*)Y, incx)); } template <> @@ -259,7 +260,7 @@ void gemv_op::operator()( else if (trans == 'C'){ cutrans = CUBLAS_OP_C; } - cublasErrcheck(cublasZgemv(diag_handle, cutrans, m, n, (double2*)alpha, (double2*)A, lda, (double2*)X, incx, (double2*)beta, (double2*)Y, incx)); + cublasErrcheck(cublasZgemv(cublas_handle, cutrans, m, n, (double2*)alpha, (double2*)A, lda, (double2*)X, incx, (double2*)beta, (double2*)Y, incx)); } @@ -271,7 +272,7 @@ void scal_op::operator()(const psi::DEVICE_GPU* d, std::complex* X, const int& incx) { - cublasErrcheck(cublasZscal(diag_handle, N, (double2*)alpha, (double2*)X, incx)); + cublasErrcheck(cublasZscal(cublas_handle, N, (double2*)alpha, (double2*)X, incx)); } template <> @@ -282,7 +283,7 @@ void scal_op::operator()(const psi::DEVICE_GPU* d, const int& incx) { - cublasErrcheck(cublasCscal(diag_handle, N, (float2*)alpha, (float2*)X, incx)); + cublasErrcheck(cublasCscal(cublas_handle, N, (float2*)alpha, (float2*)X, incx)); } @@ -324,7 +325,7 @@ void gemm_op::operator()(const psi::DEVICE_GPU* d, else if (transb == 'C'){ cutransB = CUBLAS_OP_C; } - cublasErrcheck(cublasCgemm(diag_handle, cutransA, cutransB, m, n ,k, (float2*)alpha, (float2*)a , lda, (float2*)b, ldb, (float2*)beta, (float2*)c, ldc)); + cublasErrcheck(cublasCgemm(cublas_handle, cutransA, cutransB, m, n ,k, (float2*)alpha, (float2*)a , lda, (float2*)b, ldb, (float2*)beta, (float2*)c, ldc)); } template <> @@ -365,7 +366,7 @@ void gemm_op::operator()(const psi::DEVICE_GPU* d, else if (transb == 'C'){ cutransB = CUBLAS_OP_C; } - cublasErrcheck(cublasZgemm(diag_handle, cutransA, cutransB, m, n ,k, (double2*)alpha, (double2*)a , lda, (double2*)b, ldb, (double2*)beta, (double2*)c, ldc)); + cublasErrcheck(cublasZgemm(cublas_handle, cutransA, cutransB, m, n ,k, (double2*)alpha, (double2*)a , lda, (double2*)b, ldb, (double2*)beta, (double2*)c, ldc)); } @@ -404,7 +405,7 @@ void matrixTranspose_op::operator()(const psi::DEVICE_G ZERO.x = ZERO.y = 0.0; // use 'geam' API todo transpose. - cublasErrcheck(cublasZgeam(diag_handle, CUBLAS_OP_T, CUBLAS_OP_N, col, row, &ONE, (double2*)input_matrix, col, &ZERO, (double2*)input_matrix, col, (double2*)device_temp, col)); + cublasErrcheck(cublasZgeam(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, col, row, &ONE, (double2*)input_matrix, col, &ZERO, (double2*)input_matrix, col, (double2*)device_temp, col)); } else { int thread = 1024; diff --git a/source/module_hsolver/src/dngvd_op.cpp b/source/module_hsolver/src/dngvd_op.cpp index 62ef44c326b..41e5372c2bc 100644 --- a/source/module_hsolver/src/dngvd_op.cpp +++ b/source/module_hsolver/src/dngvd_op.cpp @@ -7,108 +7,181 @@ namespace hsolver template <> void dngvx_op::operator()(const psi::DEVICE_CPU* d, - const int row, - const int col, - const std::complex* A, - const std::complex* B, - const int m, - double* W, - std::complex* V) + const int nstart, // nstart + const int ldh, // ldh + const std::complex* hcc, // hcc + const std::complex* scc, // scc + const int nbands, // nbands + double* eigenvalue, // eigenvalue + std::complex* vcc) // vcc { - int lwork; - int info = 0; + int lwork = 0; + psi::DEVICE_CPU * cpu_ctx = {}; + + ModuleBase::ComplexMatrix sdum(nstart, ldh); + ModuleBase::ComplexMatrix hdum; + + ModuleBase::ComplexMatrix hc(nstart, nstart); + ModuleBase::ComplexMatrix hvec(nstart, nbands); + psi::memory::synchronize_memory_op, psi::DEVICE_CPU, psi::DEVICE_CPU>()( + cpu_ctx, + cpu_ctx, + hc.c, + hcc, + nstart * nstart + ); + + ModuleBase::ComplexMatrix sc(nstart, nstart); + psi::memory::synchronize_memory_op, psi::DEVICE_CPU, psi::DEVICE_CPU>()( + cpu_ctx, + cpu_ctx, + sc.c, + scc, + nstart * nstart + ); + sdum = sc; - std::string name1 = "ZHETRD"; - std::string name2 = "L"; + // workspace query + int nb = LapackConnector::ilaenv(1, "ZHETRD", "U", nstart, -1, -1, -1); - int nb = LapackConnector::ilaenv(1, name1.c_str(), name2.c_str(), col, -1, -1, -1); if (nb < 1) { - nb = std::max(1, col); + nb = std::max(1, nstart); } - if (nb == 1 || nb >= col) + if (nb == 1 || nb >= nstart) { - lwork = 2 * col; // qianrui fix a bug 2021-7-25 : lwork should be at least max(1,2*n) + lwork = 2 * nstart; // mohan modify 2009-08-02 } else { - lwork = (nb + 1) * col; + lwork = (nb + 1) * nstart; } - std::complex* work = new std::complex[2 * lwork]; - assert(work != 0); - double* rwork = new double[7 * col]; - assert(rwork != 0); - int* iwork = new int[5 * col]; - assert(iwork != 0); - int* ifail = new int[col]; - assert(ifail != 0); - ModuleBase::GlobalFunc::ZEROS(work, lwork); // qianrui change it, only first lwork numbers are used in zhegvx - ModuleBase::GlobalFunc::ZEROS(rwork, 7 * col); - ModuleBase::GlobalFunc::ZEROS(iwork, 5 * col); - ModuleBase::GlobalFunc::ZEROS(ifail, col); - - LapackConnector::zhegvx( - 1, // ITYPE = 1: A*x = (lambda)*B*x - 'V', // JOBZ = 'V': Compute eigenvalues and eigenvectors. - 'I', // RANGE = 'I': the IL-th through IU-th eigenvalues will be found. - 'L', // UPLO = 'L': Lower triangles of A and B are stored. - col, // N = base - A, // A is COMPLEX*16 array dimension (LDA, N) - col, // LDA = base - B, // B is COMPLEX*16 array, dimension (LDB, N) - col, // LDB = base - 0.0, // Not referenced if RANGE = 'A' or 'I'. - 0.0, // Not referenced if RANGE = 'A' or 'I'. - 1, // IL: If RANGE='I', the index of the smallest eigenvalue to be returned. 1 <= IL <= IU <= N, - m, // IU: If RANGE='I', the index of the largest eigenvalue to be returned. 1 <= IL <= IU <= N, - 0.0, // ABSTOL - m, // M: The total number of eigenvalues found. 0 <= M <= N. if RANGE = 'I', M = IU-IL+1. - W, // W store eigenvalues - V, // store eigenvector - col, // LDZ: The leading dimension of the array Z. - work, - lwork, - rwork, - iwork, - ifail, - info, - row); - - assert(0 == info); + std::complex *work = new std::complex[lwork]; + ModuleBase::GlobalFunc::ZEROS(work, lwork); - delete[] work; - delete[] rwork; + //===================================================================== + // input s and (see below) h are copied so that they are not destroyed + //===================================================================== + + int info = 0; + int rwork_dim = 7 * nstart; + + double *rwork = new double[rwork_dim]; + ModuleBase::GlobalFunc::ZEROS(rwork, rwork_dim); + + //===================================== + // calculate only m lowest eigenvalues + //===================================== + int *iwork = new int[5 * nstart]; + int *ifail = new int[nstart]; + + ModuleBase::GlobalFunc::ZEROS(rwork, 7 * nstart); + ModuleBase::GlobalFunc::ZEROS(iwork, 5 * nstart); + ModuleBase::GlobalFunc::ZEROS(ifail, nstart); + + hdum.create(nstart, ldh); + hdum = hc; + + //============================= + // Number of calculated bands + //============================= + int mm = nbands; + + LapackConnector::zhegvx(1, // INTEGER + 'V', // CHARACTER*1 + 'I', // CHARACTER*1 + 'U', // CHARACTER*1 + nstart, // INTEGER + hdum, // COMPLEX*16 array + ldh, // INTEGER + sdum, // COMPLEX*16 array + ldh, // INTEGER + 0.0, // DOUBLE PRECISION + 0.0, // DOUBLE PRECISION + 1, // INTEGER + nbands, // INTEGER + 0.0, // DOUBLE PRECISION + mm, // INTEGER + eigenvalue, // DOUBLE PRECISION array + hvec, // COMPLEX*16 array + ldh, // INTEGER + work, // DOUBLE array, dimension (MAX(1,LWORK)) + lwork, // INTEGER + rwork, // DOUBLE PRECISION array, dimension (7*N) + iwork, // INTEGER array, dimension (5*N) + ifail, // INTEGER array, dimension (N) + info // INTEGER + ); + // dngvx_op()(ctx, nstart, ldh, hc.c, sc.c, nbands, e, hvec.c); + + psi::memory::synchronize_memory_op, psi::DEVICE_CPU, psi::DEVICE_CPU>()( + cpu_ctx, + cpu_ctx, + vcc, + hvec.c, + nstart * nbands + ); delete[] iwork; delete[] ifail; + delete[] rwork; + delete[] work; }; template <> void dngv_op::operator()(const psi::DEVICE_CPU* d, - const int row, - const int col, - const std::complex* A, - const std::complex* B, - double* W, - std::complex* V) + const int nstart, + const int ldh, + const std::complex* hcc, + const std::complex* scc, + double* eigenvalue, + std::complex* vcc) { int lwork = 0; - int nb = LapackConnector::ilaenv(1, "ZHETRD", "U", col, -1, -1, -1); + psi::DEVICE_CPU * cpu_ctx = {}; + + ModuleBase::ComplexMatrix sdum(nstart, ldh); + ModuleBase::ComplexMatrix hdum; + + ModuleBase::ComplexMatrix hc(nstart, nstart); + ModuleBase::ComplexMatrix hvec(nstart, nstart); + psi::memory::synchronize_memory_op, psi::DEVICE_CPU, psi::DEVICE_CPU>()( + cpu_ctx, + cpu_ctx, + hc.c, + hcc, + nstart * nstart + ); + + ModuleBase::ComplexMatrix sc(nstart, nstart); + psi::memory::synchronize_memory_op, psi::DEVICE_CPU, psi::DEVICE_CPU>()( + cpu_ctx, + cpu_ctx, + sc.c, + scc, + nstart * nstart + ); + sdum = sc; + + // workspace query + int nb = LapackConnector::ilaenv(1, "ZHETRD", "U", nstart, -1, -1, -1); + if (nb < 1) { - nb = std::max(1, col); + nb = std::max(1, nstart); } - if (nb == 1 || nb >= col) + + if (nb == 1 || nb >= nstart) { - lwork = 2 * col; // mohan modify 2009-08-02 + lwork = 2 * nstart; // mohan modify 2009-08-02 } else { - lwork = (nb + 1) * col; + lwork = (nb + 1) * nstart; } - std::complex* work = new std::complex[lwork]; + std::complex *work = new std::complex[lwork]; ModuleBase::GlobalFunc::ZEROS(work, lwork); //===================================================================== @@ -116,19 +189,24 @@ void dngv_op::operator()(const psi::DEVICE_CPU* d, //===================================================================== int info = 0; - int rwork_dim; - rwork_dim = 3 * col - 2; - double* rwork = new double[rwork_dim]; - ModuleBase::GlobalFunc::ZEROS(rwork, rwork_dim); + int rwork_dim = 3 * nstart - 2; - for (int i = 0; i < row * col; i++) - { - V[i] = A[i]; - } - LapackConnector::zhegv(1, 'V', 'U', col, V, col, B, col, W, work, lwork, rwork, info, row); + double *rwork = new double[rwork_dim]; + ModuleBase::GlobalFunc::ZEROS(rwork, rwork_dim); - assert(0 == info); + //=========================== + // calculate all eigenvalues + //=========================== + hvec = hc; + LapackConnector::zhegv(1, 'V', 'U', nstart, hvec, ldh, sdum, ldh, eigenvalue, work, lwork, rwork, info); + psi::memory::synchronize_memory_op, psi::DEVICE_CPU, psi::DEVICE_CPU>()( + cpu_ctx, + cpu_ctx, + vcc, + hvec.c, + nstart * nstart + ); delete[] rwork; delete[] work; } diff --git a/source/module_psi/include/device.h b/source/module_psi/include/device.h index 50ab187ff0f..c84231bf112 100644 --- a/source/module_psi/include/device.h +++ b/source/module_psi/include/device.h @@ -15,6 +15,19 @@ template void print_device_info (const Device* dev, std::ofstre template void record_device_memory (const Device* dev, std::ofstream& ofs_device, std::string str, size_t size) {return;} +int get_device_kpar(const int& kpar); +std::string get_device_flag(const std::string& device, const std::string& ks_solver, const std::string& basis_type); + +#if __MPI +int get_node_rank(); +int stringCmp(const void *a, const void* b); +#endif + +#if ((defined __CUDA) || (defined __ROCM)) +int get_device_num(); +void set_device(const int rank); +#endif + } // end of namespace device } // end of namespace psi diff --git a/source/module_psi/src/device.cpp b/source/module_psi/src/device.cpp index c00bc0cd59d..a8ea9a59243 100644 --- a/source/module_psi/src/device.cpp +++ b/source/module_psi/src/device.cpp @@ -1,9 +1,6 @@ -#include -#include -#include #include -#include "module_psi/psi.h" +#include #include "module_psi/include/types.h" #include "module_psi/include/device.h" #include "module_base/tool_quit.h" @@ -12,6 +9,10 @@ #include #endif +#ifdef __MPI +#include "mpi.h" +#endif + namespace psi{ namespace device{ @@ -27,6 +28,16 @@ template<> AbacusDevice_t get_device_type (const DEVICE_CPU* dev) { template<> AbacusDevice_t get_device_type (const DEVICE_GPU* dev) { return GpuDevice; } + +void set_device(const int rank) { + cudaSetDevice(rank); +} + +int get_device_num() { + int device_num = -1; + cudaGetDeviceCount(&device_num); + return device_num; +} #endif #if defined(__CUDA) || defined(__ROCM) @@ -54,7 +65,7 @@ template<> void print_device_info (const DEVICE_GPU* ctx, std::ofst // Console log cudaDriverGetVersion(&driverVersion); cudaRuntimeGetVersion(&runtimeVersion); - char msg[256]; + char msg[1024]; sprintf(msg, " CUDA Driver Version / Runtime Version %d.%d / %d.%d\n", driverVersion / 1000, (driverVersion % 100) / 10, @@ -257,5 +268,116 @@ template<> void record_device_memory (const DEVICE_GPU* ctx, std::of #endif +#if __MPI +int stringCmp(const void *a, const void* b) +{ + char* m = (char*)a; + char* n = (char*)b; + int i, sum = 0; + + for(i = 0; i < MPI_MAX_PROCESSOR_NAME; i++) + if (m[i] == n[i]) + continue; + else + { + sum = m[i] - n[i]; + break; + } + return sum; +} + +int get_node_rank() { + char host_name[MPI_MAX_PROCESSOR_NAME]; + memset(host_name, '\0', sizeof(char) * MPI_MAX_PROCESSOR_NAME); + char (*host_names)[MPI_MAX_PROCESSOR_NAME]; + int n, namelen, color, rank, nprocs, myrank; + size_t bytes; + MPI_Comm nodeComm; + + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Comm_size(MPI_COMM_WORLD, &nprocs); + MPI_Get_processor_name(host_name,&namelen); + + bytes = nprocs * sizeof(char[MPI_MAX_PROCESSOR_NAME]); + host_names = (char (*)[MPI_MAX_PROCESSOR_NAME]) malloc(bytes); + for (int ii = 0; ii < nprocs; ii++) { + memset(host_names[ii], '\0', sizeof(char) * MPI_MAX_PROCESSOR_NAME); + } + + strcpy(host_names[rank], host_name); + + for (n=0; n> { if(phm_in!= nullptr) { + // hsolver::DiagoIterAssist::diagH_subspace_init(phm_in, + // wfcatom, + // wvf, + // etatom.data()); hsolver::DiagoIterAssist::diagH_subspace_init(phm_in, - wfcatom, - wvf, - etatom.data()); + wfcatom.c, + wfcatom.nr, + wfcatom.nc, + wvf, + etatom.data()); return; } else @@ -263,6 +270,90 @@ void wavefunc::diago_PAO_in_pw_k2(const int &ik, psi::Psi> } } +void wavefunc::diago_PAO_in_pw_k2_device(const psi::DEVICE_CPU* ctx, const int &ik, psi::Psi, psi::DEVICE_CPU> &wvf, hamilt::Hamilt* phm_in) +{ + this->diago_PAO_in_pw_k2(ik, wvf, phm_in); +} + +#if ((defined __CUDA) || (defined __ROCM)) +void wavefunc::diago_PAO_in_pw_k2_device(const psi::DEVICE_GPU* ctx, const int &ik, psi::Psi, psi::DEVICE_GPU> &wvf, hamilt::Hamilt* phm_in) +{ + ModuleBase::TITLE("wavefunc","diago_PAO_in_pw_k2"); + // (6) Prepare for atmoic orbitals or random orbitals + const int starting_nw = this->get_starting_nw(); + if(starting_nw == 0) return; + assert(starting_nw > 0); + + const int nbasis = wvf.get_nbasis(); + const int nbands = wvf.get_nbands(); + const int current_nbasis = GlobalC::kv.ngk[ik]; + + ModuleBase::ComplexMatrix wfcatom(starting_nw, nbasis);//added by zhengdy-soc + if(GlobalV::test_wf)ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running, "starting_nw", starting_nw); + if(init_wfc.substr(0,6)=="atomic") + { + this->atomic_wfc(ik, current_nbasis, GlobalC::ucell.lmax_ppwf, wfcatom, GlobalC::ppcell.tab_at, GlobalV::NQX, GlobalV::DQ); + if( init_wfc == "atomic+random" && starting_nw == GlobalC::ucell.natomwfc )//added by qianrui 2021-5-16 + { + this->atomicrandom(wfcatom,0,starting_nw,ik, GlobalC::wfcpw); + } + + //==================================================== + // If not enough atomic wfc are available, complete + // with random wfcs + //==================================================== + this->random(wfcatom, GlobalC::ucell.natomwfc, nbands, ik, GlobalC::wfcpw); + } + else if(init_wfc=="random") + { + this->random(wfcatom,0,nbands,ik, GlobalC::wfcpw); + } + + // (7) Diago with cg method. + std::vector etatom(starting_nw, 0.0); + //if(GlobalV::DIAGO_TYPE == "cg") xiaohui modify 2013-09-02 + if(GlobalV::KS_SOLVER=="cg") //xiaohui add 2013-09-02 + { + if(phm_in!= nullptr) + { + // hsolver::DiagoIterAssist::diagH_subspace_init(phm_in, + // wfcatom, + // wvf, + // etatom.data()); + std::complex *d_wfcatom = nullptr; + psi::DEVICE_CPU * cpu_ctx = {}; + psi::DEVICE_GPU * gpu_ctx = {}; + resmem_complex_op()(gpu_ctx, d_wfcatom, wfcatom.nr * wfcatom.nc); + syncmem_complex_h2d_op()(gpu_ctx, cpu_ctx, d_wfcatom, wfcatom.c, wfcatom.nr * wfcatom.nc); + hsolver::DiagoIterAssist::diagH_subspace_init( + phm_in, + d_wfcatom, + wfcatom.nr, + wfcatom.nc, + wvf, + etatom.data()); + delmem_complex_op()(gpu_ctx, d_wfcatom); + return; + } + else + { + //this diagonalization method is obsoleted now + //GlobalC::hm.diagH_subspace(ik ,starting_nw, nbands, wfcatom, wfcatom, etatom.data()); + } + } + + assert(nbands <= wfcatom.nr); + for (int ib=0; ib>* psi_in) { ModuleBase::TITLE("wavefunc","wfcinit_k"); @@ -271,9 +362,9 @@ void wavefunc::wfcinit_k(psi::Psi>* psi_in) { this->irindex = new int [GlobalC::wfcpw->fftnxy]; GlobalC::wfcpw->getfftixy2is(this->irindex); - #if defined(__CUDA) || defined(__UT_USE_CUDA) - GlobalC::wfcpw->get_ig2ixyz_k(); - #endif + #if defined(__CUDA) || defined(__UT_USE_CUDA) + GlobalC::wfcpw->get_ig2ixyz_k(); + #endif } if(GlobalV::CALCULATION=="nscf") { diff --git a/source/src_pw/wavefunc.h b/source/src_pw/wavefunc.h index 410512050de..19f59eff79f 100644 --- a/source/src_pw/wavefunc.h +++ b/source/src_pw/wavefunc.h @@ -41,7 +41,11 @@ class wavefunc : public WF_atomic void diago_PAO_in_pw_k2(const int &ik, psi::Psi> &wvf, hamilt::Hamilt* phm_in = nullptr); void diago_PAO_in_pw_k2(const int &ik, ModuleBase::ComplexMatrix &wvf); - + + void diago_PAO_in_pw_k2_device(const psi::DEVICE_CPU* ctx, const int &ik, psi::Psi, psi::DEVICE_CPU> &wvf, hamilt::Hamilt* phm_in = nullptr); +#if ((defined __CUDA) || (defined __ROCM)) + void diago_PAO_in_pw_k2_device(const psi::DEVICE_GPU* ctx, const int &ik, psi::Psi, psi::DEVICE_GPU> &wvf, hamilt::Hamilt* phm_in = nullptr); +#endif int get_R(int ix, int iy, int iz); // pengfei 2016-11-23 int iw2it(int iw); @@ -53,6 +57,10 @@ class wavefunc : public WF_atomic ModuleBase::Vector3 ***R; int ** Rmax; + + using resmem_complex_op = psi::memory::resize_memory_op, psi::DEVICE_GPU>; + using delmem_complex_op = psi::memory::delete_memory_op, psi::DEVICE_GPU>; + using syncmem_complex_h2d_op = psi::memory::synchronize_memory_op, psi::DEVICE_GPU, psi::DEVICE_CPU>; }; #endif //wavefunc