lapack.h

Go to the documentation of this file.

#ifndef ATEN_KERNELS_LAPACK_H_
#define ATEN_KERNELS_LAPACK_H_
 
#include "source_base/macros.h"
#include <ATen/core/tensor.h>
#include <ATen/core/tensor_types.h>
 
#include <base/third_party/lapack.h>
 
namespace container {
namespace kernels {
 
 
template <typename T, typename Device>
struct set_matrix {
    void operator() (
        const char& uplo,
        T* A,
        const int& dim);
};
 
 
// --- 1. Matrix Decomposition ---
template <typename T, typename Device>
struct lapack_trtri {
    void operator()(
        const char& uplo,
        const char& diag,
        const int& dim,
        T* Mat,
        const int& lda);
};
 
 
template <typename T, typename Device>
struct lapack_potrf {
    void operator()(
        const char& uplo,
        const int& dim,
        T* Mat,
        const int& lda);
};
 
template <typename T, typename Device>
struct lapack_getrf {
    void operator()(
        const int& m,
        const int& n,
        T* Mat,
        const int& lda,
        int* ipiv);
};
 
 
template <typename T, typename Device>
struct lapack_getri {
    void operator()(
        const int& n,
        T* Mat,
        const int& lda,
        const int* ipiv,
        T* work,
        const int& lwork);
};
 
// This is QR factorization in-place
// that will change input Mat A to orthogonal/unitary matrix Q
template <typename T, typename Device>
struct lapack_geqrf_inplace {
    void operator()(
        const int m,
        const int n,
        T *A,
        const int lda);
};
 
// This is QR factorization
// where [in]Mat will be kept and the results are stored in separate matrix Q
// template <typename T, typename Device>
// struct lapack_geqrf{
//     /**
//      * Perform QR factorization of a matrix using LAPACK's geqrf function.
//      *
//      * @param m The number of rows in the matrix.
//      * @param n The number of columns in the matrix.
//      * @param Mat The matrix to be factorized.
//      *        On exit, the upper triangle contains the upper triangular matrix R,
//      *        and the elements below the diagonal, with the array TAU, represent
//      *        the unitary matrix Q as a product of min(m,n) elementary reflectors.
//      * @param lda The leading dimension of the matrix.
//      * @param tau Array of size min(m,n) containing the Householder reflectors.
//      */
//     void operator()(
//         const int m,
//         const int n,
//         T *Mat,
//         const int lda,
//         T *tau);
// };
 
 
// --- 2. Linear System Solvers ---
template <typename T, typename Device>
struct lapack_getrs {
    void operator()(
        const char& trans,
        const int& n,
        const int& nrhs,
        T* A,
        const int& lda,
        const int* ipiv,
        T* B,
        const int& ldb);
};
 
 
 
// --- 3. Standard & Generalized Eigenvalue ---
 
// ============================================================================
// Standard Hermitian Eigenvalue Problem Solvers
// ============================================================================
// The following structures (lapack_heevd and lapack_heevx) implement solvers
// for standard Hermitian eigenvalue problems of the form:
//      A * x = lambda * x
// where:
//   - A is a Hermitian matrix
//   - lambda are the eigenvalues to be computed
//   - x are the corresponding eigenvectors
//
// ============================================================================
template <typename T, typename Device>
struct lapack_heevd {
    // !> ZHEEVD computes all eigenvalues and, optionally, eigenvectors of a
    // !> complex Hermitian matrix A.  If eigenvectors are desired, it uses a
    // !> divide and conquer algorithm.
    // !>          On exit, if JOBZ = 'V', then if INFO = 0, A contains the
    // !>          orthonormal eigenvectors of the matrix A.
    using Real = typename GetTypeReal<T>::type;
    void operator()(
        const int dim,
        T* Mat,
        const int lda,
        Real* eigen_val);
};
 
template <typename T, typename Device>
struct lapack_heevx {
    using Real = typename GetTypeReal<T>::type;
    void operator()(
        const int dim,
        const int lda,
        const T *Mat,
        const int neig,
        Real *eigen_val,
        T *eigen_vec);
};
 
 
// ============================================================================
// Generalized Hermitian-definite Eigenvalue Problem Solvers
// ============================================================================
// The following structures (lapack_hegvd and lapack_hegvx) implement solvers
// for generalized Hermitian-definite eigenvalue problems of the form:
//      A * x = lambda * B * x
// where:
//   - A is a Hermitian matrix
//   - B is a Hermitian positive definite matrix
//   - lambda are the eigenvalues to be computed
//   - x are the corresponding eigenvectors
//
// ============================================================================
 
template <typename T, typename Device>
struct lapack_hegvd {
    using Real = typename GetTypeReal<T>::type;
    void operator()(
        const int n,
        const int lda,
        T *Mat_A,
        T *Mat_B,
        Real *eigen_val,
        T *eigen_vec);
};
 
template <typename T, typename Device>
struct lapack_hegvx {
    using Real = typename GetTypeReal<T>::type;
    void operator()(
        const int n,
        const int lda,
        T *Mat_A,
        T *Mat_B,
        const int m,
        Real *eigen_val,
        T *eigen_vec);
};
 
 
#if defined(__CUDA) || defined(__ROCM)
// TODO: Use C++ singleton to manage the GPU handles
void createGpuSolverHandle();  // create cusolver handle
void destroyGpuSolverHandle(); // destroy cusolver handle
#endif
 
} // namespace container
} // namespace kernels
 
#endif // ATEN_KERNELS_LAPACK_H_