Add initializer and GPU tests

Author: Critsium-xy

source/module_base/blas_connector.h

@@ -327,6 +327,15 @@ class BlasConnector
 	void vector_add_vector(const int& dim, std::complex<double> *result, const std::complex<double> *vector1, const double constant1, const std::complex<double> *vector2, const double constant2, base_device::AbacusDevice_t device_type = base_device::AbacusDevice_t::CpuDevice);
 };
 
+#ifdef __CUDA
+
+namespace BlasUtils{
+	void createGpuBlasHandle();
+	void destoryBLAShandle();
+}
+
+#endif
+
 // If GATHER_INFO is defined, the original function is replaced with a "i" suffix,
 // preventing changes on the original code.
 // The real function call is at gather_math_lib_info.cpp

source/module_base/test/blas_connector_test.cpp

@@ -93,6 +93,33 @@ TEST(blas_connector, Scal) {
     }
 }
 
+#ifdef __CUDA
+
+TEST(blas_connector, ScalGpu) {
+    const int size = 8;
+    const std::complex<double> scale = {2, 3};
+    const int incx = 1;
+    std::complex<double> result[8], answer[8];
+    std::complex<double>* result_gpu = nullptr;
+    resmem_zd_op()(gpu_ctx, result_gpu, 8 * sizeof(std::complex<double>));
+    for (int i=0; i< size; i++) {
+        result[i] = std::complex<double>{static_cast<double>(std::rand() / double(RAND_MAX)),
+                 static_cast<double>(std::rand() / double(RAND_MAX))};
+    };
+    for (int i = 0; i < size; i++)
+        answer[i] = result[i] * scale;
+    syncmem_z2z_h2d_op()(gpu_ctx, cpu_ctx, result_gpu, result, sizeof(std::complex<double>) * 8);
+    BlasConnector::scal(size,scale,result_gpu,incx,base_device::AbacusDevice_t::GpuDevice);
+    syncmem_z2z_d2h_op()(cpu_ctx, gpu_ctx, result, result_gpu, sizeof(std::complex<double>) * 8);
+    delmem_zd_op()(gpu_ctx, result_gpu);
+    // incx is the spacing between elements if result
+    for (int i = 0; i < size; i++) {
+        EXPECT_DOUBLE_EQ(answer[i].real(), result[i].real());
+        EXPECT_DOUBLE_EQ(answer[i].imag(), result[i].imag());
+    }
+}
+
+#endif
 
 TEST(blas_connector, daxpy_) {
     typedef double T;
@@ -136,6 +163,67 @@ TEST(blas_connector, zaxpy_) {
     }
 }
 
+TEST(blas_connector, Axpy) {
+    typedef std::complex<double> T;
+    const int size = 8;
+    const T scale = {2, 3};
+    const int incx = 1;
+    const int incy = 1;
+    std::array<T, size> x_const, result, answer;
+    std::generate(x_const.begin(), x_const.end(), []() {
+        return T{static_cast<double>(std::rand() / double(RAND_MAX)),
+                 static_cast<double>(std::rand() / double(RAND_MAX))};
+    });
+    std::generate(result.begin(), result.end(), []() {
+        return T{static_cast<double>(std::rand() / double(RAND_MAX)),
+                 static_cast<double>(std::rand() / double(RAND_MAX))};
+    });
+    for (int i = 0; i < size; i++)
+        answer[i] = x_const[i] * scale + result[i];
+    BlasConnector::axpy(size, scale, x_const.data(), incx, result.data(), incy);
+    for (int i = 0; i < size; i++) {
+        EXPECT_DOUBLE_EQ(answer[i].real(), result[i].real());
+        EXPECT_DOUBLE_EQ(answer[i].imag(), result[i].imag());
+    }
+}
+
+#ifdef __CUDA
+
+TEST(blas_connector, AxpyGpu) {
+    typedef std::complex<double> T;
+    const int size = 8;
+    const T scale = {2, 3};
+    const int incx = 1;
+    const int incy = 1;
+    std::array<T, size> x_const, result, answer;
+    T* x_gpu = nullptr;
+    T* result_gpu = nullptr;
+    resmem_zd_op()(gpu_ctx, x_gpu, size * sizeof(std::complex<double>));
+    resmem_zd_op()(gpu_ctx, result_gpu, size * sizeof(std::complex<double>));
+    std::generate(x_const.begin(), x_const.end(), []() {
+        return T{static_cast<double>(std::rand() / double(RAND_MAX)),
+                 static_cast<double>(std::rand() / double(RAND_MAX))};
+    });
+    std::generate(result.begin(), result.end(), []() {
+        return T{static_cast<double>(std::rand() / double(RAND_MAX)),
+                 static_cast<double>(std::rand() / double(RAND_MAX))};
+    });
+    for (int i = 0; i < size; i++)
+        answer[i] = x_const[i] * scale + result[i];
+    syncmem_z2z_h2d_op()(gpu_ctx, cpu_ctx, result_gpu, result.data(), sizeof(std::complex<double>) * size);
+    syncmem_z2z_h2d_op()(gpu_ctx, cpu_ctx, x_gpu, x_const.data(), sizeof(std::complex<double>) * size);
+    BlasConnector::axpy(size, scale, x_gpu, incx, result_gpu, incy, base_device::AbacusDevice_t::GpuDevice);
+    syncmem_z2z_d2h_op()(cpu_ctx, gpu_ctx, result.data(), result_gpu, sizeof(std::complex<double>) * size);
+    delmem_zd_op()(gpu_ctx, result_gpu);
+    delmem_zd_op()(gpu_ctx, x_gpu);
+    for (int i = 0; i < size; i++) {
+        EXPECT_DOUBLE_EQ(answer[i].real(), result[i].real());
+        EXPECT_DOUBLE_EQ(answer[i].imag(), result[i].imag());
+    }
+}
+
+#endif
+
 TEST(blas_connector, dcopy_) {
     typedef double T;
     long const size = 8;
@@ -532,7 +620,78 @@ TEST(blas_connector, Gemm) {
         }
 }
 
+#ifdef __CUDA
+
+TEST(blas_connector, GemmGpu) {
+    typedef std::complex<double> T;
+    const char transa_m = 'N';
+    const char transb_m = 'N';
+    const int size_m = 3;
+    const int size_n = 4;
+    const int size_k = 5;
+    const T alpha_const = {2, 3};
+    const T beta_const = {3, 4};
+    const int lda = 6;
+    const int ldb = 5;
+    const int ldc = 4;
+    std::array<T, size_k * lda> a_const;
+    std::array<T, size_n * ldb> b_const;
+    std::array<T, size_n * ldc> c_dot{}, answer, result;
+    std::complex<double>* a_gpu = nullptr;
+    std::complex<double>* b_gpu = nullptr;
+    std::complex<double>* result_gpu = nullptr;
+    resmem_zd_op()(gpu_ctx, a_gpu, size_k * lda * sizeof(std::complex<double>));
+    resmem_zd_op()(gpu_ctx, b_gpu, size_n * ldb * sizeof(std::complex<double>));
+    resmem_zd_op()(gpu_ctx, result_gpu, size_n * ldc * sizeof(std::complex<double>));
+    std::generate(a_const.begin(), a_const.end(), []() {
+        return T{static_cast<double>(std::rand() / double(RAND_MAX)),
+                 static_cast<double>(std::rand() / double(RAND_MAX))};
+    });
+    std::generate(b_const.begin(), b_const.end(), []() {
+        return T{static_cast<double>(std::rand() / double(RAND_MAX)),
+                 static_cast<double>(std::rand() / double(RAND_MAX))};
+    });
+    std::generate(result.begin(), result.end(), []() {
+        return T{static_cast<double>(std::rand() / double(RAND_MAX)),
+                 static_cast<double>(std::rand() / double(RAND_MAX))};
+    });
+    for (int i = 0; i < size_m; i++) {
+        for (int j = 0; j < size_n; j++) {
+            for (int k = 0; k < size_k; k++) {
+                c_dot[i + j * ldc] +=
+                    a_const[i + k * lda] * b_const[k + j * ldb];
+            }
+            answer[i + j * ldc] = alpha_const * c_dot[i + j * ldc] +
+                                  beta_const * result[i + j * ldc];
+        }
+    }
+    syncmem_z2z_h2d_op()(gpu_ctx, cpu_ctx, a_gpu, a_const.data(), sizeof(std::complex<double>) * size_k * lda);
+    syncmem_z2z_h2d_op()(gpu_ctx, cpu_ctx, b_gpu, b_const.data(), sizeof(std::complex<double>) * size_n * ldb);
+    syncmem_z2z_h2d_op()(gpu_ctx, cpu_ctx, result_gpu, result.data(), sizeof(std::complex<double>) * size_n * ldc);
+    BlasConnector::gemm_cm(transa_m, transb_m, size_m, size_n, size_k, alpha_const,
+           a_gpu, lda, b_gpu, ldb, beta_const,
+           result_gpu, ldc, base_device::AbacusDevice_t::GpuDevice);
+    syncmem_z2z_d2h_op()(cpu_ctx, gpu_ctx, result.data(), result_gpu, sizeof(std::complex<double>) * size_n * ldc);
+    delmem_zd_op()(gpu_ctx, result_gpu);
+    delmem_zd_op()(gpu_ctx, a_gpu);
+    delmem_zd_op()(gpu_ctx, b_gpu);
+    for (int i = 0; i < size_m; i++)
+        for (int j = 0; j < size_n; j++) {
+            EXPECT_DOUBLE_EQ(answer[i + j * ldc].real(),
+                             result[i + j * ldc].real());
+            EXPECT_DOUBLE_EQ(answer[i + j * ldc].imag(),
+                             result[i + j * ldc].imag());
+        }
+}
+
+#endif
+
 int main(int argc, char **argv) {
+#ifdef __CUDA
+    std::cout << "Initializing CublasHandle..." << std::endl;
+    BlasUtils::createGpuBlasHandle();
+    std::cout << "Initializing CublasHandle Done." << std::endl;
+#endif
     testing::InitGoogleTest(&argc, argv);
     return RUN_ALL_TESTS();
 }