stress_op.h

Go to the documentation of this file.

#ifndef SRC_PW_STRESS_MULTI_DEVICE_H
#define SRC_PW_STRESS_MULTI_DEVICE_H
#include "source_io/module_parameter/parameter.h"
 
#include "source_psi/psi.h"
 
#include <complex>
#include <source_base/macros.h>
 
namespace hamilt
{
 
template <typename FPTYPE, typename Device>
struct cal_dbecp_noevc_nl_op
{
    void operator()(const Device* ctx,
                    const int& ipol,
                    const int& jpol,
                    const int& nkb,
                    const int& npw,
                    const int& npwx,
                    const int& ik,
                    const FPTYPE& tpiba,
                    const FPTYPE* gcar,
                    const FPTYPE* kvec_c,
                    std::complex<FPTYPE>* vkbi,
                    std::complex<FPTYPE>* vkbj,
                    std::complex<FPTYPE>* vkb,
                    std::complex<FPTYPE>* vkb1,
                    std::complex<FPTYPE>* vkb2,
                    std::complex<FPTYPE>* dbecp_noevc);
};
 
template <typename FPTYPE, typename Device>
struct cal_stress_nl_op
{
    void operator()(const Device* ctx,
                    const bool& nondiagonal,
                    const int& ipol,
                    const int& jpol,
                    const int& nkb,
                    const int& nbands_occ,
                    const int& ntype,
                    const int& spin,
                    const int& deeq_2,
                    const int& deeq_3,
                    const int& deeq_4,
                    const int* atom_nh,
                    const int* atom_na,
                    const FPTYPE* d_wg,
                    const bool& occ,
                    const FPTYPE* d_ekb,
                    const FPTYPE* qq_nt,
                    const FPTYPE* deeq,
                    const std::complex<FPTYPE>* becp,
                    const std::complex<FPTYPE>* dbecp,
                    FPTYPE* stress);
    // interface for nspin=4 only
    void operator()(const Device* ctx,
                    const int& ipol,
                    const int& jpol,
                    const int& nkb,
                    const int& nbands_occ,
                    const int& ntype,
                    const int& deeq_2,
                    const int& deeq_3,
                    const int& deeq_4,
                    const int* atom_nh,
                    const int* atom_na,
                    const FPTYPE* d_wg,
                    const bool& occ,
                    const FPTYPE* d_ekb,
                    const FPTYPE* qq_nt,
                    const std::complex<FPTYPE>* deeq_nc,
                    const std::complex<FPTYPE>* becp,
                    const std::complex<FPTYPE>* dbecp,
                    FPTYPE* stress);
    // kernel for DFT+U
    void operator()(const base_device::DEVICE_CPU* ctx,
                    const int& nkb,
                    const int& nbands_occ,
                    const int& ntype,
                    const int& wg_nc,
                    const int& ik,
                    const int* atom_nh,
                    const int* atom_na,
                    const FPTYPE* d_wg,
                    const std::complex<FPTYPE>* vu,
                    const int* orbital_corr,
                    const std::complex<FPTYPE>* becp,
                    const std::complex<FPTYPE>* dbecp,
                    FPTYPE* stress);
    // kernel for DeltaSpin
    void operator()(const base_device::DEVICE_CPU* ctx,
                    const int& nkb,
                    const int& nbands_occ,
                    const int& ntype,
                    const int& wg_nc,
                    const int& ik,
                    const int* atom_nh,
                    const int* atom_na,
                    const FPTYPE* d_wg,
                    const double* lambda,
                    const std::complex<FPTYPE>* becp,
                    const std::complex<FPTYPE>* dbecp,
                    FPTYPE* stress);
};
 
template <typename T, typename Device>
struct cal_stress_mgga_op
{
    using Real = typename GetTypeReal<T>::type;
    void operator()(const int& spin, const int& nrxx, const Real& w1, const T* gradwfc, Real* crosstaus);
};
 
// cpu version first, gpu version later
template <typename FPTYPE, typename Device>
struct cal_vkb_op
{
    void operator()(const Device* ctx,
                    const int nh,
                    const int npw,
                    const int* indexes,
                    const FPTYPE* vqs_in,
                    const FPTYPE* ylms_in,
                    const std::complex<FPTYPE>* sk_in,
                    const std::complex<FPTYPE>* pref_in,
                    std::complex<FPTYPE>* vkbs_out);
};
 
// cpu version first, gpu version later
template <typename FPTYPE, typename Device>
struct cal_vkb_deri_op
{
    void operator()(const Device* ctx,
                    const int nh,
                    const int npw,
                    const int ipol,
                    const int jpol,
                    const int* indexes,
                    const FPTYPE* vqs_in,
                    const FPTYPE* vqs_deri_in,
                    const FPTYPE* ylms_in,
                    const FPTYPE* ylms_deri_in,
                    const std::complex<FPTYPE>* sk_in,
                    const std::complex<FPTYPE>* pref_in,
                    const FPTYPE* gk_in,
                    std::complex<FPTYPE>* vkbs_out);
};
 
// cpu version first, gpu version later
template <typename FPTYPE, typename Device>
struct cal_vq_op
{
    void operator()(const Device* ctx,
                    const FPTYPE* tab,
                    int it,
                    const FPTYPE* gk,
                    int npw,
                    const int tab_2,
                    const int tab_3,
                    const FPTYPE table_interval,
                    const int nbeta,
                    FPTYPE* vq);
};
 
// cpu version first, gpu version later
template <typename FPTYPE, typename Device>
struct cal_vq_deri_op
{
    void operator()(const Device* ctx,
                    const FPTYPE* tab,
                    int it,
                    const FPTYPE* gk,
                    int npw,
                    const int tab_2,
                    const int tab_3,
                    const FPTYPE table_interval,
                    const int nbeta,
                    FPTYPE* vq);
};
 
 
template <typename FPTYPE, typename Device>
struct cal_stress_drhoc_aux_op{
    void operator()(
        const FPTYPE* r, const FPTYPE* rhoc, 
        const FPTYPE *gx_arr, const FPTYPE *rab, FPTYPE *drhocg, 
        const int mesh, const int igl0, const int ngg, const double omega,
        int type
    );
};
 
template <typename FPTYPE, typename Device>
struct cal_force_npw_op{
    void operator()(const std::complex<FPTYPE> *psiv,
                    const FPTYPE* gv_x, const FPTYPE* gv_y, const FPTYPE* gv_z,
                    const FPTYPE* rhocgigg_vec,
                    FPTYPE* force,
                    const FPTYPE pos_x, const FPTYPE pos_y, const FPTYPE pos_xz,
                    const int npw,
                    const FPTYPE omega, const FPTYPE tpiba
    );
};
 
template <typename FPTYPE, typename Device>
struct cal_multi_dot_op{
    FPTYPE operator()(const int& npw,
                    const FPTYPE& fac,
                    const FPTYPE* gk1,
                    const FPTYPE* gk2,
                    const FPTYPE* d_kfac,
                    const std::complex<FPTYPE>* psi);
};
 
 
#if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
template <typename FPTYPE>
struct cal_dbecp_noevc_nl_op<FPTYPE, base_device::DEVICE_GPU>
{
    void operator()(const base_device::DEVICE_GPU* ctx,
                    const int& ipol,
                    const int& jpol,
                    const int& nkb,
                    const int& npw,
                    const int& npwx,
                    const int& ik,
                    const FPTYPE& tpiba,
                    const FPTYPE* gcar,
                    const FPTYPE* kvec_c,
                    std::complex<FPTYPE>* vkbi,
                    std::complex<FPTYPE>* vkbj,
                    std::complex<FPTYPE>* vkb,
                    std::complex<FPTYPE>* vkb1,
                    std::complex<FPTYPE>* vkb2,
                    std::complex<FPTYPE>* dbecp_noevc);
};
 
template <typename FPTYPE>
struct cal_stress_nl_op<FPTYPE, base_device::DEVICE_GPU>
{
    void operator()(const base_device::DEVICE_GPU* ctx,
                    const bool& nondiagonal,
                    const int& ipol,
                    const int& jpol,
                    const int& nkb,
                    const int& nbands_occ,
                    const int& ntype,
                    const int& spin,
                    const int& deeq_2,
                    const int& deeq_3,
                    const int& deeq_4,
                    const int* atom_nh,
                    const int* atom_na,
                    const FPTYPE* d_wg,
                    const bool& occ,
                    const FPTYPE* d_ekb,
                    const FPTYPE* qq_nt,
                    const FPTYPE* deeq,
                    const std::complex<FPTYPE>* becp,
                    const std::complex<FPTYPE>* dbecp,
                    FPTYPE* stress);
    // interface for nspin=4 only
    void operator()(const base_device::DEVICE_GPU* ctx,
                    const int& ipol,
                    const int& jpol,
                    const int& nkb,
                    const int& nbands_occ,
                    const int& ntype,
                    const int& deeq_2,
                    const int& deeq_3,
                    const int& deeq_4,
                    const int* atom_nh,
                    const int* atom_na,
                    const FPTYPE* d_wg,
                    const bool& occ,
                    const FPTYPE* d_ekb,
                    const FPTYPE* qq_nt,
                    const std::complex<FPTYPE>* deeq_nc,
                    const std::complex<FPTYPE>* becp,
                    const std::complex<FPTYPE>* dbecp,
                    FPTYPE* stress);
    // kernel for DFT+U
    void operator()(const base_device::DEVICE_GPU* ctx,
                    const int& nkb,
                    const int& nbands_occ,
                    const int& ntype,
                    const int& wg_nc,
                    const int& ik,
                    const int* atom_nh,
                    const int* atom_na,
                    const FPTYPE* d_wg,
                    const std::complex<FPTYPE>* vu,
                    const int* orbital_corr,
                    const std::complex<FPTYPE>* becp,
                    const std::complex<FPTYPE>* dbecp,
                    FPTYPE* stress);
    // kernel for DeltaSpin
    void operator()(const base_device::DEVICE_GPU* ctx,
                    const int& nkb,
                    const int& nbands_occ,
                    const int& ntype,
                    const int& wg_nc,
                    const int& ik,
                    const int* atom_nh,
                    const int* atom_na,
                    const FPTYPE* d_wg,
                    const double* lambda,
                    const std::complex<FPTYPE>* becp,
                    const std::complex<FPTYPE>* dbecp,
                    FPTYPE* stress);
};
 
// cpu version first, gpu version later
template <typename FPTYPE>
struct cal_vkb_op<FPTYPE, base_device::DEVICE_GPU>
{
    void operator()(const base_device::DEVICE_GPU* ctx,
                    const int nh,
                    const int npw,
                    const int* indexes,
                    const FPTYPE* vqs_in,
                    const FPTYPE* ylms_in,
                    const std::complex<FPTYPE>* sk_in,
                    const std::complex<FPTYPE>* pref_in,
                    std::complex<FPTYPE>* vkbs_out);
};
 
template <typename FPTYPE>
struct cal_vkb_deri_op<FPTYPE, base_device::DEVICE_GPU>
{
    void operator()(const base_device::DEVICE_GPU* ctx,
                    const int nh,
                    const int npw,
                    const int ipol,
                    const int jpol,
                    const int* indexes,
                    const FPTYPE* vqs_in,
                    const FPTYPE* vqs_deri_in,
                    const FPTYPE* ylms_in,
                    const FPTYPE* ylms_deri_in,
                    const std::complex<FPTYPE>* sk_in,
                    const std::complex<FPTYPE>* pref_in,
                    const FPTYPE* gk_in,
                    std::complex<FPTYPE>* vkbs_out);
};
 
// cpu version first, gpu version later
template <typename FPTYPE>
struct cal_vq_op<FPTYPE, base_device::DEVICE_GPU>
{
    void operator()(const base_device::DEVICE_GPU* ctx,
                    const FPTYPE* tab,
                    int it,
                    const FPTYPE* gk,
                    int npw,
                    const int tab_2,
                    const int tab_3,
                    const FPTYPE table_interval,
                    const int nbeta,
                    FPTYPE* vq);
};
 
// cpu version first, gpu version later
template <typename FPTYPE>
struct cal_vq_deri_op<FPTYPE, base_device::DEVICE_GPU>
{
    void operator()(const base_device::DEVICE_GPU* ctx,
                    const FPTYPE* tab,
                    int it,
                    const FPTYPE* gk,
                    int npw,
                    const int tab_2,
                    const int tab_3,
                    const FPTYPE table_interval,
                    const int nbeta,
                    FPTYPE* vq);
};
 
template <typename FPTYPE>
struct cal_multi_dot_op<FPTYPE, base_device::DEVICE_GPU>{
    FPTYPE operator()(const int& npw,
                    const FPTYPE& fac,
                    const FPTYPE* gk1,
                    const FPTYPE* gk2,
                    const FPTYPE* d_kfac,
                    const std::complex<FPTYPE>* psi);
};
 
template <typename FPTYPE>
struct cal_stress_drhoc_aux_op<FPTYPE, base_device::DEVICE_GPU>{
    void operator()(
        const FPTYPE* r, const FPTYPE* rhoc, 
        const FPTYPE *gx_arr, const FPTYPE *rab, FPTYPE *drhocg, 
        const int mesh, const int igl0, const int ngg, const double omega,
        int type
    );
};
 
 
template <typename FPTYPE>
struct cal_force_npw_op<FPTYPE, base_device::DEVICE_GPU>{
    void operator()(const std::complex<FPTYPE> *psiv,
                    const FPTYPE* gv_x, const FPTYPE* gv_y, const FPTYPE* gv_z,
                    const FPTYPE* rhocgigg_vec,
                    FPTYPE* force,
                    const FPTYPE pos_x, const FPTYPE pos_y, const FPTYPE pos_xz,
                    const int npw,
                    const FPTYPE omega, const FPTYPE tpiba
    );
};
 
 
 
#endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
 
template <typename Device>
struct pointer_array_malloc
{
    void operator()(void** ptr, const int n);
};
 
template <typename Device>
struct synchronize_ptrs
{
    void operator()(void** ptr_out, const void** ptr_in, const int size);
};
 
} // namespace hamilt
#endif // SRC_PW_STRESS_MULTI_DEVICE_H