Fix: port sdft memory op to DSP (#7203)

* Fix DSP sdft memory op

* Fix include guard

Author: Cstandardlib

source/source_pw/module_stodft/sto_che.h

@@ -26,9 +26,14 @@ class StoChe
 
   private:
     Device* ctx = {};
+#ifdef __DSP
+    using resmem_var_op = base_device::memory::resize_memory_op_mt<REAL, Device>;
+    using delmem_var_op = base_device::memory::delete_memory_op_mt<REAL, Device>;
+#else
     using resmem_var_op = base_device::memory::resize_memory_op<REAL, Device>;
-    using syncmem_var_d2h_op = base_device::memory::synchronize_memory_op<REAL, base_device::DEVICE_CPU, Device>;
     using delmem_var_op = base_device::memory::delete_memory_op<REAL, Device>;
+#endif
+    using syncmem_var_d2h_op = base_device::memory::synchronize_memory_op<REAL, base_device::DEVICE_CPU, Device>;
 };
 
 /**
@@ -64,4 +69,4 @@ REAL vTMv(const REAL* v, const REAL* M, const int n)
     return result;
 }
 
-#endif
+#endif

source/source_pw/module_stodft/sto_elecond.h

@@ -6,7 +6,7 @@
 #include "source_pw/module_stodft/sto_wf.h"
 #include "source_hsolver/hsolver_pw_sdft.h"
 
-template <typename FPTYPE, typename Device> 
+template <typename FPTYPE, typename Device>
 class Sto_EleCond : protected EleCond<FPTYPE, Device>
 {
   public:
@@ -16,8 +16,13 @@ class Sto_EleCond : protected EleCond<FPTYPE, Device>
     using lowTYPE = double;
 #endif
     using lcomplex = std::complex<lowTYPE>;
-    using resmem_lcomplex_op = base_device::memory::resize_memory_op<std::complex<lowTYPE>, Device>; 
+#ifdef __DSP
+    using resmem_lcomplex_op = base_device::memory::resize_memory_op_mt<std::complex<lowTYPE>, Device>;
+    using delmem_lcomplex_op = base_device::memory::delete_memory_op_mt<std::complex<lowTYPE>, Device>;
+#else
+    using resmem_lcomplex_op = base_device::memory::resize_memory_op<std::complex<lowTYPE>, Device>;
     using delmem_lcomplex_op = base_device::memory::delete_memory_op<std::complex<lowTYPE>, Device>;
+#endif
     using cpymem_lcomplex_op = base_device::memory::synchronize_memory_op<std::complex<lowTYPE>, Device, Device>;
     using castmem_lcomplex_op = base_device::memory::cast_memory_op<std::complex<lowTYPE>, std::complex<FPTYPE>, Device, Device>;
     using cpymem_complex_op = base_device::memory::synchronize_memory_op<std::complex<FPTYPE>, Device, Device>;
@@ -114,4 +119,4 @@ class Sto_EleCond : protected EleCond<FPTYPE, Device>
                      const std::complex<lowTYPE>& factor,
                      const int bandinfo[6]);
 };
-#endif // ELECOND_H
+#endif // STOELECOND_H

source/source_pw/module_stodft/sto_forces.h

@@ -44,9 +44,13 @@ class Sto_Forces : public Forces<FPTYPE, Device>
                           const UnitCell& ucell,
                           const psi::Psi<std::complex<FPTYPE>, Device>& psi_in,
                           const Stochastic_WF<std::complex<FPTYPE>, Device>& stowf);
-
+#ifdef __DSP
+    using resmem_var_op = base_device::memory::resize_memory_op_mt<FPTYPE, Device>;
+    using delmem_var_op = base_device::memory::delete_memory_op_mt<FPTYPE, Device>;
+#else
     using resmem_var_op = base_device::memory::resize_memory_op<FPTYPE, Device>;
     using delmem_var_op = base_device::memory::delete_memory_op<FPTYPE, Device>;
+#endif
     using syncmem_var_h2d_op = base_device::memory::synchronize_memory_op<FPTYPE, Device, base_device::DEVICE_CPU>;
     using syncmem_var_d2h_op = base_device::memory::synchronize_memory_op<FPTYPE, base_device::DEVICE_CPU, Device>;
 };

source/source_pw/module_stodft/sto_iter.h

@@ -46,7 +46,7 @@ class Stochastic_Iter
 
     /**
      * @brief sum demet and eband energies for each k point and each band
-     * 
+     *
      * @param stowf stochastic wave function
      * @param pes elecstate
      * @param pHamilt hamiltonian
@@ -59,7 +59,7 @@ class Stochastic_Iter
 
     /**
      * @brief calculate the density
-     * 
+     *
      * @param ucell reference to unit cell
      * @param stowf stochastic wave function
      * @param pes elecstate
@@ -72,23 +72,23 @@ class Stochastic_Iter
 
     /**
      * @brief calculate total number of electrons
-     * 
+     *
      * @param pes elecstate
-     * @return double 
+     * @return double
      */
     double calne(elecstate::ElecState* pes);
 
     /**
      * @brief solve ne(mu) = ne_target and get chemical potential mu
-     * 
+     *
      * @param iter scf iteration index
      * @param pes elecstate
      */
     void itermu(const int iter, elecstate::ElecState* pes);
 
     /**
      * @brief orthogonalize stochastic wave functions with KS wave functions
-     * 
+     *
      * @param ik k point index
      * @param psi KS wave functions
      * @param stowf stochastic wave functions
@@ -97,7 +97,7 @@ class Stochastic_Iter
 
     /**
      * @brief check emax and emin
-     * 
+     *
      * @param ik k point index
      * @param istep ion step index
      * @param iter scf iteration index
@@ -107,7 +107,7 @@ class Stochastic_Iter
 
     /**
      * @brief check precision of Chebyshev expansion
-     * 
+     *
      * @param ref reference value
      * @param thr threshold
      * @param info information
@@ -153,14 +153,22 @@ class Stochastic_Iter
     const Device* ctx = {};
     const base_device::DEVICE_CPU* cpu_ctx = {};
     using ct_Device = typename container::PsiToContainer<Device>::type;
+#ifdef __DSP
+    using setmem_var_op = base_device::memory::set_memory_op_mt<Real, Device>;
+    using resmem_var_op = base_device::memory::resize_memory_op_mt<Real, Device>;
+    using delmem_var_op = base_device::memory::delete_memory_op_mt<Real, Device>;
+    using resmem_complex_op = base_device::memory::resize_memory_op_mt<T, Device>;
+    using delmem_complex_op = base_device::memory::delete_memory_op_mt<T, Device>;
+#else
     using setmem_var_op = base_device::memory::set_memory_op<Real, Device>;
-    using syncmem_var_h2d_op = base_device::memory::synchronize_memory_op<Real, Device, base_device::DEVICE_CPU>;
-    using syncmem_var_d2h_op = base_device::memory::synchronize_memory_op<Real, base_device::DEVICE_CPU, Device>;
-    using cpymem_complex_op = base_device::memory::synchronize_memory_op<T, Device, Device>;
     using resmem_var_op = base_device::memory::resize_memory_op<Real, Device>;
     using delmem_var_op = base_device::memory::delete_memory_op<Real, Device>;
     using resmem_complex_op = base_device::memory::resize_memory_op<T, Device>;
     using delmem_complex_op = base_device::memory::delete_memory_op<T, Device>;
+#endif
+    using syncmem_var_h2d_op = base_device::memory::synchronize_memory_op<Real, Device, base_device::DEVICE_CPU>;
+    using syncmem_var_d2h_op = base_device::memory::synchronize_memory_op<Real, base_device::DEVICE_CPU, Device>;
+    using cpymem_complex_op = base_device::memory::synchronize_memory_op<T, Device, Device>;
     using castmem_d2z_op = base_device::memory::cast_memory_op<T, Real, Device, Device>;
     using castmem_var_d2h_op = base_device::memory::cast_memory_op<double, Real, base_device::DEVICE_CPU, Device>;
     using gemv_op = ModuleBase::gemv_op<T, Device>;

source/source_pw/module_stodft/sto_stress_pw.h

@@ -52,9 +52,16 @@ class Sto_Stress_PW : public Stress_Func<FPTYPE, Device>
                        const Stochastic_WF<std::complex<FPTYPE>, Device>& stowf);
 
   private:
+#ifdef __DSP
+    using resmem_var_op = base_device::memory::resize_memory_op_mt<FPTYPE, Device>;
+    using setmem_var_op = base_device::memory::set_memory_op_mt<FPTYPE, Device>;
+    using delmem_var_op = base_device::memory::delete_memory_op_mt<FPTYPE, Device>;
+
+#else
     using resmem_var_op = base_device::memory::resize_memory_op<FPTYPE, Device>;
     using setmem_var_op = base_device::memory::set_memory_op<FPTYPE, Device>;
     using delmem_var_op = base_device::memory::delete_memory_op<FPTYPE, Device>;
+#endif
     using syncmem_var_d2h_op = base_device::memory::synchronize_memory_op<FPTYPE, base_device::DEVICE_CPU, Device>;
 };
 #endif