sto_che.h

Go to the documentation of this file.

#ifndef STO_CHE_H
#define STO_CHE_H
#include "source_base/math_chebyshev.h"
#include "source_base/kernels/math_kernel_op.h"
#include "source_base/module_container/ATen/kernels/blas.h"
 
template <typename REAL, typename Device = base_device::DEVICE_CPU>
class StoChe
{
  public:
    StoChe(const int& nche, const int& method, const REAL& emax_sto, const REAL& emin_sto);
    ~StoChe();
 
  public:
    int nche = 0;               
    REAL* spolyv = nullptr;     
    REAL* spolyv_cpu = nullptr; 
    int method_sto = 0;         
 
    // Chebyshev expansion
    // It stores the plan of FFTW and should be initialized at the beginning of the calculation
    ModuleBase::Chebyshev<REAL, Device>* p_che = nullptr;
 
    REAL emax_sto = 0.0; 
    REAL emin_sto = 0.0; 
 
  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 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>;
};
 
template <typename REAL, typename Device>
REAL vTMv(const REAL* v, const REAL* M, const int n)
{
    Device* ctx = {};
    base_device::DEVICE_CPU* cpu_ctx = {};
    using ct_Device = typename container::PsiToContainer<Device>::type;
    const char normal = 'N';
    const REAL one = 1;
    const int inc = 1;
    const REAL zero = 0;
    REAL* y = nullptr;
    base_device::memory::resize_memory_op<REAL, Device>()(y, n);
    ModuleBase::gemv_op<REAL, Device>()(normal, n, n, &one, M, n, v, inc, &zero, y, inc);
    REAL result = 0;
    REAL* dot_device = nullptr;
    base_device::memory::resize_memory_op<REAL, Device>()(dot_device, 1);
    container::kernels::blas_dot<REAL, ct_Device>()(n, y, 1, v, 1, dot_device);
    base_device::memory::synchronize_memory_op<REAL, base_device::DEVICE_CPU, Device>()(&result,
                                                                                        dot_device,
                                                                                        1);
    base_device::memory::delete_memory_op<REAL, Device>()(y);
    base_device::memory::delete_memory_op<REAL, Device>()(dot_device);
    return result;
}
 
#endif