the segfault caused by device gpu

Author: kirk0830

source/module_psi/psi_init.cpp

@@ -27,6 +27,19 @@ PSIInit<T, Device>::PSIInit(const std::string& init_wfc_in,
     this->pw_wfc = pw_wfc_in;
 }
 
+template <typename T, typename Device>
+PSIInit<T, Device>::~PSIInit()
+{
+    if (this->use_psiinitializer)
+    {
+        {
+            this->psi_init->deallocate_psig();
+            // delete this->psi_init;
+            // this->psi_init = nullptr;  
+        }
+    }
+}
+
 template <typename T, typename Device>
 void PSIInit<T, Device>::prepare_init(Structure_Factor* p_sf,
                                       UnitCell* p_ucell,

source/module_psi/psi_init.h

@@ -17,7 +17,7 @@ class PSIInit
             const std::string& basis_type_in,
             const bool& use_psiinitializer_in,
             ModulePW::PW_Basis_K* pw_wfc_in);
-    ~PSIInit(){ this->psi_init->deallocate_psig(); }
+    ~PSIInit();
 
     // prepare the wavefunction initialization
     void prepare_init(Structure_Factor* p_sf, //< structure factor

source/module_psi/psi_initializer.cpp

@@ -11,6 +11,15 @@
 #include "module_base/parallel_reduce.h"
 #endif
 
+template<typename T, typename Device>
+PsiInitializer<T, Device>::PsiInitializer()
+{
+    // the following line will simultaneously run by CPU and GPU
+    // For CPU, this->device will get value base_device::CpuDevice
+    // For GPU, this->device will get value base_device::GpuDevice
+    this->device = base_device::get_device_type<Device>(this->ctx);
+}
+
 template<typename T, typename Device>
 psi::Psi<std::complex<double>>* PsiInitializer<T, Device>::allocate(const bool only_psig)
 {
@@ -92,7 +101,7 @@ psi::Psi<std::complex<double>>* PsiInitializer<T, Device>::allocate(const bool o
     this->d_psig_ = new psi::Psi<T, Device>(nks_psig, 
                                             nbands_actual, 
                                             nbasis_actual, 
-                                            this->pw_wfc_->npwk);
+                                            this->pw_wfc_->npwk); // psig can be directly allocate on GPU
 
     double memory_cost_psig = 
             nks_psig * nbands_actual * this->pw_wfc_->npwk_max * PARAM.globalv.npol * sizeof(T);

source/module_psi/psi_initializer.h

@@ -68,7 +68,7 @@ class PsiInitializer
     public:
         // technical notes:
         // Polymorphism is used to implement different methods, and achieved by pointers and virtual functions
-        PsiInitializer() {};
+        PsiInitializer();
         virtual ~PsiInitializer() {};
         #ifdef __MPI // MPI additional implementation
         /// @brief initialize the PsiInitializer with external data and methods
@@ -197,7 +197,10 @@ class PsiInitializer
         // refactored psig, in old version it is of datatype Psi<T, Device>*, use std::shared_ptr to
         // avoid memory leak
         // std::shared_ptr<psi::Psi<T, Device>> psig_;
-
+        
+        Device* ctx = {}; // device information carrier. The ctx can have type DEVICE_CPU or DEVICE_GPU
+        base_device::DEVICE_CPU* cpu_ctx = {};
+        base_device::AbacusDevice_t device = {};
         psi::Psi<T, Device>* d_psig_ = nullptr; // is it possible to directly operate data on GPU?
     private:
         int mem_saver_ = 0;