fix: memory leak when precision=single (#5839)

Qianruipku · web-flow · commit 74b29543f421 · 2025-01-09T20:46:21.000+08:00
* fix: memory leak when precision=single

* change op

* fix wrong logic of atomic+random
diff --git a/source/module_esolver/esolver_ks_pw.cpp b/source/module_esolver/esolver_ks_pw.cpp
@@ -106,15 +106,18 @@ ESolver_KS_PW<T, Device>::~ESolver_KS_PW()
         container::kernels::destroyGpuBlasHandle();
         container::kernels::destroyGpuSolverHandle();
 #endif
-        delete reinterpret_cast<psi::Psi<T, Device>*>(this->kspw_psi);
     }
 #ifdef __DSP
     std::cout << " ** Closing DSP Hardware..." << std::endl;
     dspDestoryHandle(GlobalV::MY_RANK);
 #endif
+    if(PARAM.inp.device == "gpu" || PARAM.inp.precision == "single")
+    {
+        delete this->kspw_psi;
+    }
     if (PARAM.inp.precision == "single")
     {
-        delete reinterpret_cast<psi::Psi<std::complex<double>, Device>*>(this->__kspw_psi);
+        delete this->__kspw_psi;
     }
 
     delete this->psi;
diff --git a/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.cpp b/source/module_hamilt_pw/hamilt_pwdft/VNL_in_pw.cpp
@@ -532,9 +532,12 @@ void pseudopot_cell_vnl::getvnl(Device* ctx,
     delmem_var_op()(ctx, ylm);
     delmem_var_op()(ctx, vkb1);
     delmem_complex_op()(ctx, sk);
-    if (base_device::get_device_type<Device>(ctx) == base_device::GpuDevice)
+    if (PARAM.inp.device == "gpu" || PARAM.inp.precision == "single")
     {
         delmem_var_op()(ctx, gk);
+    }
+    if (PARAM.inp.device == "gpu")
+    {
         delmem_int_op()(ctx, atom_nh);
         delmem_int_op()(ctx, atom_nb);
         delmem_int_op()(ctx, atom_na);
diff --git a/source/module_io/read_input_item_system.cpp b/source/module_io/read_input_item_system.cpp
@@ -775,12 +775,28 @@ void ReadInput::item_system()
             para.input.device=base_device::information::get_device_flag(
                                 para.inp.device, para.inp.basis_type);
         };
+        item.check_value = [](const Input_Item& item, const Parameter& para) {
+            std::vector<std::string> avail_list = {"cpu", "gpu"};
+            if (std::find(avail_list.begin(), avail_list.end(), para.input.device) == avail_list.end())
+            {
+                const std::string warningstr = nofound_str(avail_list, "device");
+                ModuleBase::WARNING_QUIT("ReadInput", warningstr);
+            }
+        };
         this->add_item(item);
     }
     {
         Input_Item item("precision");
         item.annotation = "the computing precision for ABACUS";
         read_sync_string(input.precision);
+        item.check_value = [](const Input_Item& item, const Parameter& para) {
+            std::vector<std::string> avail_list = {"single", "double"};
+            if (std::find(avail_list.begin(), avail_list.end(), para.input.precision) == avail_list.end())
+            {
+                const std::string warningstr = nofound_str(avail_list, "precision");
+                ModuleBase::WARNING_QUIT("ReadInput", warningstr);
+            }
+        };
         this->add_item(item);
     }
 }
diff --git a/source/module_psi/psi_init.cpp b/source/module_psi/psi_init.cpp
@@ -53,7 +53,7 @@ void PSIInit<T, Device>::prepare_init(const int& random_seed)
         this->psi_initer = std::unique_ptr<psi_initializer<T>>(new psi_initializer_random<T>());
     }
     else if (this->init_wfc == "atomic"
-             || (this->init_wfc == "atomic+random" && this->ucell.natomwfc != PARAM.inp.nbands))
+             || (this->init_wfc == "atomic+random" && this->ucell.natomwfc < PARAM.inp.nbands))
     {
         this->psi_initer = std::unique_ptr<psi_initializer<T>>(new psi_initializer_atomic<T>());
     }
@@ -99,17 +99,30 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
     const int nbands_start = this->psi_initer->nbands_start();
     const int nbands = psi->get_nbands();
     const int nbasis = psi->get_nbasis();
-    const bool another_psi_space = (nbands_start != nbands || PARAM.inp.precision == "single");
+    const bool not_equal = (nbands_start != nbands);
 
     Psi<T>* psi_cpu = reinterpret_cast<psi::Psi<T>*>(psi);
     Psi<T, Device>* psi_device = kspw_psi;
 
-    if (another_psi_space)
+    if (not_equal)
     {
         psi_cpu = new Psi<T>(1, nbands_start, nbasis, nullptr);
         psi_device = PARAM.inp.device == "gpu" ? new psi::Psi<T, Device>(psi_cpu[0])
                                                : reinterpret_cast<psi::Psi<T, Device>*>(psi_cpu);
     }
+    else if (PARAM.inp.precision == "single")
+    {
+        if (PARAM.inp.device == "cpu")
+        {
+            psi_cpu = reinterpret_cast<psi::Psi<T>*>(kspw_psi);
+            psi_device = kspw_psi;
+        }
+        else
+        {
+            psi_cpu = new Psi<T>(1, nbands_start, nbasis, nullptr);
+            psi_device = kspw_psi;  
+        }
+    }
 
     // loop over kpoints, make it possible to only allocate memory for psig at the only one kpt
     // like (1, nbands, npwx), in which npwx is the maximal npw of all kpoints
@@ -126,16 +139,16 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
         this->psi_initer->init_psig(psi_cpu->get_pointer(), ik);
         if (psi_device->get_pointer() != psi_cpu->get_pointer())
         {
-            castmem_h2d_op()(ctx, cpu_ctx, psi_device->get_pointer(), psi_cpu->get_pointer(), nbands_start * nbasis);
+            syncmem_h2d_op()(ctx, cpu_ctx, psi_device->get_pointer(), psi_cpu->get_pointer(), nbands_start * nbasis);
         }
 
         std::vector<typename GetTypeReal<T>::type> etatom(nbands_start, 0.0);
 
         if (this->ks_solver == "cg")
         {
-            if (another_psi_space)
+            if (not_equal)
             {
-                // for diagH_subspace_init, psi_cpu->get_pointer() and kspw_psi->get_pointer() should be different
+                // for diagH_subspace_init, psi_device->get_pointer() and kspw_psi->get_pointer() should be different
                 hsolver::DiagoIterAssist<T, Device>::diagH_subspace_init(p_hamilt,
                                                                          psi_device->get_pointer(),
                                                                          nbands_start,
@@ -145,7 +158,7 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
             }
             else
             {
-                // for diagH_subspace_init, psi_cpu->get_pointer() and kspw_psi->get_pointer() can be the same
+                // for diagH_subspace, psi_device->get_pointer() and kspw_psi->get_pointer() can be the same
                 hsolver::DiagoIterAssist<T, Device>::diagH_subspace(p_hamilt,
                                                                     *psi_device,
                                                                     *kspw_psi,
@@ -155,21 +168,25 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,
         }
         else // dav, bpcg
         {
-            if (another_psi_space)
+            if (psi_device->get_pointer() != kspw_psi->get_pointer())
             {
                 syncmem_complex_op()(ctx, ctx, kspw_psi->get_pointer(), psi_device->get_pointer(), nbands * nbasis);
             }
         }
     } // end k-point loop
 
-    if (another_psi_space)
+    if (not_equal)
     {
         delete psi_cpu;
         if(PARAM.inp.device == "gpu")
         {
             delete psi_device;
         }
     }
+    else if (PARAM.inp.precision == "single" && PARAM.inp.device == "gpu")
+    {
+        delete psi_cpu;
+    }
 
     ModuleBase::timer::tick("PSIInit", "initialize_psi");
 }
diff --git a/source/module_psi/psi_init.h b/source/module_psi/psi_init.h
@@ -82,8 +82,7 @@ class PSIInit
 
     //-------------------------OP--------------------------------------------
     using syncmem_complex_op = base_device::memory::synchronize_memory_op<T, Device, Device>;
-    using castmem_h2d_op
-        = base_device::memory::cast_memory_op<T, T, Device, base_device::DEVICE_CPU>;
+    using syncmem_h2d_op = base_device::memory::synchronize_memory_op<T, Device, base_device::DEVICE_CPU>;
 };
 
 ///@brief allocate the wavefunction

Original file line number	Diff line number	Diff line change
`@@ -53,7 +53,7 @@ void PSIInit<T, Device>::prepare_init(const int& random_seed)`
`53`	`53`	`this->psi_initer = std::unique_ptr<psi_initializer<T>>(new psi_initializer_random<T>());`
`54`	`54`	`}`
`55`	`55`	`else if (this->init_wfc == "atomic"`
`56`		`- \|\| (this->init_wfc == "atomic+random" && this->ucell.natomwfc != PARAM.inp.nbands))`
	`56`	`+ \|\| (this->init_wfc == "atomic+random" && this->ucell.natomwfc < PARAM.inp.nbands))`
`57`	`57`	`{`
`58`	`58`	`this->psi_initer = std::unique_ptr<psi_initializer<T>>(new psi_initializer_atomic<T>());`
`59`	`59`	`}`
`@@ -99,17 +99,30 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,`
`99`	`99`	`const int nbands_start = this->psi_initer->nbands_start();`
`100`	`100`	`const int nbands = psi->get_nbands();`
`101`	`101`	`const int nbasis = psi->get_nbasis();`
`102`		`- const bool another_psi_space = (nbands_start != nbands \|\| PARAM.inp.precision == "single");`
	`102`	`+ const bool not_equal = (nbands_start != nbands);`
`103`	`103`
`104`	`104`	`Psi<T>* psi_cpu = reinterpret_cast<psi::Psi<T>*>(psi);`
`105`	`105`	`Psi<T, Device>* psi_device = kspw_psi;`
`106`	`106`
`107`		`- if (another_psi_space)`
	`107`	`+ if (not_equal)`
`108`	`108`	`{`
`109`	`109`	`psi_cpu = new Psi<T>(1, nbands_start, nbasis, nullptr);`
`110`	`110`	`psi_device = PARAM.inp.device == "gpu" ? new psi::Psi<T, Device>(psi_cpu[0])`
`111`	`111`	`: reinterpret_cast<psi::Psi<T, Device>*>(psi_cpu);`
`112`	`112`	`}`
	`113`	`+ else if (PARAM.inp.precision == "single")`
	`114`	`+ {`
	`115`	`+ if (PARAM.inp.device == "cpu")`
	`116`	`+ {`
	`117`	`+ psi_cpu = reinterpret_cast<psi::Psi<T>*>(kspw_psi);`
	`118`	`+ psi_device = kspw_psi;`
	`119`	`+ }`
	`120`	`+ else`
	`121`	`+ {`
	`122`	`+ psi_cpu = new Psi<T>(1, nbands_start, nbasis, nullptr);`
	`123`	`+ psi_device = kspw_psi;`
	`124`	`+ }`
	`125`	`+ }`
`113`	`126`
`114`	`127`	`// loop over kpoints, make it possible to only allocate memory for psig at the only one kpt`
`115`	`128`	`// like (1, nbands, npwx), in which npwx is the maximal npw of all kpoints`
`@@ -126,16 +139,16 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,`
`126`	`139`	`this->psi_initer->init_psig(psi_cpu->get_pointer(), ik);`
`127`	`140`	`if (psi_device->get_pointer() != psi_cpu->get_pointer())`
`128`	`141`	`{`
`129`		`- castmem_h2d_op()(ctx, cpu_ctx, psi_device->get_pointer(), psi_cpu->get_pointer(), nbands_start * nbasis);`
	`142`	`+ syncmem_h2d_op()(ctx, cpu_ctx, psi_device->get_pointer(), psi_cpu->get_pointer(), nbands_start * nbasis);`
`130`	`143`	`}`
`131`	`144`
`132`	`145`	`std::vector<typename GetTypeReal<T>::type> etatom(nbands_start, 0.0);`
`133`	`146`
`134`	`147`	`if (this->ks_solver == "cg")`
`135`	`148`	`{`
`136`		`- if (another_psi_space)`
	`149`	`+ if (not_equal)`
`137`	`150`	`{`
`138`		`- // for diagH_subspace_init, psi_cpu->get_pointer() and kspw_psi->get_pointer() should be different`
	`151`	`+ // for diagH_subspace_init, psi_device->get_pointer() and kspw_psi->get_pointer() should be different`
`139`	`152`	`hsolver::DiagoIterAssist<T, Device>::diagH_subspace_init(p_hamilt,`
`140`	`153`	`psi_device->get_pointer(),`
`141`	`154`	`nbands_start,`
`@@ -145,7 +158,7 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,`
`145`	`158`	`}`
`146`	`159`	`else`
`147`	`160`	`{`
`148`		`- // for diagH_subspace_init, psi_cpu->get_pointer() and kspw_psi->get_pointer() can be the same`
	`161`	`+ // for diagH_subspace, psi_device->get_pointer() and kspw_psi->get_pointer() can be the same`
`149`	`162`	`hsolver::DiagoIterAssist<T, Device>::diagH_subspace(p_hamilt,`
`150`	`163`	`*psi_device,`
`151`	`164`	`*kspw_psi,`
`@@ -155,21 +168,25 @@ void PSIInit<T, Device>::initialize_psi(Psi<std::complex<double>>* psi,`
`155`	`168`	`}`
`156`	`169`	`else // dav, bpcg`
`157`	`170`	`{`
`158`		`- if (another_psi_space)`
	`171`	`+ if (psi_device->get_pointer() != kspw_psi->get_pointer())`
`159`	`172`	`{`
`160`	`173`	`syncmem_complex_op()(ctx, ctx, kspw_psi->get_pointer(), psi_device->get_pointer(), nbands * nbasis);`
`161`	`174`	`}`
`162`	`175`	`}`
`163`	`176`	`} // end k-point loop`
`164`	`177`
`165`		`- if (another_psi_space)`
	`178`	`+ if (not_equal)`
`166`	`179`	`{`
`167`	`180`	`delete psi_cpu;`
`168`	`181`	`if(PARAM.inp.device == "gpu")`
`169`	`182`	`{`
`170`	`183`	`delete psi_device;`
`171`	`184`	`}`
`172`	`185`	`}`
	`186`	`+ else if (PARAM.inp.precision == "single" && PARAM.inp.device == "gpu")`
	`187`	`+ {`
	`188`	`+ delete psi_cpu;`
	`189`	`+ }`
`173`	`190`
`174`	`191`	`ModuleBase::timer::tick("PSIInit", "initialize_psi");`
`175`	`192`	`}`