RT-TDDFT GPU Acceleration: RT-TD now fully support GPU computation

source/module_base/lapack_connector.h

@@ -133,8 +133,8 @@ extern "C"
 
     // zgetrf computes the LU factorization of a general matrix
     // while zgetri takes its output to perform matrix inversion
-    void zgetrf_(const int* m, const int *n, const std::complex<double> *A, const int *lda, int *ipiv, const int* info);
-    void zgetri_(const int* n, std::complex<double> *A, const int *lda, int *ipiv, std::complex<double> *work, int *lwork, const int *info);
+    void zgetrf_(const int* m, const int *n, std::complex<double> *A, const int *lda, int *ipiv, int* info);
+    void zgetri_(const int* n, std::complex<double>* A, const int* lda, const int* ipiv, std::complex<double>* work, const int* lwork, int* info);
 
     // if trans=='N':	C = alpha * A * A.H + beta * C
 	// if trans=='C':	C = alpha * A.H * A + beta * C

source/module_base/module_container/ATen/kernels/lapack.cpp

@@ -124,6 +124,41 @@ struct lapack_dngvd<T, DEVICE_CPU> {
     }
 };
 
+template <typename T>
+struct lapack_getrf<T, DEVICE_CPU> {
+    void operator()(
+        const int& m,
+        const int& n,
+        T* Mat,
+        const int& lda,
+        int* ipiv,
+        int& info)
+    {
+        lapackConnector::getrf(m, n, Mat, lda, ipiv, info);
+        if (info != 0) {
+            throw std::runtime_error("getrf failed with info = " + std::to_string(info));
+        }
+    }
+};
+
+template <typename T>
+struct lapack_getri<T, DEVICE_CPU> {
+    void operator()(
+        const int& n,
+        T* Mat,
+        const int& lda,
+        const int* ipiv,
+        T* work,
+        const int& lwork,
+        int& info)
+    {
+        lapackConnector::getri(n, Mat, lda, ipiv, work, lwork, info);
+        if (info != 0) {
+            throw std::runtime_error("getri failed with info = " + std::to_string(info));
+        }
+    }
+};
+
 template struct set_matrix<float,  DEVICE_CPU>;
 template struct set_matrix<double, DEVICE_CPU>;
 template struct set_matrix<std::complex<float>,  DEVICE_CPU>;
@@ -149,5 +184,15 @@ template struct lapack_dngvd<double, DEVICE_CPU>;
 template struct lapack_dngvd<std::complex<float>,  DEVICE_CPU>;
 template struct lapack_dngvd<std::complex<double>, DEVICE_CPU>;
 
+template struct lapack_getrf<float,  DEVICE_CPU>;
+template struct lapack_getrf<double, DEVICE_CPU>;
+template struct lapack_getrf<std::complex<float>,  DEVICE_CPU>;
+template struct lapack_getrf<std::complex<double>, DEVICE_CPU>;
+
+template struct lapack_getri<float, DEVICE_CPU>;
+template struct lapack_getri<double, DEVICE_CPU>;
+template struct lapack_getri<std::complex<float>, DEVICE_CPU>;
+template struct lapack_getri<std::complex<double>, DEVICE_CPU>;
+
 } // namespace kernels
 } // namespace container

source/module_base/module_container/ATen/kernels/lapack.h

@@ -65,6 +65,32 @@ struct lapack_dngvd {
         Real* eigen_val);
 };
 
+
+template <typename T, typename Device>
+struct lapack_getrf {
+    void operator()(
+        const int& m,
+        const int& n,
+        T* Mat,
+        const int& lda,
+        int* ipiv,
+        int& info);
+};
+
+
+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,
+        int& info);
+};
+
+
 #if defined(__CUDA) || defined(__ROCM)
 // TODO: Use C++ singleton to manage the GPU handles
 void createGpuSolverHandle();  // create cusolver handle

source/module_base/module_container/base/third_party/lapack.h

@@ -105,6 +105,15 @@ void dtrtri_(const char* uplo, const char* diag, const int* n, double* a, const
 void ctrtri_(const char* uplo, const char* diag, const int* n, std::complex<float>* a, const int* lda, int* info);
 void ztrtri_(const char* uplo, const char* diag, const int* n, std::complex<double>* a, const int* lda, int* info);
 
+void sgetrf_(const int* m, const int* n, float* a, const int* lda, int* ipiv, int* info);
+void dgetrf_(const int* m, const int* n, double* a, const int* lda, int* ipiv, int* info);
+void cgetrf_(const int* m, const int* n, std::complex<float>* a, const int* lda, int* ipiv, int* info);
+void zgetrf_(const int* m, const int* n, std::complex<double>* a, const int* lda, int* ipiv, int* info);
+
+void sgetri_(const int* n, float* A, const int* lda, const int* ipiv, float* work, const int* lwork, int* info);
+void dgetri_(const int* n, double* A, const int* lda, const int* ipiv, double* work, const int* lwork, int* info);
+void cgetri_(const int* n, std::complex<float>* A, const int* lda, const int* ipiv, std::complex<float>* work, const int* lwork, int* info);
+void zgetri_(const int* n, std::complex<double>* A, const int* lda, const int* ipiv, std::complex<double>* work, const int* lwork, int* info);
 }
 
 // Class LapackConnector provide the connector to fortran lapack routine.
@@ -321,6 +330,48 @@ void trtri( const char &uplo, const char &diag, const int &n, std::complex<doubl
     ztrtri_( &uplo, &diag, &n, A, &lda, &info);
 }
 
+static inline
+void getrf(const int &m, const int &n, float* A, const int &lda, int* ipiv, int &info)
+{
+    sgetrf_(&m, &n, A, &lda, ipiv, &info);
+}
+static inline
+void getrf(const int &m, const int &n, double* A, const int &lda, int* ipiv, int &info)
+{
+    dgetrf_(&m, &n, A, &lda, ipiv, &info);
+}
+static inline
+void getrf(const int &m, const int &n, std::complex<float>* A, const int &lda, int* ipiv, int &info)
+{
+    cgetrf_(&m, &n, A, &lda, ipiv, &info);
+}
+static inline
+void getrf(const int &m, const int &n, std::complex<double>* A, const int &lda, int* ipiv, int &info)
+{
+    zgetrf_(&m, &n, A, &lda, ipiv, &info);
+}
+
+static inline
+void getri(const int& n, float* A, const int& lda, const int* ipiv, float* work, const int& lwork, int& info)
+{
+    sgetri_(&n, A, &lda, ipiv, work, &lwork, &info);
+}
+static inline
+void getri(const int& n, double* A, const int& lda, const int* ipiv, double* work, const int& lwork, int& info)
+{
+    dgetri_(&n, A, &lda, ipiv, work, &lwork, &info);
+}
+static inline
+void getri(const int& n, std::complex<float>* A, const int& lda, const int* ipiv, std::complex<float>* work, const int& lwork, int& info)
+{
+    cgetri_(&n, A, &lda, ipiv, work, &lwork, &info);
+}
+static inline
+void getri(const int& n, std::complex<double>* A, const int& lda, const int* ipiv, std::complex<double>* work, const int& lwork, int& info)
+{
+    zgetri_(&n, A, &lda, ipiv, work, &lwork, &info);
+}
+
 } // namespace lapackConnector
 } // namespace container