Refactor device parameter in esolver_ks_pw

source/source_base/math_chebyshev.cpp

@@ -61,7 +61,7 @@ Chebyshev<REAL, Device>::Chebyshev(const int norder_in) : fftw(2 * EXTEND * nord
     }
     coefr_cpu = new REAL[norder];
     coefc_cpu = new std::complex<REAL>[norder];
-    if (base_device::get_device_type<Device>(this->ctx) == base_device::GpuDevice)
+    if (base_device::get_device_type(this->ctx) == base_device::GpuDevice)
     {
         resmem_var_op()(this->coef_real, norder);
         resmem_complex_op()(this->coef_complex, norder);
@@ -82,7 +82,7 @@ template <typename REAL, typename Device>
 Chebyshev<REAL, Device>::~Chebyshev()
 {
     delete[] polytrace;
-    if (base_device::get_device_type<Device>(this->ctx) == base_device::GpuDevice)
+    if (base_device::get_device_type(this->ctx) == base_device::GpuDevice)
     {
         delmem_var_op()(this->coef_real);
         delmem_complex_op()(this->coef_complex);
@@ -209,7 +209,7 @@ void Chebyshev<REAL, Device>::calcoef_real(std::function<REAL(REAL)> fun)
         }
     }
 
-    if (base_device::get_device_type<Device>(this->ctx) == base_device::GpuDevice)
+    if (base_device::get_device_type(this->ctx) == base_device::GpuDevice)
     {
         syncmem_var_h2d_op()(coef_real, coefr_cpu, norder);
     }
@@ -299,7 +299,7 @@ void Chebyshev<REAL, Device>::calcoef_complex(std::function<std::complex<REAL>(s
             coefc_cpu[i].imag(imag(coefc_cpu[i]) + real(pcoef[i]) / norder2 * 2 / 3);
         }
     }
-    if (base_device::get_device_type<Device>(this->ctx) == base_device::GpuDevice)
+    if (base_device::get_device_type(this->ctx) == base_device::GpuDevice)
     {
         syncmem_complex_h2d_op()(coef_complex, coefc_cpu, norder);
     }
@@ -390,7 +390,7 @@ void Chebyshev<REAL, Device>::calcoef_pair(std::function<REAL(REAL)> fun1, std::
         }
     }
 
-    if (base_device::get_device_type<Device>(this->ctx) == base_device::GpuDevice)
+    if (base_device::get_device_type(this->ctx) == base_device::GpuDevice)
     {
         syncmem_complex_h2d_op()(coef_complex, coefc_cpu, norder);
     }
@@ -684,7 +684,7 @@ bool Chebyshev<REAL, Device>::checkconverge(
     funA(arrayn_1, arrayn, 1);
     REAL sum1, sum2;
     REAL t;
-    if (base_device::get_device_type<Device>(this->ctx) == base_device::GpuDevice)
+    if (base_device::get_device_type(this->ctx) == base_device::GpuDevice)
     {
         sum1 = this->ddot_real(arrayn_1, arrayn_1, N);
         sum2 = this->ddot_real(arrayn_1, arrayn, N);
@@ -714,7 +714,7 @@ bool Chebyshev<REAL, Device>::checkconverge(
     for (int ior = 2; ior < norder; ++ior)
     {
         funA(arrayn, arraynp1, 1);
-        if (base_device::get_device_type<Device>(this->ctx) == base_device::GpuDevice)
+        if (base_device::get_device_type(this->ctx) == base_device::GpuDevice)
         {
             sum1 = this->ddot_real(arrayn, arrayn, N);
             sum2 = this->ddot_real(arrayn, arraynp1, N);

source/source_base/module_device/device.h

@@ -145,6 +145,36 @@ class DeviceContext {
      */
     int get_local_rank() const { return local_rank_; }
 
+    /**
+     * @brief Set the device type (CpuDevice, GpuDevice, or DspDevice)
+     * @param type The device type
+     */
+    void set_device_type(AbacusDevice_t type) { device_type_ = type; }
+
+    /**
+     * @brief Get the device type
+     * @return AbacusDevice_t The device type
+     */
+    AbacusDevice_t get_device_type() const { return device_type_; }
+
+    /**
+     * @brief Check if the device is CPU
+     * @return true if the device is CPU
+     */
+    bool is_cpu() const { return device_type_ == CpuDevice; }
+
+    /**
+     * @brief Check if the device is GPU
+     * @return true if the device is GPU
+     */
+    bool is_gpu() const { return device_type_ == GpuDevice; }
+
+    /**
+     * @brief Check if the device is DSP
+     * @return true if the device is DSP
+     */
+    bool is_dsp() const { return device_type_ == DspDevice; }
+
     // Disable copy and assignment
     DeviceContext(const DeviceContext&) = delete;
     DeviceContext& operator=(const DeviceContext&) = delete;
@@ -158,10 +188,21 @@ class DeviceContext {
     int device_id_ = -1;
     int device_count_ = 0;
     int local_rank_ = 0;
+    AbacusDevice_t device_type_ = CpuDevice;
 
     std::mutex init_mutex_;
 };
 
+/**
+ * @brief Get the device type enum from DeviceContext (runtime version).
+ * @param ctx Pointer to DeviceContext
+ * @return AbacusDevice_t enum value
+ */
+inline AbacusDevice_t get_device_type(const DeviceContext* ctx)
+{
+    return ctx->get_device_type();
+}
+
 } // end of namespace base_device
 
 #endif // MODULE_DEVICE_H_

source/source_base/module_device/device_helpers.cpp

@@ -3,19 +3,6 @@
 namespace base_device
 {
 
-// Device type specializations
-template <>
-AbacusDevice_t get_device_type<DEVICE_CPU>(const DEVICE_CPU* dev)
-{
-    return CpuDevice;
-}
-
-template <>
-AbacusDevice_t get_device_type<DEVICE_GPU>(const DEVICE_GPU* dev)
-{
-    return GpuDevice;
-}
-
 // Precision specializations
 template <>
 std::string get_current_precision<float>(const float* var)

source/source_base/module_device/device_helpers.h

@@ -13,25 +13,35 @@
 #include "types.h"
 #include <complex>
 #include <string>
+#include <type_traits>
 
 namespace base_device
 {
 
+// Forward declaration
+class DeviceContext;
+
+/**
+ * @brief Get the device type enum from DeviceContext (runtime version).
+ * @param ctx Pointer to DeviceContext
+ * @return AbacusDevice_t enum value
+ */
+inline AbacusDevice_t get_device_type(const DeviceContext* ctx);
+
 /**
- * @brief Get the device type enum for a given device type.
+ * @brief Get the device type enum for a given device type (compile-time version).
  * @tparam Device The device type (DEVICE_CPU or DEVICE_GPU)
  * @param dev Pointer to device (used for template deduction)
  * @return AbacusDevice_t enum value
  */
 template <typename Device>
-AbacusDevice_t get_device_type(const Device* dev);
-
-// Template specialization declarations
-template <>
-AbacusDevice_t get_device_type<DEVICE_CPU>(const DEVICE_CPU* dev);
-
-template <>
-AbacusDevice_t get_device_type<DEVICE_GPU>(const DEVICE_GPU* dev);
+AbacusDevice_t get_device_type(const Device* dev)
+{
+    if (std::is_same<Device, DEVICE_CPU>::value) return CpuDevice;
+    else if (std::is_same<Device, DEVICE_GPU>::value) return GpuDevice;
+    else if (std::is_same<Device, DEVICE_DSP>::value) return DspDevice;
+    else return UnKnown;
+}
 
 /**
  * @brief Get the precision string for a given numeric type.

source/source_base/module_device/test/device_test.cpp

@@ -20,14 +20,14 @@ class TestModulePsiDevice : public ::testing::Test
 
 TEST_F(TestModulePsiDevice, get_device_type_cpu)
 {
-    base_device::AbacusDevice_t device = base_device::get_device_type<base_device::DEVICE_CPU>(cpu_ctx);
+    base_device::AbacusDevice_t device = base_device::get_device_type(cpu_ctx);
     EXPECT_EQ(device, base_device::CpuDevice);
 }
 
 #if __UT_USE_CUDA || __UT_USE_ROCM
 TEST_F(TestModulePsiDevice, get_device_type_gpu)
 {
-    base_device::AbacusDevice_t device = base_device::get_device_type<base_device::DEVICE_GPU>(gpu_ctx);
+    base_device::AbacusDevice_t device = base_device::get_device_type(gpu_ctx);
     EXPECT_EQ(device, base_device::GpuDevice);
 }
 #endif // __UT_USE_CUDA || __UT_USE_ROCM

source/source_esolver/esolver_ks_pw.cpp

@@ -43,7 +43,7 @@ ESolver_KS_PW<T, Device>::ESolver_KS_PW()
 {
     this->classname = "ESolver_KS_PW";
     this->basisname = "PW";
-    this->device = base_device::get_device_type<Device>(this->ctx);
+    this->ctx = nullptr;
 }
 
 template <typename T, typename Device>
@@ -251,7 +251,7 @@ void ESolver_KS_PW<T, Device>::after_scf(UnitCell& ucell, const int istep, const
     // Output quantities
     ModuleIO::ctrl_scf_pw<T, Device>(istep, ucell, this->pelec, this->chr, this->kv, this->pw_wfc,
               this->pw_rho, this->pw_rhod, this->pw_big, this->stp,
-              this->ctx, this->device, this->Pgrid, PARAM.inp);
+              this->ctx, this->Pgrid, PARAM.inp);
 
     ModuleBase::timer::tick("ESolver_KS_PW", "after_scf");
 }

source/source_esolver/esolver_ks_pw.h

@@ -60,9 +60,6 @@ class ESolver_KS_PW : public ESolver_KS
     // for get_pchg and get_wf, use ctx as input of fft
     Device* ctx = {};
 
-    // for device to host data transformation
-    base_device::AbacusDevice_t device = {};
-
 };
 } // namespace ModuleESolver
 #endif

source/source_hsolver/diago_dav_subspace.cpp

@@ -36,7 +36,7 @@ Diago_DavSubspace<T, Device>::Diago_DavSubspace(const std::vector<Real>& precond
       diag_thr(diag_thr_in), iter_nmax(diag_nmax_in), diag_comm(diag_comm_in),
         diag_subspace(diag_subspace_in), diago_subspace_bs(diago_subspace_bs_in)
 {
-    this->device = base_device::get_device_type<Device>(this->ctx);
+    this->device = base_device::get_device_type(this->ctx);
 
     this->one = &one_;
     this->zero = &zero_;

source/source_hsolver/diago_david.cpp

@@ -20,7 +20,7 @@ DiagoDavid<T, Device>::DiagoDavid(const Real* precondition_in,
                                   const diag_comm_info& diag_comm_in)
     : nband(nband_in), dim(dim_in), nbase_x(david_ndim_in * nband_in), david_ndim(david_ndim_in), use_paw(use_paw_in), diag_comm(diag_comm_in)
 {
-    this->device = base_device::get_device_type<Device>(this->ctx);
+    this->device = base_device::get_device_type(this->ctx);
     this->precondition = precondition_in;
 
     this->one = &one_;

source/source_hsolver/diago_iter_assist.cpp

@@ -400,14 +400,14 @@ void DiagoIterAssist<T, Device>::diag_heevx(const int matrix_size,
     // (const Device *d, const int matrix_size, const int lda, const T *A, const int num_eigenpairs, Real *eigenvalues, T *eigenvectors);
     heevx_op<T, Device>()(ctx, matrix_size, ldh, h, num_eigenpairs, eigenvalues, v);
 
-    if (base_device::get_device_type<Device>(ctx) == base_device::GpuDevice)
+    if (base_device::get_device_type(ctx) == base_device::GpuDevice)
     {
 #if ((defined __CUDA) || (defined __ROCM))
         // eigenvalues to e, from device to host
         syncmem_var_d2h_op()(e, eigenvalues, num_eigenpairs);
 #endif
     }
-    else if (base_device::get_device_type<Device>(ctx) == base_device::CpuDevice)
+    else if (base_device::get_device_type(ctx) == base_device::CpuDevice)
     {
         // eigenvalues to e
         syncmem_var_op()(e, eigenvalues, num_eigenpairs);
@@ -436,14 +436,14 @@ void DiagoIterAssist<T, Device>::diag_hegvd(const int nstart,
 
     hegvd_op<T, Device>()(ctx, nstart, ldh, hcc, scc, eigenvalues, vcc);
 
-    if (base_device::get_device_type<Device>(ctx) == base_device::GpuDevice)
+    if (base_device::get_device_type(ctx) == base_device::GpuDevice)
     {
 #if ((defined __CUDA) || (defined __ROCM))
         // set eigenvalues in GPU to e in CPU
         syncmem_var_d2h_op()(e, eigenvalues, nbands);
 #endif
     }
-    else if (base_device::get_device_type<Device>(ctx) == base_device::CpuDevice)
+    else if (base_device::get_device_type(ctx) == base_device::CpuDevice)
     {
         // set eigenvalues in CPU to e in CPU
         syncmem_var_op()(e, eigenvalues, nbands);

source/source_hsolver/test/hsolver_pw_sup.h

@@ -126,7 +126,7 @@ DiagoDavid<T, Device>::DiagoDavid(const Real* precondition_in,
                                   const bool use_paw_in,
                                   const diag_comm_info& diag_comm_in)
     : nband(nband_in), dim(dim_in), nbase_x(david_ndim_in * nband_in), david_ndim(david_ndim_in), use_paw(use_paw_in), diag_comm(diag_comm_in) {
-    this->device = base_device::get_device_type<Device>(this->ctx);
+    this->device = base_device::get_device_type(this->ctx);
     this->precondition = precondition_in;
 
     test_david = 2;

source/source_io/module_ctrl/ctrl_output_pw.cpp

@@ -92,15 +92,14 @@ void ModuleIO::ctrl_scf_pw(const int istep,
 		const ModulePW::PW_Basis_Big *pw_big,
         Setup_Psi_pw<T, Device> &stp,
         const Device* ctx,
-        const base_device::AbacusDevice_t &device,
         const Parallel_Grid &para_grid,
         const Input_para& inp)
 {
     ModuleBase::TITLE("ModuleIO", "ctrl_scf_pw");
     ModuleBase::timer::tick("ModuleIO", "ctrl_scf_pw");
 
     // Transfer data from device (GPU) to host (CPU) in pw basis
-    stp.copy_d2h(device);
+    stp.copy_d2h(ctx);
 
     //----------------------------------------------------------
     //! 4) Compute density of states (DOS)
@@ -386,7 +385,6 @@ template void ModuleIO::ctrl_scf_pw<std::complex<float>, base_device::DEVICE_CPU
     const ModulePW::PW_Basis_Big *pw_big,
     Setup_Psi_pw<std::complex<float>, base_device::DEVICE_CPU> &stp,
     const base_device::DEVICE_CPU* ctx,
-    const base_device::AbacusDevice_t &device,
     const Parallel_Grid &para_grid,
     const Input_para& inp);
 
@@ -403,7 +401,6 @@ template void ModuleIO::ctrl_scf_pw<std::complex<double>, base_device::DEVICE_CP
     const ModulePW::PW_Basis_Big *pw_big,
     Setup_Psi_pw<std::complex<double>, base_device::DEVICE_CPU> &stp,
     const base_device::DEVICE_CPU* ctx,
-    const base_device::AbacusDevice_t &device,
     const Parallel_Grid &para_grid,
     const Input_para& inp);
 
@@ -421,7 +418,6 @@ template void ModuleIO::ctrl_scf_pw<std::complex<float>, base_device::DEVICE_GPU
     const ModulePW::PW_Basis_Big *pw_big,
     Setup_Psi_pw<std::complex<float>, base_device::DEVICE_GPU> &stp,
     const base_device::DEVICE_GPU* ctx,
-    const base_device::AbacusDevice_t &device,
     const Parallel_Grid &para_grid,
     const Input_para& inp);
 
@@ -438,7 +434,6 @@ template void ModuleIO::ctrl_scf_pw<std::complex<double>, base_device::DEVICE_GP
     const ModulePW::PW_Basis_Big *pw_big,
     Setup_Psi_pw<std::complex<double>, base_device::DEVICE_GPU> &stp,
     const base_device::DEVICE_GPU* ctx,
-    const base_device::AbacusDevice_t &device,
     const Parallel_Grid &para_grid,
     const Input_para& inp);
 #endif