Make init stochastic WF support GPU

Author: Qianruipku

source/module_esolver/esolver_sdft_pw.cpp

@@ -54,13 +54,13 @@ void ESolver_SDFT_PW<T, Device>::before_all_runners(const Input_para& inp, UnitC
 
     // 3) initialize the pointer for electronic states of SDFT
     this->pelec = new elecstate::ElecStatePW_SDFT<T, Device>(this->pw_wfc,
-                                                          &(this->chr),
-                                                          &(this->kv),
-                                                          &ucell,
-                                                          &(GlobalC::ppcell),
-                                                          this->pw_rhod,
-                                                          this->pw_rho,
-                                                          this->pw_big);
+                                                             &(this->chr),
+                                                             &(this->kv),
+                                                             &ucell,
+                                                             &(GlobalC::ppcell),
+                                                             this->pw_rhod,
+                                                             this->pw_rho,
+                                                             this->pw_big);
 
     // 4) inititlize the charge density.
     this->pelec->charge->allocate(PARAM.inp.nspin);
@@ -80,11 +80,11 @@ void ESolver_SDFT_PW<T, Device>::before_all_runners(const Input_para& inp, UnitC
 
     // 6) prepare some parameters for electronic wave functions initilization
     this->p_wf_init = new psi::WFInit<T, Device>(PARAM.inp.init_wfc,
-                                                                    PARAM.inp.ks_solver,
-                                                                    PARAM.inp.basis_type,
-                                                                    PARAM.inp.psi_initializer,
-                                                                    &this->wf,
-                                                                    this->pw_wfc);
+                                                 PARAM.inp.ks_solver,
+                                                 PARAM.inp.basis_type,
+                                                 PARAM.inp.psi_initializer,
+                                                 &this->wf,
+                                                 this->pw_wfc);
     // 7) set occupatio, redundant?
     if (PARAM.inp.ocp)
     {
@@ -95,42 +95,38 @@ void ESolver_SDFT_PW<T, Device>::before_all_runners(const Input_para& inp, UnitC
     this->Init_GlobalC(inp, ucell, GlobalC::ppcell); // temporary
 
     // 9) initialize the stochastic wave functions
-    stowf.init(&this->kv, this->pw_wfc->npwk_max);
+    this->stowf.init(&this->kv, this->pw_wfc->npwk_max);
     if (inp.nbands_sto != 0)
     {
         if (inp.initsto_ecut < inp.ecutwfc)
         {
-            Init_Sto_Orbitals(this->stowf, inp.seed_sto);
+            this->stowf.init_sto_orbitals(inp.seed_sto);
         }
         else
         {
-            Init_Sto_Orbitals_Ecut(this->stowf, inp.seed_sto, this->kv, *this->pw_wfc, inp.initsto_ecut);
+            this->stowf.init_sto_orbitals_Ecut(inp.seed_sto, this->kv, *this->pw_wfc, inp.initsto_ecut);
         }
     }
     else
     {
-        Init_Com_Orbitals(this->stowf);
+        this->stowf.init_com_orbitals();
     }
     if (this->method_sto == 2)
     {
-        stowf.allocate_chiallorder(this->nche_sto);
+        this->stowf.allocate_chiallorder(this->nche_sto);
     }
+    this->stowf.sync_chi0();
 
+    // 10) allocate spaces for \sqrt(f(H))|chi> and |\tilde{chi}>
     size_t size = stowf.chi0->size();
-
-    this->stowf.shchi = new psi::Psi<T>(this->kv.get_nks(),
-                                                           this->stowf.nchip_max,
-                                                           this->wf.npwx,
-                                                           this->kv.ngk.data());
-
+    this->stowf.shchi
+        = new psi::Psi<T, Device>(this->kv.get_nks(), this->stowf.nchip_max, this->wf.npwx, this->kv.ngk.data());
     ModuleBase::Memory::record("SDFT::shchi", size * sizeof(T));
 
     if (PARAM.inp.nbands > 0)
     {
-        this->stowf.chiortho = new psi::Psi<T>(this->kv.get_nks(),
-                                                                  this->stowf.nchip_max,
-                                                                  this->wf.npwx,
-                                                                  this->kv.ngk.data());
+        this->stowf.chiortho
+            = new psi::Psi<T, Device>(this->kv.get_nks(), this->stowf.nchip_max, this->wf.npwx, this->kv.ngk.data());
         ModuleBase::Memory::record("SDFT::chiortho", size * sizeof(T));
     }
 
@@ -143,16 +139,16 @@ void ESolver_SDFT_PW<T, Device>::before_scf(const int istep)
     ESolver_KS_PW<T, Device>::before_scf(istep);
     delete reinterpret_cast<hamilt::HamiltPW<double>*>(this->p_hamilt);
     this->p_hamilt = new hamilt::HamiltSdftPW<T, Device>(this->pelec->pot,
-                                                                            this->pw_wfc,
-                                                                            &this->kv,
-                                                                            PARAM.globalv.npol,
-                                                                            &this->stoche.emin_sto,
-                                                                            &this->stoche.emax_sto);
+                                                         this->pw_wfc,
+                                                         &this->kv,
+                                                         PARAM.globalv.npol,
+                                                         &this->stoche.emin_sto,
+                                                         &this->stoche.emax_sto);
     this->p_hamilt_sto = static_cast<hamilt::HamiltSdftPW<T, Device>*>(this->p_hamilt);
 
     if (istep > 0 && PARAM.inp.nbands_sto != 0 && PARAM.inp.initsto_freq > 0 && istep % PARAM.inp.initsto_freq == 0)
     {
-        Update_Sto_Orbitals(this->stowf, PARAM.inp.seed_sto);
+        this->stowf.update_sto_orbitals(PARAM.inp.seed_sto);
     }
 }
 
@@ -192,24 +188,23 @@ void ESolver_SDFT_PW<T, Device>::hamilt2density(int istep, int iter, double ethr
     hsolver::DiagoIterAssist<T, Device>::PW_DIAG_NMAX = PARAM.inp.pw_diag_nmax;
 
     // hsolver only exists in this function
-    hsolver::HSolverPW_SDFT<T, Device> hsolver_pw_sdft_obj(
-        &this->kv,
-        this->pw_wfc,
-        &this->wf,
-        this->stowf,
-        this->stoche,
-        this->p_hamilt_sto,
-        PARAM.inp.calculation,
-        PARAM.inp.basis_type,
-        PARAM.inp.ks_solver,
-        PARAM.inp.use_paw,
-        PARAM.globalv.use_uspp,
-        PARAM.inp.nspin,
-        hsolver::DiagoIterAssist<T, Device>::SCF_ITER,
-        hsolver::DiagoIterAssist<T, Device>::PW_DIAG_NMAX,
-        hsolver::DiagoIterAssist<T, Device>::PW_DIAG_THR,
-        hsolver::DiagoIterAssist<T, Device>::need_subspace,
-        this->init_psi);
+    hsolver::HSolverPW_SDFT<T, Device> hsolver_pw_sdft_obj(&this->kv,
+                                                           this->pw_wfc,
+                                                           &this->wf,
+                                                           this->stowf,
+                                                           this->stoche,
+                                                           this->p_hamilt_sto,
+                                                           PARAM.inp.calculation,
+                                                           PARAM.inp.basis_type,
+                                                           PARAM.inp.ks_solver,
+                                                           PARAM.inp.use_paw,
+                                                           PARAM.globalv.use_uspp,
+                                                           PARAM.inp.nspin,
+                                                           hsolver::DiagoIterAssist<T, Device>::SCF_ITER,
+                                                           hsolver::DiagoIterAssist<T, Device>::PW_DIAG_NMAX,
+                                                           hsolver::DiagoIterAssist<T, Device>::PW_DIAG_THR,
+                                                           hsolver::DiagoIterAssist<T, Device>::need_subspace,
+                                                           this->init_psi);
 
     hsolver_pw_sdft_obj.solve(this->p_hamilt, this->psi[0], this->pelec, this->pw_wfc, this->stowf, istep, iter, false);
     this->init_psi = true;

source/module_esolver/esolver_sdft_pw.h

@@ -2,15 +2,15 @@
 #define ESOLVER_SDFT_PW_H
 
 #include "esolver_ks_pw.h"
+#include "module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.h"
+#include "module_hamilt_pw/hamilt_stodft/sto_che.h"
 #include "module_hamilt_pw/hamilt_stodft/sto_iter.h"
 #include "module_hamilt_pw/hamilt_stodft/sto_wf.h"
-#include "module_hamilt_pw/hamilt_stodft/sto_che.h"
-#include "module_hamilt_pw/hamilt_stodft/hamilt_sdft_pw.h"
 
 namespace ModuleESolver
 {
 
-template <typename T, typename Device>
+template <typename T, typename Device = base_device::DEVICE_CPU>
 class ESolver_SDFT_PW : public ESolver_KS_PW<T, Device>
 {
   public:
@@ -26,7 +26,7 @@ class ESolver_SDFT_PW : public ESolver_KS_PW<T, Device>
     void cal_stress(ModuleBase::matrix& stress) override;
 
   public:
-    Stochastic_WF stowf;
+    Stochastic_WF<T, Device> stowf;
     StoChe<double> stoche;
     hamilt::HamiltSdftPW<T, Device>* p_hamilt_sto = nullptr;
 

source/module_hamilt_pw/hamilt_stodft/sto_dos.cpp

@@ -1,8 +1,8 @@
 #include "sto_dos.h"
 
-#include "module_parameter/parameter.h"
 #include "module_base/timer.h"
 #include "module_base/tool_title.h"
+#include "module_parameter/parameter.h"
 #include "sto_tool.h"
 Sto_DOS::~Sto_DOS()
 {
@@ -14,7 +14,7 @@ Sto_DOS::Sto_DOS(ModulePW::PW_Basis_K* p_wfcpw_in,
                  psi::Psi<std::complex<double>>* p_psi_in,
                  hamilt::Hamilt<std::complex<double>>* p_hamilt_in,
                  StoChe<double>& stoche,
-                 Stochastic_WF* p_stowf_in)
+                 Stochastic_WF<std::complex<double>, base_device::DEVICE_CPU>* p_stowf_in)
 {
     this->p_wfcpw = p_wfcpw_in;
     this->p_kv = p_kv_in;
@@ -38,13 +38,7 @@ void Sto_DOS::decide_param(const int& dos_nche,
                            const double& dos_scale)
 {
     this->dos_nche = dos_nche;
-    check_che(this->dos_nche,
-              emin_sto,
-              emax_sto,
-              this->nbands_sto,
-              this->p_kv,
-              this->p_stowf,
-              this->p_hamilt_sto);
+    check_che(this->dos_nche, emin_sto, emax_sto, this->nbands_sto, this->p_kv, this->p_stowf, this->p_hamilt_sto);
     if (dos_setemax)
     {
         this->emax = dos_emax_ev;
@@ -147,12 +141,7 @@ void Sto_DOS::caldos(const double sigmain, const double de, const int npart)
                 }
                 ModuleBase::GlobalFunc::ZEROS(allorderchi.data(), nchipk_new * npwx * dos_nche);
                 std::complex<double>* tmpchi = pchi + start_nchipk * npwx;
-                che.calpolyvec_complex(hchi_norm,
-                                       tmpchi,
-                                       allorderchi.data(),
-                                       npw,
-                                       npwx,
-                                       nchipk_new);
+                che.calpolyvec_complex(hchi_norm, tmpchi, allorderchi.data(), npw, npwx, nchipk_new);
                 double* vec_all = (double*)allorderchi.data();
                 int LDA = npwx * nchipk_new * 2;
                 int M = npwx * nchipk_new * 2;

source/module_hamilt_pw/hamilt_stodft/sto_dos.h

@@ -15,7 +15,7 @@ class Sto_DOS
             psi::Psi<std::complex<double>>* p_psi_in,
             hamilt::Hamilt<std::complex<double>>* p_hamilt_in,
             StoChe<double>& stoche,
-            Stochastic_WF* p_stowf_in);
+            Stochastic_WF<std::complex<double>, base_device::DEVICE_CPU>* p_stowf_in);
     ~Sto_DOS();
 
     /**
@@ -59,8 +59,10 @@ class Sto_DOS
     elecstate::ElecState* p_elec = nullptr;          ///< pointer to the electronic state
     psi::Psi<std::complex<double>>* p_psi = nullptr; ///< pointer to the wavefunction
     hamilt::Hamilt<std::complex<double>>* p_hamilt;  ///< pointer to the Hamiltonian
-    Stochastic_WF* p_stowf = nullptr;                ///< pointer to the stochastic wavefunctions
-    Sto_Func<double> stofunc;                        ///< functions
+
+    Stochastic_WF<std::complex<double>, base_device::DEVICE_CPU>* p_stowf
+        = nullptr;            ///< pointer to the stochastic wavefunctions
+    Sto_Func<double> stofunc; ///< functions
 
     hamilt::HamiltSdftPW<std::complex<double>>* p_hamilt_sto = nullptr; ///< pointer to the Hamiltonian for sDFT
 };

source/module_hamilt_pw/hamilt_stodft/sto_elecond.cpp

@@ -1,11 +1,11 @@
 #include "sto_elecond.h"
 
-#include "module_parameter/parameter.h"
 #include "module_base/complexmatrix.h"
 #include "module_base/constants.h"
 #include "module_base/memory.h"
 #include "module_base/timer.h"
 #include "module_base/vector3.h"
+#include "module_parameter/parameter.h"
 #include "sto_tool.h"
 
 #include <chrono>
@@ -21,7 +21,7 @@ Sto_EleCond::Sto_EleCond(UnitCell* p_ucell_in,
                          pseudopot_cell_vnl* p_ppcell_in,
                          hamilt::Hamilt<std::complex<double>>* p_hamilt_in,
                          StoChe<double>& stoche,
-                         Stochastic_WF* p_stowf_in)
+                         Stochastic_WF<std::complex<double>, base_device::DEVICE_CPU>* p_stowf_in)
     : EleCond(p_ucell_in, p_kv_in, p_elec_in, p_wfcpw_in, p_psi_in, p_ppcell_in)
 {
     this->p_hamilt = p_hamilt_in;
@@ -386,9 +386,10 @@ void Sto_EleCond::cal_jmatrix(const psi::Psi<std::complex<float>>& kspsi_all,
 
         remain -= tmpnb;
         startnb += tmpnb;
-        if (remain == 0) {
+        if (remain == 0)
+        {
             break;
-}
+        }
     }
 
     for (int id = 0; id < ndim; ++id)

source/module_hamilt_pw/hamilt_stodft/sto_elecond.h

@@ -17,7 +17,7 @@ class Sto_EleCond : protected EleCond
                 pseudopot_cell_vnl* p_ppcell_in,
                 hamilt::Hamilt<std::complex<double>>* p_hamilt_in,
                 StoChe<double>& stoche,
-                Stochastic_WF* p_stowf_in);
+                Stochastic_WF<std::complex<double>, base_device::DEVICE_CPU>* p_stowf_in);
     ~Sto_EleCond(){};
     /**
      * @brief Set the N order of Chebyshev expansion for conductivities
@@ -59,8 +59,9 @@ class Sto_EleCond : protected EleCond
     int fd_nche = 0;                                ///< number of Chebyshev orders for Fermi-Dirac function
     int cond_dtbatch = 0;                           ///< number of time steps in a batch
     hamilt::Hamilt<std::complex<double>>* p_hamilt; ///< pointer to the Hamiltonian
-    Stochastic_WF* p_stowf = nullptr;               ///< pointer to the stochastic wavefunctions
-    Sto_Func<double> stofunc;                       ///< functions
+    Stochastic_WF<std::complex<double>, base_device::DEVICE_CPU>* p_stowf
+        = nullptr;            ///< pointer to the stochastic wavefunctions
+    Sto_Func<double> stofunc; ///< functions
 
     hamilt::HamiltSdftPW<std::complex<double>>* p_hamilt_sto = nullptr; ///< pointer to the Hamiltonian for sDFT