initial commit

Author: kirk0830

source/module_esolver/esolver_ks_pw.h

@@ -52,7 +52,7 @@ class ESolver_KS_PW : public ESolver_KS<T, Device>
     //! hide the psi in ESolver_KS for tmp use
     psi::Psi<std::complex<double>, base_device::DEVICE_CPU>* psi = nullptr;
 
-    // psi_initializer controller
+    // PsiInitializer controller
     psi::PSIInit<T, Device>* p_wf_init = nullptr;
 
     Device* ctx = {};

source/module_io/read_input_item_postprocess.cpp

@@ -196,7 +196,7 @@ void ReadInput::item_postprocess()
 
                        In the future lcao_in_pw will have its own ESolver.
 
-                       2023/12/22 use new psi_initializer to expand numerical
+                       2023/12/22 use new PsiInitializer to expand numerical
                       atomic orbitals, ykhuang
                    */
             if (para.input.towannier90 && para.input.basis_type == "lcao_in_pw")

source/module_io/read_input_item_system.cpp

@@ -502,7 +502,7 @@ void ReadInput::item_system()
     }
     {
         Input_Item item("psi_initializer");
-        item.annotation = "whether to use psi_initializer";
+        item.annotation = "whether to use PsiInitializer";
         item.reset_value = [](const Input_Item& item, Parameter& para) {
             if (para.input.basis_type == "lcao_in_pw")
             {

source/module_io/to_wannier90_lcao_in_pw.cpp

@@ -43,7 +43,7 @@ void toWannier90_LCAO_IN_PW::calculate(
 
     Structure_Factor* sf_ptr = const_cast<Structure_Factor*>(&sf);
     ModulePW::PW_Basis_K* wfcpw_ptr = const_cast<ModulePW::PW_Basis_K*>(wfcpw);
-    this->psi_init_ = new psi_initializer_nao<std::complex<double>, base_device::DEVICE_CPU>();
+    this->psi_init_ = new PsiInitializerNAO<std::complex<double>, base_device::DEVICE_CPU>();
 #ifdef __MPI
     this->psi_init_->initialize(sf_ptr, wfcpw_ptr, &(GlobalC::ucell), &(GlobalC::Pkpoints), 1, nullptr, GlobalV::MY_RANK);
     #else

source/module_io/to_wannier90_lcao_in_pw.h

@@ -68,7 +68,7 @@ class toWannier90_LCAO_IN_PW : public toWannier90_PW
   protected:
     const Parallel_Orbitals* ParaV;
     /// @brief psi initializer for expanding nao in planewave basis
-    psi_initializer<std::complex<double>, base_device::DEVICE_CPU>* psi_init_;
+    PsiInitializer<std::complex<double>, base_device::DEVICE_CPU>* psi_init_;
 
     /// @brief get Bloch function from LCAO wavefunction
     /// @param psi_in

source/module_psi/psi_init.cpp

@@ -5,11 +5,11 @@
 #include "module_base/tool_quit.h"
 #include "module_hsolver/diago_iter_assist.h"
 #include "module_parameter/parameter.h"
+#include "module_psi/psi_initializer_random.h"
 #include "module_psi/psi_initializer_atomic.h"
 #include "module_psi/psi_initializer_atomic_random.h"
 #include "module_psi/psi_initializer_nao.h"
 #include "module_psi/psi_initializer_nao_random.h"
-#include "module_psi/psi_initializer_random.h"
 namespace psi
 {
 
@@ -46,27 +46,27 @@ void PSIInit<T, Device>::prepare_init(Structure_Factor* p_sf,
     ModuleBase::timer::tick("PSIInit", "prepare_init");
     if ((this->init_wfc.substr(0, 6) == "atomic") && (p_ucell->natomwfc == 0))
     {
-        this->psi_init = std::unique_ptr<psi_initializer<T, Device>>(new psi_initializer_random<T, Device>());
+        this->psi_init = std::unique_ptr<PsiInitializer<T, Device>>(new PsiInitializerRandom<T, Device>());
     }
     else if (this->init_wfc == "atomic")
     {
-        this->psi_init = std::unique_ptr<psi_initializer<T, Device>>(new psi_initializer_atomic<T, Device>());
+        this->psi_init = std::unique_ptr<PsiInitializer<T, Device>>(new PsiInitializerAtomic<T, Device>());
     }
     else if (this->init_wfc == "random")
     {
-        this->psi_init = std::unique_ptr<psi_initializer<T, Device>>(new psi_initializer_random<T, Device>());
+        this->psi_init = std::unique_ptr<PsiInitializer<T, Device>>(new PsiInitializerRandom<T, Device>());
     }
     else if (this->init_wfc == "nao")
     {
-        this->psi_init = std::unique_ptr<psi_initializer<T, Device>>(new psi_initializer_nao<T, Device>());
+        this->psi_init = std::unique_ptr<PsiInitializer<T, Device>>(new PsiInitializerNAO<T, Device>());
     }
     else if (this->init_wfc == "atomic+random")
     {
-        this->psi_init = std::unique_ptr<psi_initializer<T, Device>>(new psi_initializer_atomic_random<T, Device>());
+        this->psi_init = std::unique_ptr<PsiInitializer<T, Device>>(new PsiInitializerAtomicRandom<T, Device>());
     }
     else if (this->init_wfc == "nao+random")
     {
-        this->psi_init = std::unique_ptr<psi_initializer<T, Device>>(new psi_initializer_nao_random<T, Device>());
+        this->psi_init = std::unique_ptr<PsiInitializer<T, Device>>(new PsiInitializerNAORandom<T, Device>());
     }
     else
     {
@@ -110,7 +110,7 @@ void PSIInit<T, Device>::allocate_psi(Psi<std::complex<double>>*& psi,
     // is not ready yet.
     if (this->use_psiinitializer) // new method
     {
-        // psi_initializer drag initialization of pw wavefunction out of HSolver, make psi
+        // PsiInitializer drag initialization of pw wavefunction out of HSolver, make psi
         // initialization decoupled with HSolver (diagonalization) procedure.
         // However, due to EXX is hard to maintain, we still use the old method for EXX.
         // LCAOINPW in version >= 3.5.0 uses this new method.
@@ -147,9 +147,10 @@ void PSIInit<T, Device>::make_table(const int nks, Structure_Factor* p_sf, pseud
     }
 }
 
+// in the following function, the psi on Device will be initialized with the CPU psi
 template <typename T, typename Device>
-void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
-                                        psi::Psi<T, Device>* kspw_psi,
+void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,      // the one always on CPU
+                                        psi::Psi<T, Device>* kspw_psi,       // the one may be on GPU. In CPU case, it is the same as psi
                                         hamilt::Hamilt<T, Device>* p_hamilt,
                                         const pseudopot_cell_vnl& nlpp,
                                         std::ofstream& ofs_running,
@@ -169,7 +170,8 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
         // like (1, nbands, npwx), in which npwx is the maximal npw of all kpoints
         for (int ik = 0; ik < this->pw_wfc->nks; ik++)
         {
-            //! Fix the wavefunction to initialize at given kpoint
+            //! Fix the wavefunction to initialize at given kpoint. 
+            // This will fix the kpoint for CPU case. For GPU, we should additionally call fix_k for kspw_psi
             psi->fix_k(ik);
 
             //! Update Hamiltonian from other kpoint to the given one
@@ -179,20 +181,20 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
             //! and G is wavevector of the peroiodic part of the Bloch function
             this->psi_init->proj_ao_onkG(ik);
 
-            //! psi_initializer manages memory of psig with shared pointer,
+            //! PsiInitializer manages memory of psig with shared pointer,
             //! its access to use is shared here via weak pointer
-            //! therefore once the psi_initializer is destructed, psig will be destructed, too
+            //! therefore once the PsiInitializer is destructed, psig will be destructed, too
             //! this way, we can avoid memory leak and undefined behavior
-            std::weak_ptr<psi::Psi<T, Device>> psig = this->psi_init->share_psig();
-
-            if (psig.expired())
+            // std::weak_ptr<psi::Psi<T, Device>> psig = this->psi_init->share_psig();
+            psi::Psi<T, Device>* psig_ = this->psi_init->share_psig();
+            if (/*psig.expired()*/ psig_ == nullptr)
             {
                 ModuleBase::WARNING_QUIT("PSIInit::initialize_psi", "psig lifetime is expired");
             }
 
             //! to use psig, we need to lock it to get a shared pointer version,
             //! then switch kpoint of psig to the given one
-            auto psig_ = psig.lock();
+            // auto psig_ = psig.lock();
             // CHANGE LOG: if not lcaoinpw, the psig will only be used in psi-initialization
             // so we can only allocate memory for one kpoint with the maximal number of pw
             // over all kpoints, then the memory space will be always enough. Then for each
@@ -210,6 +212,7 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
                 if (((this->ks_solver == "cg") || (this->ks_solver == "lapack")) && (this->basis_type == "pw"))
                 {
                     // the following function is only run serially, to be improved
+                    // For GPU: this psig_ should be on GPU before calling the following function
                     hsolver::DiagoIterAssist<T, Device>::diagH_subspace_init(p_hamilt,
                                                                              psig_->get_pointer(),
                                                                              psig_->get_nbands(),
@@ -218,6 +221,7 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
                                                                              etatom.data());
                     continue;
                 }
+                // do nothing in LCAO_IN_PW case because psig is used to do transformation instead of initialization
                 else if ((this->ks_solver == "lapack") && (this->basis_type == "lcao_in_pw"))
                 {
                     if (ik == 0)
@@ -239,6 +243,7 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
             }
 
             // for the Davidson method, we just copy the wavefunction (partially)
+            // For GPU: although this is simply the copy operation, if GPU present, this should be a data sending operation
             for (int iband = 0; iband < kspw_psi->get_nbands(); iband++)
             {
                 for (int ibasis = 0; ibasis < kspw_psi->get_nbasis(); ibasis++)
@@ -248,7 +253,7 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
             }
         } // end k-point loop
 
-        if (this->basis_type != "lcao_in_pw")
+        if (this->basis_type != "lcao_in_pw") // if not LCAO_IN_PW case, we can release the memory of psig after initailization is done.
         {
             this->psi_init->deallocate_psig();
         }

source/module_psi/psi_init.h

@@ -41,7 +41,7 @@ class PSIInit
     // make interpolate table
     void make_table(const int nks, Structure_Factor* p_sf, pseudopot_cell_vnl* p_ppcell);
 
-    //------------------------ only for psi_initializer --------------------
+    //------------------------ only for PsiInitializer --------------------
     /**
      * @brief initialize the wavefunction
      *
@@ -58,27 +58,27 @@ class PSIInit
                         const bool is_already_initpsi);
 
     /**
-     * @brief get the psi_initializer
+     * @brief get the PsiInitializer
      *
-     * @return psi_initializer<T, Device>*
+     * @return PsiInitializer<T, Device>*
      */
-    std::weak_ptr<psi::Psi<T, Device>> get_psig() const
+    psi::Psi<T, Device>* get_psig() const
     {
         return this->psi_init->share_psig();
     }
     //----------------------------------------------------------------------
 
   private:
-    // psi_initializer<T, Device>* psi_init = nullptr;
+    // PsiInitializer<T, Device>* psi_init = nullptr;
     // change to use smart pointer to manage the memory, and avoid memory leak
     // while the std::make_unique() is not supported till C++14,
     // so use the new and std::unique_ptr to manage the memory, but this makes new-delete not symmetric
-    std::unique_ptr<psi_initializer<T, Device>> psi_init;
+    std::unique_ptr<PsiInitializer<T, Device>> psi_init;
 
     //! temporary: wave functions, this one may be deleted in future
     wavefunc wf_old;
 
-    // whether to use psi_initializer
+    // whether to usePsiInitializer
     bool use_psiinitializer = false;
 
     // wavefunction initialization type

source/module_psi/psi_initializer.cpp

@@ -12,9 +12,9 @@
 #endif
 
 template<typename T, typename Device>
-psi::Psi<std::complex<double>>* psi_initializer<T, Device>::allocate(const bool only_psig)
+psi::Psi<std::complex<double>>* PsiInitializer<T, Device>::allocate(const bool only_psig)
 {
-    ModuleBase::timer::tick("psi_initializer", "allocate");
+    ModuleBase::timer::tick("PsiInitializer", "allocate");
     /*
         WARNING: when basis_type = "pw", the variable PARAM.globalv.nlocal will also be set, in this case, it is set to
         9 = 1 + 3 + 5, which is the maximal number of orbitals spd, I don't think it is reasonable
@@ -82,13 +82,17 @@ psi::Psi<std::complex<double>>* psi_initializer<T, Device>::allocate(const bool
         // std::cout << " MEMORY FOR PSI PER PROCESSOR (MB)  : " << double(memory_cost_psi)/1024.0/1024.0 << std::endl;
         ModuleBase::Memory::record("Psi_PW", memory_cost_psi);
     }
-    // psi_initializer also works for basis transformation tasks. In this case, psig needs to allocate memory for 
+    // PsiInitializer also works for basis transformation tasks. In this case, psig needs to allocate memory for 
     // each kpoint, otherwise, for initializing pw wavefunction, only one kpoint's space is enough.
     const int nks_psig = (PARAM.inp.basis_type == "pw")? 1 : nks_psi;
-    this->psig_ = std::make_shared<psi::Psi<T, Device>>(nks_psig, 
-                                                        nbands_actual, 
-                                                        nbasis_actual, 
-                                                        this->pw_wfc_->npwk);
+    // this->psig_ = std::make_shared<psi::Psi<T, Device>>(nks_psig, 
+    //                                                     nbands_actual, 
+    //                                                     nbasis_actual, 
+    //                                                     this->pw_wfc_->npwk);
+    this->d_psig_ = new psi::Psi<T, Device>(nks_psig, 
+                                            nbands_actual, 
+                                            nbasis_actual, 
+                                            this->pw_wfc_->npwk);
 
     double memory_cost_psig = 
             nks_psig * nbands_actual * this->pw_wfc_->npwk_max * PARAM.globalv.npol * sizeof(T);
@@ -111,14 +115,14 @@ psi::Psi<std::complex<double>>* psi_initializer<T, Device>::allocate(const bool
                          << "npwk_max = " << this->pw_wfc_->npwk_max << "\n"
                          << "npol = " << PARAM.globalv.npol << "\n";
     ModuleBase::Memory::record("psigPW", memory_cost_psig);
-    ModuleBase::timer::tick("psi_initializer", "allocate");
+    ModuleBase::timer::tick("PsiInitializer", "allocate");
     return psi_out;
 }
 
 template<typename T, typename Device>
-void psi_initializer<T, Device>::random_t(T* psi, const int iw_start, const int iw_end, const int ik)
+void PsiInitializer<T, Device>::random_t(T* psi, const int iw_start, const int iw_end, const int ik)
 {
-    ModuleBase::timer::tick("psi_initializer", "random_t");
+    ModuleBase::timer::tick("PsiInitializer", "random_t");
     assert(iw_start >= 0);
     const int ng = this->pw_wfc_->npwk[ik];
 
@@ -213,14 +217,14 @@ void psi_initializer<T, Device>::random_t(T* psi, const int iw_start, const int
             }
         }
     }
-    ModuleBase::timer::tick("psi_initializer_random", "random_t");
+    ModuleBase::timer::tick("PsiInitializer", "random_t");
 }
 
 #ifdef __MPI
 template<typename T, typename Device>
-void psi_initializer<T, Device>::stick_to_pool(Real* stick, const int& ir, Real* out) const
+void PsiInitializer<T, Device>::stick_to_pool(Real* stick, const int& ir, Real* out) const
 {	
-    ModuleBase::timer::tick("psi_initializer", "stick_to_pool");
+    ModuleBase::timer::tick("PsiInitializer", "stick_to_pool");
 	MPI_Status ierror;
     const int is = this->ixy2is_[ir];
 	const int ip = this->pw_wfc_->fftixy2ip[ir];
@@ -245,7 +249,7 @@ void psi_initializer<T, Device>::stick_to_pool(Real* stick, const int& ir, Real*
         }
         else
         {
-            ModuleBase::WARNING_QUIT("psi_initializer", "stick_to_pool: Real type not supported");
+            ModuleBase::WARNING_QUIT("PsiInitializer", "stick_to_pool: Real type not supported");
         }
 		for(int iz=0; iz<nz; iz++)
 		{
@@ -264,25 +268,25 @@ void psi_initializer<T, Device>::stick_to_pool(Real* stick, const int& ir, Real*
         }
         else
         {
-            ModuleBase::WARNING_QUIT("psi_initializer", "stick_to_pool: Real type not supported");
+            ModuleBase::WARNING_QUIT("PsiInitializer", "stick_to_pool: Real type not supported");
         }
 	}
 
 	return;	
-    ModuleBase::timer::tick("psi_initializer", "stick_to_pool");
+    ModuleBase::timer::tick("PsiInitializer", "stick_to_pool");
 }
 #endif
 
 // explicit instantiation
-template class psi_initializer<std::complex<double>, base_device::DEVICE_CPU>;
-template class psi_initializer<std::complex<float>, base_device::DEVICE_CPU>;
+template class PsiInitializer<std::complex<double>, base_device::DEVICE_CPU>;
+template class PsiInitializer<std::complex<float>, base_device::DEVICE_CPU>;
 // gamma point calculation
-template class psi_initializer<double, base_device::DEVICE_CPU>;
-template class psi_initializer<float, base_device::DEVICE_CPU>;
+template class PsiInitializer<double, base_device::DEVICE_CPU>;
+template class PsiInitializer<float, base_device::DEVICE_CPU>;
 #if ((defined __CUDA) || (defined __ROCM))
-template class psi_initializer<std::complex<double>, base_device::DEVICE_GPU>;
-template class psi_initializer<std::complex<float>, base_device::DEVICE_GPU>;
+template class PsiInitializer<std::complex<double>, base_device::DEVICE_GPU>;
+template class PsiInitializer<std::complex<float>, base_device::DEVICE_GPU>;
 // gamma point calculation
-template class psi_initializer<double, base_device::DEVICE_GPU>;
-template class psi_initializer<float, base_device::DEVICE_GPU>;
+template class PsiInitializer<double, base_device::DEVICE_GPU>;
+template class PsiInitializer<float, base_device::DEVICE_GPU>;
 #endif