diago_dav_subspace.h

Go to the documentation of this file.

#ifndef DIAGO_NEW_DAV_H
#define DIAGO_NEW_DAV_H
 
#include "source_base/macros.h"   // GetRealType
#include "source_base/module_device/device.h"   // base_device
#include "source_base/module_device/memory_op.h"// base_device::memory"
 
#include "source_base/module_container/ATen/kernels/lapack.h"
 
#include "source_hsolver/diag_comm_info.h"
#include "source_hsolver/diag_const_nums.h"
 
#include <vector>
#include <functional>
 
namespace hsolver
{
 
template <typename T = std::complex<double>, typename Device = base_device::DEVICE_CPU>
class Diago_DavSubspace
{
  private:
    // Note GetTypeReal<T>::type will
    // return T if T is real type(float, double),
    // otherwise return the real type of T(complex<float>, std::complex<double>)
    using Real = typename GetTypeReal<T>::type;
 
  public:
    Diago_DavSubspace(const std::vector<Real>& precondition_in,
                      const int& nband_in,
                      const int& nbasis_in,
                      const int& david_ndim_in,
                      const double& diag_thr_in,
                      const int& diag_nmax_in,
                      const diag_comm_info& diag_comm_in,
                      const int diago_dav_method_in,
                      const int block_size_in);
 
    ~Diago_DavSubspace();
 
    // See diago_david.h for information on the HPsiFunc function type
    using HPsiFunc = std::function<void(T*, T*, const int, const int)>;
 
    int diag(const HPsiFunc& hpsi_func,
             const HPsiFunc& spsi_func,
             T* psi_in,
             const int psi_in_dmax,
             Real* eigenvalue_in,
             const std::vector<double>& ethr_band,
             const bool& scf_type);
 
  private:
    const diag_comm_info diag_comm;
 
    const double diag_thr;
 
    const int iter_nmax;
 
    const int n_band = 0;
 
    const int dim = 0;
 
    const int nbase_x = 0;
 
    const std::vector<Real>& precondition;
    Real* d_precondition = nullptr;
 
    int notconv = 0;
 
    T* psi_in_iter = nullptr;
 
    T* hpsi = nullptr;
 
    T* spsi = nullptr;
 
    T* hcc = nullptr;
 
    T* scc = nullptr;
 
    T* vcc = nullptr;
 
    T* d_scc = nullptr;
    Real* d_eigenvalue = nullptr;
 
    Device* ctx = {};
    base_device::DEVICE_CPU* cpu_ctx = {};
    base_device::AbacusDevice_t device = {};
 
    void cal_grad(const HPsiFunc& hpsi_func,
                  const HPsiFunc& spsi_func,
                  const int& dim,
                  const int& nbase,
                  const int& notconv,
                  T* psi_iter,
                  T* hpsi,
                  T* spsi,
                  T* vcc,
                  const int* unconv,
                  std::vector<Real>* eigenvalue_iter);
 
    void cal_elem(const int& dim,
                  int& nbase,
                  const int& notconv,
                  const T* psi_iter,
                  const T* spsi,
                  const T* hpsi,
                  T* hcc,
                  T* scc);
 
    void refresh(const int& dim,
                 const int& nband,
                 int& nbase,
                 const Real* eigenvalue,
                 T* psi_iter,
                 T* hpsi,
                 T* spsi,
                 T* hcc,
                 T* scc,
                 T* vcc);
 
    // void diagH_LAPACK(const int nstart,
    //                   const int nbands,
    //                   const T* hcc,
    //                   const T* sc,
    //                   const int ldh, // nstart
    //                   Real* e,
    //                   T* vcc);
 
    void diag_zhegvx(const int& nbase,
                     const int& nband,
                     T* hcc,
                     T* scc,
                     const int& nbase_x,
                     std::vector<Real>* eigenvalue_iter,
                     T* vcc);
 
    int diag_once(const HPsiFunc& hpsi_func,
                  const HPsiFunc& spsi_func,
                  T* psi_in,
                  const int psi_in_dmax,
                  Real* eigenvalue_in,
                  const std::vector<double>& ethr_band);
 
    bool test_exit_cond(const int& ntry, const int& notconv, const bool& scf);
 
    int diag_subspace; // 0: LAPACK, 1: Gen-ELPA, 2: ScaLAPACK
    int diago_subspace_bs = 0; // the block size in 2d block cyclic distribution if use elpa or scalapack
 
#ifdef __DSP
    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 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 setmem_complex_op = base_device::memory::set_memory_op<T, Device>;
 
#ifdef __DSP
    using resmem_real_op = base_device::memory::resize_memory_op_mt<Real, Device>;
    using delmem_real_op = base_device::memory::delete_memory_op_mt<Real, Device>;
#else
    using resmem_real_op = base_device::memory::resize_memory_op<Real, Device>;
    using delmem_real_op = base_device::memory::delete_memory_op<Real, Device>;
#endif
    using setmem_real_op = base_device::memory::set_memory_op<Real, Device>;
    using setmem_complex_2d_op = base_device::memory::set_memory_2d_op<T, Device>;
 
    using resmem_real_h_op = base_device::memory::resize_memory_op<Real, base_device::DEVICE_CPU>;
    using delmem_real_h_op = base_device::memory::delete_memory_op<Real, base_device::DEVICE_CPU>;
    using setmem_real_h_op = base_device::memory::set_memory_op<Real, base_device::DEVICE_CPU>;
 
    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 syncmem_complex_op = base_device::memory::synchronize_memory_op<T, Device, Device>;
    using syncmem_complex_2d_op = base_device::memory::synchronize_memory_2d_op<T, Device, Device>;
    using castmem_complex_op = base_device::memory::cast_memory_op<std::complex<double>, T, Device, Device>;
    using syncmem_h2d_op = base_device::memory::synchronize_memory_op<T, Device, base_device::DEVICE_CPU>;
    using syncmem_d2h_op = base_device::memory::synchronize_memory_op<T, base_device::DEVICE_CPU, Device>;
 
    // Note that ct_Device is different from base_device!
    using ct_Device = typename ct::PsiToContainer<Device>::type;
    // using hegvd_op = container::kernels::lapack_hegvd<T, ct_Device>;
 
    const T *one = nullptr, *zero = nullptr, *neg_one = nullptr;
    const T one_ = static_cast<T>(1.0), zero_ = static_cast<T>(0.0), neg_one_ = static_cast<T>(-1.0);
};
 
} // namespace hsolver
 
#endif