math_kernel_op.h

Go to the documentation of this file.

// TODO: This is a temperary location for these functions.
// And will be moved to a global module(module base) later.
#ifndef MODULE_HSOLVER_MATH_KERNEL_H
#define MODULE_HSOLVER_MATH_KERNEL_H
 
#include "source_base/macros.h"
 
#include "source_base/parallel_reduce.h"
 
#include "source_base/module_device/memory_op.h"
#include "source_base/module_device/types.h"
 
#if defined(__CUDA) || defined(__UT_USE_CUDA)
#include <cuda_runtime.h>
 
#include "cublas_v2.h"
#endif //__CUDA || __UT_USE_CUDA
 
namespace ModuleBase {
 
//---------------------------------------------------------------------------------
//-----------------------------0. Tool Functions-----------------------------------
//---------------------------------------------------------------------------------
inline std::complex<double> set_real_tocomplex(const std::complex<double> &x) {
  return {x.real(), 0.0};
}
 
inline std::complex<float> set_real_tocomplex(const std::complex<float> &x) {
  return {x.real(), 0.0};
}
 
inline double set_real_tocomplex(const double &x) { return x; }
 
inline float set_real_tocomplex(const float &x) { return x; }
 
inline std::complex<double> get_conj(const std::complex<double> &x) {
  return {x.real(), -x.imag()};
}
 
inline std::complex<float> get_conj(const std::complex<float> &x) {
  return {x.real(), -x.imag()};
}
 
inline double get_conj(const double &x) { return x; }
 
inline float get_conj(const float &x) { return x; }
 
 
//---------------------------------------------------------------------------------
//-----------------------------1. Vector Operations--------------------------------
//---------------------------------------------------------------------------------
template <typename FPTYPE, typename Device> struct scal_op {
  void operator()(const int &N,
                  const std::complex<FPTYPE> *alpha, std::complex<FPTYPE> *X,
                  const int &incx);
};
 
template <typename T, typename Device> struct vector_mul_real_op {
  using Real = typename GetTypeReal<T>::type;
  void operator()(const int dim, T* result, const T* vector, const Real constant);
};
 
// vector operator: result[i] = vector1[i](complex) * vector2[i](not complex)
template <typename T, typename Device> struct vector_mul_vector_op {
  using Real = typename GetTypeReal<T>::type;
  void operator()(const int& dim, T* result, const T* vector1, const Real* vector2, const bool& add = false);
};
 
// vector operator: result[i] = vector[i] / constant
template <typename T, typename Device> struct vector_div_constant_op {
  using Real = typename GetTypeReal<T>::type;
  void operator()(const int& dim, T* result, const T* vector, const Real constant);
};
 
// vector operator: result[i] = vector1[i](complex) / vector2[i](not complex)
template <typename T, typename Device> struct vector_div_vector_op {
  using Real = typename GetTypeReal<T>::type;
  void operator()(const int &dim, T *result, const T *vector1,
                  const Real *vector2);
};
 
//  compute Y = alpha * X + Y
template <typename T, typename Device> struct axpy_op {
  void operator()(const int &N, const T *alpha, const T *X,
                  const int &incX, T *Y, const int &incY);
};
 
// vector operator: result[i] = vector1[i] * constant1 + vector2[i] * constant2
template <typename T, typename Device>
struct vector_add_vector_op {
  using Real = typename GetTypeReal<T>::type;
  void operator()(const int &dim, T *result, const T *vector1,
                  const Real constant1, const T *vector2, const Real constant2);
};
 
template <typename T, typename Device> struct dot_real_op {
  using Real = typename GetTypeReal<T>::type;
  Real operator()(const int &dim, const T *psi_L,
                  const T *psi_R, const bool reduce = true);
};
 
 
//---------------------------------------------------------------------------------
//-----------------------------2. Matrix Operations--------------------------------
//---------------------------------------------------------------------------------
 
// compute y = alpha * op(A) * x + beta * y
template <typename T, typename Device> struct gemv_op {
  void operator()(const char &trans, const int &m,
                  const int &n, const T *alpha, const T *A, const int &lda,
                  const T *X, const int &incx, const T *beta, T *Y,
                  const int &incy);
};
 
// compute C = alpha * op(A) * op(B) + beta * C
template <typename T, typename Device> struct gemm_op {
  void operator()(const char &transa, const char &transb,
                  const int &m, const int &n, const int &k, const T *alpha,
                  const T *a, const int &lda, const T *b, const int &ldb,
                  const T *beta, T *c, const int &ldc);
};
 
#ifdef __DSP
// compute C = alpha * op(A) * op(B) + beta * C on DSP Hardware
template <typename T, typename Device> struct gemm_op_mt {
  void operator()(const char &transa, const char &transb,
                  const int &m, const int &n, const int &k, const T *alpha,
                  const T *a, const int &lda, const T *b, const int &ldb,
                  const T *beta, T *c, const int &ldc);
};
#endif
 
template <typename T, typename Device> struct matrixTranspose_op {
  void operator()(const int &row, const int &col,
                  const T *input_matrix, T *output_matrix);
};
 
template <typename T, typename Device> struct matrixCopy {
  void operator()(const int& n1, const int& n2, const T* A, const int& LDA, T* B, const int& LDB);
};
 
template <typename T, typename Device>
struct matrix_mul_vector_op {
    using Real = typename GetTypeReal<T>::type;
  void operator()(const int &m, const int &n,
                  T *a,
                  const int &lda,
                  const Real *b,
                  const Real alpha,
                  T *c,
                  const int &ldc);
};
 
template <typename T, typename Device>
struct apply_eigenvalues_op {
    using Real = typename GetTypeReal<T>::type;
 
    void operator()(const Device *d, const int &nbase, const int &nbase_x, const int &notconv,
                    T *result, const T *vectors, const Real *eigenvalues);
};
 
template <typename T, typename Device>
struct precondition_op {
    using Real = typename GetTypeReal<T>::type;
    void operator()(const Device* d,
                   const int& dim,
                   T* psi_iter,
                   const int& nbase,
                   const int& notconv,
                   const Real* precondition,
                   const Real* eigenvalues);
};
 
template <typename T, typename Device>
struct normalize_op {
    using Real = typename GetTypeReal<T>::type;
    void operator()(const Device* d,
                   const int& dim,
                   T* psi_iter,
                   const int& nbase,
                   const int& notconv,
                   Real* psi_norm = nullptr);
};
 
template <typename T>
struct normalize_op<T, base_device::DEVICE_GPU> {
    using Real = typename GetTypeReal<T>::type;
    void operator()(const base_device::DEVICE_GPU* d,
                   const int& dim,
                   T* psi_iter,
                   const int& nbase,
                   const int& notconv,
                   Real* psi_norm);
};
 
#if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
// Partially specialize functor for base_device::GpuDevice.
template <typename T> struct dot_real_op<T, base_device::DEVICE_GPU> {
  using Real = typename GetTypeReal<T>::type;
  Real operator()(const int &dim,
                  const T *psi_L, const T *psi_R, const bool reduce = true);
};
 
// vector operator: result[i] = vector[i] / constant
template <typename T>
struct vector_mul_real_op<T, base_device::DEVICE_GPU>
{
  using Real = typename GetTypeReal<T>::type;
  void operator()(const int dim, T* result, const T* vector, const Real constant);
};
 
// vector operator: result[i] = vector1[i](complex) * vector2[i](not complex)
template <typename T> struct vector_mul_vector_op<T, base_device::DEVICE_GPU> {
  using Real = typename GetTypeReal<T>::type;
  void operator()(const int& dim, T* result, const T* vector1, const Real* vector2, const bool& add = false);
};
 
// vector operator: result[i] = vector[i] / constant
template <typename T> struct vector_div_constant_op<T, base_device::DEVICE_GPU> {
  using Real = typename GetTypeReal<T>::type;
  void operator()(const int& dim, T* result, const T* vector, const Real constant);
};
 
// vector operator: result[i] = vector1[i](complex) / vector2[i](not complex)
template <typename T> struct vector_div_vector_op<T, base_device::DEVICE_GPU> {
  using Real = typename GetTypeReal<T>::type;
  void operator()(const int &dim, T *result,
                  const T *vector1, const Real *vector2);
};
 
// vector operator: result[i] = vector1[i] * constant1 + vector2[i] * constant2
template <typename T>
struct vector_add_vector_op<T, base_device::DEVICE_GPU> {
  using Real = typename GetTypeReal<T>::type;
  void operator()(const int &dim, T *result,
                  const T *vector1, const Real constant1, const T *vector2,
                  const Real constant2);
};
 
template <typename T> struct matrixCopy<T, base_device::DEVICE_GPU> {
    void operator()(const int& n1,
                    const int& n2,
                    const T* A, // input
                    const int& LDA,
                    T* B, // output
                    const int& LDB);
};
 
template <typename T> struct matrix_mul_vector_op<T, base_device::DEVICE_GPU> {
  using Real = typename GetTypeReal<T>::type;
  void operator()(const int &m, const int &n,
                  T *a,
                  const int &lda,
                  const Real *b,
                  const Real alpha,
                  T *c,
                  const int &ldc);
};
 
void createGpuBlasHandle();
void destoryBLAShandle();
 
// vector operator: result[i] = -lambda[i] * vector[i]
template <typename T> struct apply_eigenvalues_op<T, base_device::DEVICE_GPU> {
    using Real = typename GetTypeReal<T>::type;
 
    void operator()(const base_device::DEVICE_GPU *d, const int &nbase, const int &nbase_x, const int &notconv,
                    T *result, const T *vectors, const Real *eigenvalues);
};
 
template <typename T>
struct precondition_op<T, base_device::DEVICE_GPU> {
    using Real = typename GetTypeReal<T>::type;
    void operator()(const base_device::DEVICE_GPU* d,
                   const int& dim,
                   T* psi_iter,
                   const int& nbase,
                   const int& notconv,
                   const Real* precondition,
                   const Real* eigenvalues);
};
 
#endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
} // namespace hsolver
 
#endif // MODULE_HSOLVER_MATH_KERNEL_H