diago_bpcg.h

Go to the documentation of this file.

#ifndef DIAGO_BPCG_H_
#define DIAGO_BPCG_H_
 
#include "source_base/kernels/math_kernel_op.h"
#include "source_base/module_device/memory_op.h"
#include "source_base/module_device/types.h"
#include "source_base/para_gemm.h"
#include "source_hamilt/hamilt.h"
#include "source_pw/module_pwdft/structure_factor.h"
#include "source_hsolver/kernels/dngvd_op.h"
#include "source_hsolver/para_linear_transform.h"
 
#include <ATen/core/tensor.h>
#include <ATen/core/tensor_map.h>
#include <source_base/macros.h>
 
namespace hsolver {
 
template <typename T = std::complex<double>, typename Device = base_device::DEVICE_CPU>
class DiagoBPCG
{
  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;
  // Column major psi in this class
  public:
    explicit DiagoBPCG(const Real* precondition);
 
    ~DiagoBPCG();
 
    void init_iter(const int nband, const int nband_l, const int nbasis, const int ndim);
 
    using HPsiFunc = std::function<void(T*, T*, const int, const int)>;
 
    void diag(const HPsiFunc& hpsi_func,
              T* psi_in,
              Real* eigenvalue_in,
              const std::vector<double>& ethr_band);
 
  private:
    int n_band = 0;
    int n_band_l = 0;
    int n_basis = 0;
    int n_dim = 0;
    int nline = 4;
    ModuleBase::PGemmCN<T, Device> pmmcn;
    PLinearTransform<T, Device> plintrans;
 
 
    ct::DataType r_type  = ct::DataType::DT_INVALID;
    ct::DataType t_type  = ct::DataType::DT_INVALID;
    ct::DeviceType device_type = ct::DeviceType::UnKnown;
 
    ct::Tensor prec = {}, h_prec = {};
 
    ct::Tensor beta = {};
    ct::Tensor err_st = {};
    ct::Tensor eigen = {};
 
    ct::Tensor psi = {}, hpsi = {};
    
    ct::Tensor hsub = {};
 
    ct::Tensor grad = {}, hgrad = {}, grad_old = {};
 
    ct::Tensor work = {};
 
    // These are for hsolver gemm_op use
    Device * ctx = {};
    // Pointer to objects of 1 and 0 for gemm 
    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);
 
    void calc_prec();
 
    void calc_hpsi_with_block(
        const HPsiFunc& hpsi_func,
        T *psi_in, 
        ct::Tensor& hpsi_out);
 
    void diag_hsub(
        const ct::Tensor& psi_in, 
        const ct::Tensor& hpsi_in,
        ct::Tensor& hsub_out, 
        ct::Tensor& eigenvalue_out);
 
    void rotate_wf(
        const ct::Tensor& hsub_in,
        ct::Tensor& psi_out, 
        ct::Tensor& workspace_in);
 
    void calc_grad_with_block(
        const ct::Tensor& prec_in, 
        ct::Tensor& err_out, 
        ct::Tensor& beta_out,
        ct::Tensor& psi_in, ct::Tensor& hpsi_in,
        ct::Tensor& grad_out, ct::Tensor& grad_old_out);
 
    void calc_hsub_with_block(
        const HPsiFunc& hpsi_func,
        T *psi_in,
        ct::Tensor& psi_out, ct::Tensor& hpsi_out,
        ct::Tensor& hsub_out, ct::Tensor& workspace_in,
        ct::Tensor& eigenvalue_out);
    
    void calc_hsub_with_block_exit(
        ct::Tensor& psi_out, 
        ct::Tensor& hpsi_out,
        ct::Tensor& hsub_out, 
        ct::Tensor& workspace_in,
        ct::Tensor& eigenvalue_out);
 
    void orth_projection(
        const ct::Tensor& psi_in,
        ct::Tensor& hsub_in,
        ct::Tensor& grad_out);
 
    void line_minimize(
        ct::Tensor& grad_in,
        ct::Tensor& hgrad_in,
        ct::Tensor& psi_out,
        ct::Tensor& hpsi_out);
 
    void orth_cholesky(
        ct::Tensor& workspace_in, 
        ct::Tensor& psi_out, 
        ct::Tensor& hpsi_out, 
        ct::Tensor& hsub_out);
 
    bool test_error(const ct::Tensor& err_in, const std::vector<double>& ethr_band);
 
    using ct_Device = typename ct::PsiToContainer<Device>::type;
    using setmem_var_op = ct::kernels::set_memory<Real, ct_Device>;
    using resmem_var_op = ct::kernels::resize_memory<Real, ct_Device>;
    using delmem_var_op = ct::kernels::delete_memory<Real, ct_Device>;
    using syncmem_var_h2d_op = ct::kernels::synchronize_memory<Real, ct_Device, ct::DEVICE_CPU>;
    using syncmem_var_d2h_op = ct::kernels::synchronize_memory<Real, ct::DEVICE_CPU, ct_Device>;
 
    using setmem_complex_op = ct::kernels::set_memory<T, ct_Device>;
    using delmem_complex_op = ct::kernels::delete_memory<T, ct_Device>;
    using resmem_complex_op = ct::kernels::resize_memory<T, ct_Device>;
    using syncmem_complex_op = ct::kernels::synchronize_memory<T, ct_Device, ct_Device>;
 
    // note: these operators use template parameter base_device::Device_*
    // defined in source_base/module_device/types.h
    // different from ct_Device!
    using gemm_op = ModuleBase::gemm_op<T, Device>;
 
};
 
} // namespace hsolver
#endif // DIAGO_BPCG_H_