Feature: support double type for CG

source/Makefile.Objects

@@ -272,6 +272,7 @@ OBJS_HCONTAINER=base_matrix.o\
 
 OBJS_HSOLVER=diago_cg.o\
     diago_david.o\
+    diagh_consts.o\
     diago_bpcg.o\
     hsolver_pw.o\
     hsolver_pw_sdft.o\

source/module_hsolver/CMakeLists.txt

@@ -1,4 +1,5 @@
 list(APPEND objects
+    diagh_consts.cpp
     diago_cg.cpp
     diago_david.cpp
     diago_bpcg.cpp

source/module_hsolver/diagh.h

@@ -3,11 +3,18 @@
 
 #include <string>
 
-#include "module_base/complexmatrix.h"
 #include "module_base/macros.h"
 #include "module_hamilt_general/hamilt.h"
 #include "module_psi/psi.h"
 
+template<typename T> struct consts
+{
+    consts();
+    T zero;
+    T one;
+    T neg_one;
+};
+
 namespace hsolver
 {
 

source/module_hsolver/diagh_consts.cpp

@@ -0,0 +1,23 @@
+#include "diagh.h"
+template class consts<double>;
+template class consts<std::complex<double>>;
+template class consts<std::complex<float>>;
+
+template<> consts<double>::consts()
+{
+    zero = 0.0;
+    one = 1.0;
+    neg_one = -1.0;
+}
+template<> consts<std::complex<double>>::consts()
+{
+    zero = std::complex<double>(0.0, 0.0);
+    one = std::complex<double>(1.0, 0.0);
+    neg_one = std::complex<double>(-1.0, 0.0);
+}
+template<> consts<std::complex<float>>::consts()
+{
+    zero = std::complex<float>(0.0, 0.0);
+    one = std::complex<float>(1.0, 0.0);
+    neg_one = std::complex<float>(-1.0, 0.0);
+}

source/module_hsolver/diago_cg.cpp

@@ -19,9 +19,9 @@ DiagoCG<T, Device>::DiagoCG(const Real* precondition_in)
     this->precondition = precondition_in;
     test_cg = 0;
     reorder = false;
-    this->one = new T(1.0, 0.0);
-    this->zero = new T(0.0, 0.0);
-    this->neg_one = new T(-1.0, 0.0);
+    this->one = &this->cs.one;
+    this->zero = &this->cs.zero;
+    this->neg_one = &this->cs.neg_one;
 }
 
 template<typename T, typename Device>
@@ -35,9 +35,6 @@ DiagoCG<T, Device>::~DiagoCG() {
     delmem_complex_op()(this->ctx, this->gradient);
     delmem_complex_op()(this->ctx, this->g0);
     delmem_complex_op()(this->ctx, this->lagrange);
-    delete this->one;
-    delete this->zero;
-    delete this->neg_one;
 }
 
 template<typename T, typename Device>
@@ -250,7 +247,7 @@ void DiagoCG<T, Device>::calculate_gradient()
     //     this->gradient[i] -= lambda * this->pphi[i];
     // }
     // haozhihan replace this 2022-10-6
-    constantvector_addORsub_constantVector_op<Real, Device>()(this->ctx, this->dim, this->gradient, this->gradient, 1.0, this->pphi, (-lambda));
+    constantvector_addORsub_constantVector_op<T, Device>()(this->ctx, this->dim, this->gradient, this->gradient, 1.0, this->pphi, (-lambda));
 }
 
 template<typename T, typename Device>
@@ -313,6 +310,16 @@ void DiagoCG<T, Device>::orthogonal_gradient(hamilt::Hamilt<T, Device> *phm_in,
         1);
 }
 
+inline void init_real(const double in, double& out)
+{
+    out = in;
+}
+template<typename FPTYPE>
+inline void init_real(const FPTYPE  in, std::complex<FPTYPE>& out)
+{
+    out = std::complex<FPTYPE>(in, 0.0);
+}
+
 template<typename T, typename Device>
 void DiagoCG<T, Device>::calculate_gamma_cg(const int iter, Real &gg_last, const Real &cg_norm, const Real &theta)
 {
@@ -377,10 +384,11 @@ void DiagoCG<T, Device>::calculate_gamma_cg(const int iter, Real &gg_last, const
         //     pcg[i] = gamma * pcg[i] + this->gradient[i];
         // }
         // haozhihan replace this 2022-10-6
-        constantvector_addORsub_constantVector_op<Real, Device>()(this->ctx, this->dim, pcg, pcg, gamma, this->gradient, 1.0);
+        constantvector_addORsub_constantVector_op<T, Device>()(this->ctx, this->dim, pcg, pcg, gamma, this->gradient, 1.0);
 
         const Real norma = gamma * cg_norm * sin(theta);
-        T znorma(norma * -1, 0.0);
+        T znorma;
+        init_real(norma * -1, znorma);
 
         // haozhihan replace this 2022-10-6
         // const int one = 1;
@@ -389,7 +397,7 @@ void DiagoCG<T, Device>::calculate_gamma_cg(const int iter, Real &gg_last, const
         {
             pcg[i] -= norma * pphi_m[i];
         }*/
-        axpy_op<Real, Device>()(this->ctx, this->dim, &znorma, pphi_m, 1, pcg, 1);
+        axpy_op<T, Device>()(this->ctx, this->dim, &znorma, pphi_m, 1, pcg, 1);
     }
 }
 
@@ -436,7 +444,7 @@ bool DiagoCG<T, Device>::update_psi(Real &cg_norm, Real &theta, Real &eigenvalue
     // }
 
     // haozhihan replace this 2022-10-6
-    constantvector_addORsub_constantVector_op<Real, Device>()(this->ctx, this->dim, phi_m_pointer, phi_m_pointer, cost, pcg, sint_norm);
+    constantvector_addORsub_constantVector_op<T, Device>()(this->ctx, this->dim, phi_m_pointer, phi_m_pointer, cost, pcg, sint_norm);
 
 
     //	std::cout << "\n overlap2 = "  << this->ddot(dim, phi_m, phi_m);
@@ -455,12 +463,36 @@ bool DiagoCG<T, Device>::update_psi(Real &cg_norm, Real &theta, Real &eigenvalue
         // }
 
         // haozhihan replace this 2022-10-6
-        constantvector_addORsub_constantVector_op<Real, Device>()(this->ctx, this->dim, this->sphi, this->sphi, cost, this->scg, sint_norm);
-        constantvector_addORsub_constantVector_op<Real, Device>()(this->ctx, this->dim, this->hphi, this->hphi, cost, this->pphi, sint_norm);
+        constantvector_addORsub_constantVector_op<T, Device>()(this->ctx, this->dim, this->sphi, this->sphi, cost, this->scg, sint_norm);
+        constantvector_addORsub_constantVector_op<T, Device>()(this->ctx, this->dim, this->hphi, this->hphi, cost, this->pphi, sint_norm);
         return 0;
     }
 }
 
+inline double get_real(const double& x)
+{
+    return x;
+}
+inline double get_real(const std::complex<double>& x)
+{
+    return x.real();
+}
+inline float get_real(const std::complex<float>& x)
+{
+    return x.real();
+}
+inline double get_conj(const double& x)
+{
+    return x;
+}
+inline std::complex<double> get_conj(const std::complex<double>& x)
+{
+    return std::conj(x);
+}
+inline std::complex<float> get_conj(const std::complex<float>& x)
+{
+    return std::conj(x);
+}
 template<typename T, typename Device>
 void DiagoCG<T, Device>::schmit_orth(
                           const int &m, // end
@@ -499,9 +531,9 @@ void DiagoCG<T, Device>::schmit_orth(
     // be careful , here reduce m+1
     Parallel_Reduce::reduce_pool(lagrange_so, m + 1);
 
-    T var(0, 0);
+    T var = cs.zero;
     syncmem_complex_d2h_op()(this->cpu_ctx, this->ctx, &var, lagrange_so + m, 1);
-    Real psi_norm = var.real();
+    Real psi_norm = get_real(var);
 
     //<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
     // haozhihan replace 2022-10-6
@@ -536,7 +568,7 @@ void DiagoCG<T, Device>::schmit_orth(
         std::cout << " m = " << m << std::endl;
         for (int j = 0; j <= m; ++j)
         {
-            std::cout << "j = " << j << " lagrange norm = " << (conj(lagrange_so[j]) * lagrange_so[j]).real()
+            std::cout << "j = " << j << " lagrange norm = " << get_real(get_conj(lagrange_so[j]) * lagrange_so[j])
                       << std::endl;
         }
         std::cout << " in DiagoCG, psi norm = " << psi_norm << std::endl;
@@ -599,4 +631,11 @@ template class DiagoCG<std::complex<double>, psi::DEVICE_CPU>;
 template class DiagoCG<std::complex<float>, psi::DEVICE_GPU>;
 template class DiagoCG<std::complex<double>, psi::DEVICE_GPU>;
 #endif 
+
+#ifdef __LCAO
+template class DiagoCG<double, psi::DEVICE_CPU>;
+#if ((defined __CUDA) || (defined __ROCM))
+template class DiagoCG<double, psi::DEVICE_GPU>;
+#endif
+#endif
 } // namespace hsolver

source/module_hsolver/diago_cg.h

@@ -111,6 +111,7 @@ class DiagoCG : public DiagH<T, Device>
     using setmem_var_h_op = psi::memory::set_memory_op<Real, psi::DEVICE_CPU>;
     using syncmem_var_h2d_op = psi::memory::synchronize_memory_op<Real, Device, psi::DEVICE_CPU>;
 
+    consts<T> cs;
     const T * one = nullptr, * zero = nullptr, * neg_one = nullptr;
 };
 

source/module_hsolver/diago_david.cpp

@@ -17,29 +17,6 @@
 
 using namespace hsolver;
 
-template class consts<double>;
-template class consts<std::complex<double>>;
-template class consts<std::complex<float>>;
-
-template<> consts<double>::consts()
-{
-    zero = 0.0;
-    one = 1.0;
-    neg_one = -1.0;
-}
-template<> consts<std::complex<double>>::consts()
-{
-    zero = std::complex<double>(0.0, 0.0);
-    one = std::complex<double>(1.0, 0.0);
-    neg_one = std::complex<double>(-1.0, 0.0);
-}
-template<> consts<std::complex<float>>::consts()
-{
-    zero = std::complex<float>(0.0, 0.0);
-    one = std::complex<float>(1.0, 0.0);
-    neg_one = std::complex<float>(-1.0, 0.0);
-}
-
 inline double get_real(const double& x)
 {
     return x;

source/module_hsolver/diago_david.h

@@ -16,13 +16,6 @@
 #include "module_hamilt_pw/hamilt_pwdft/structure_factor.h"
 #include "module_psi/kernels/device.h"
 
-template<typename T> struct consts
-{
-    consts();
-    T zero;
-    T one;
-    T neg_one;
-};
 namespace hsolver
 {
 

source/module_hsolver/kernels/cuda/math_kernel_op.cu

@@ -260,14 +260,14 @@ __global__ void vector_div_vector_kernel(
     }
 }
 
-template <typename FPTYPE>
+template <typename T>
 __global__ void constantvector_addORsub_constantVector_kernel(
     const int size,
-    thrust::complex<FPTYPE>* result,
-    const thrust::complex<FPTYPE>* vector1,
-    const FPTYPE constant1,
-    const thrust::complex<FPTYPE>* vector2,
-    const FPTYPE constant2)
+    T* result,
+    const T* vector1,
+    const typename GetTypeReal<T>::type constant1,
+    const T* vector2,
+    const typename GetTypeReal<T>::type constant2)
 {
     int i = blockIdx.x * blockDim.x + threadIdx.x;
     if (i < size) 
@@ -563,8 +563,22 @@ void vector_div_vector_op<std::complex<double>, psi::DEVICE_GPU>::operator()(
     vector_div_vector_complex_wrapper(d, dim, result, vector1, vector2);
 }
 // vector operator: result[i] = vector1[i] * constant1 + vector2[i] * constant2
-template <typename FPTYPE> 
-void constantvector_addORsub_constantVector_op<FPTYPE, psi::DEVICE_GPU>::operator()(
+template <>
+void constantvector_addORsub_constantVector_op<double, psi::DEVICE_GPU>::operator()(
+    const psi::DEVICE_GPU* d,
+    const int& dim,
+    double* result,
+    const double* vector1,
+    const double constant1,
+    const double* vector2,
+    const double constant2)
+{
+    int thread = 1024;
+    int block = (dim + thread - 1) / thread;
+    constantvector_addORsub_constantVector_kernel<double> << <block, thread >> > (dim, result, vector1, constant1, vector2, constant2);
+}
+template <typename FPTYPE>
+inline void constantvector_addORsub_constantVector_complex_wrapper(
     const psi::DEVICE_GPU* d,
     const int& dim,
     std::complex<FPTYPE>* result,
@@ -579,11 +593,48 @@ void constantvector_addORsub_constantVector_op<FPTYPE, psi::DEVICE_GPU>::operato
 
     int thread = 1024;
     int block = (dim + thread - 1) / thread;
-    constantvector_addORsub_constantVector_kernel<FPTYPE><<<block, thread>>>(dim, result_tmp, vector1_tmp,constant1, vector2_tmp, constant2);
+    constantvector_addORsub_constantVector_kernel<thrust::complex<FPTYPE>> << <block, thread >> > (dim, result_tmp, vector1_tmp, constant1, vector2_tmp, constant2);
+}
+template <>
+void constantvector_addORsub_constantVector_op<std::complex<float>, psi::DEVICE_GPU>::operator()(
+    const psi::DEVICE_GPU* d,
+    const int& dim,
+    std::complex<float>* result,
+    const std::complex<float>* vector1,
+    const float constant1,
+    const std::complex<float>* vector2,
+    const float constant2)
+{
+    constantvector_addORsub_constantVector_complex_wrapper(d, dim, result, vector1, constant1, vector2, constant2);
+}
+template <>
+void constantvector_addORsub_constantVector_op<std::complex<double>, psi::DEVICE_GPU>::operator()(
+    const psi::DEVICE_GPU* d,
+    const int& dim,
+    std::complex<double>* result,
+    const std::complex<double>* vector1,
+    const double constant1,
+    const std::complex<double>* vector2,
+    const double constant2)
+{
+    constantvector_addORsub_constantVector_complex_wrapper(d, dim, result, vector1, constant1, vector2, constant2);
 }
 
-template <> 
-void axpy_op<float, psi::DEVICE_GPU>::operator()(
+template <>
+void axpy_op<double, psi::DEVICE_GPU>::operator()(
+    const psi::DEVICE_GPU* d,
+    const int& N,
+    const double* alpha,
+    const double* X,
+    const int& incX,
+    double* Y,
+    const int& incY)
+{
+    cublasErrcheck(cublasDaxpy(cublas_handle, N, alpha, X, incX, Y, incY));
+}
+
+template <>
+void axpy_op<std::complex<float>, psi::DEVICE_GPU>::operator()(
     const psi::DEVICE_GPU* d,
     const int& N,
     const std::complex<float> *alpha,
@@ -596,7 +647,7 @@ void axpy_op<float, psi::DEVICE_GPU>::operator()(
 }
 
 template <> 
-void axpy_op<double, psi::DEVICE_GPU>::operator()(
+void axpy_op<std::complex<double>, psi::DEVICE_GPU>::operator()(
     const psi::DEVICE_GPU* d,
     const int& N,
     const std::complex<double> *alpha,
@@ -969,7 +1020,7 @@ template struct line_minimize_with_block_op<std::complex<float>, psi::DEVICE_GPU
 template struct vector_div_constant_op<std::complex<float>, psi::DEVICE_GPU>;
 template struct vector_mul_vector_op<std::complex<float>, psi::DEVICE_GPU>;
 template struct vector_div_vector_op<std::complex<float>, psi::DEVICE_GPU>;
-template struct constantvector_addORsub_constantVector_op<float, psi::DEVICE_GPU>;
+template struct constantvector_addORsub_constantVector_op<std::complex<float>, psi::DEVICE_GPU>;
 template struct matrixSetToAnother<std::complex<float>, psi::DEVICE_GPU>;
 
 template struct dot_real_op<std::complex<double>, psi::DEVICE_GPU>;
@@ -978,7 +1029,7 @@ template struct line_minimize_with_block_op<std::complex<double>, psi::DEVICE_GP
 template struct vector_div_constant_op<std::complex<double>, psi::DEVICE_GPU>;
 template struct vector_mul_vector_op<std::complex<double>, psi::DEVICE_GPU>;
 template struct vector_div_vector_op<std::complex<double>, psi::DEVICE_GPU>;
-template struct constantvector_addORsub_constantVector_op<double, psi::DEVICE_GPU>;
+template struct constantvector_addORsub_constantVector_op<std::complex<double>, psi::DEVICE_GPU>;
 template struct matrixSetToAnother<std::complex<double>, psi::DEVICE_GPU>;
 
 #ifdef __LCAO
@@ -987,5 +1038,6 @@ template struct vector_div_constant_op<double, psi::DEVICE_GPU>;
 template struct vector_mul_vector_op<double, psi::DEVICE_GPU>;
 template struct vector_div_vector_op<double, psi::DEVICE_GPU>;
 template struct matrixSetToAnother<double, psi::DEVICE_GPU>;
+template struct constantvector_addORsub_constantVector_op<double, psi::DEVICE_GPU>;
 #endif
 }  // namespace hsolver