From ed440c16a2f5f9541eb22003e64d7e6c5afc0842 Mon Sep 17 00:00:00 2001
From: Anthony Scemama <scemama@irsamc.ups-tlse.fr>
Date: Sun, 16 Jul 2023 17:27:14 +0200
Subject: [PATCH] Working

---
 devel/ccsd_gpu/gpu.c | 272 +++++++++++++++++++++++++------------------
 1 file changed, 157 insertions(+), 115 deletions(-)

diff --git a/devel/ccsd_gpu/gpu.c b/devel/ccsd_gpu/gpu.c
index 47cf01b..71a2c09 100644
--- a/devel/ccsd_gpu/gpu.c
+++ b/devel/ccsd_gpu/gpu.c
@@ -1,9 +1,10 @@
 #include <stdio.h>
 #include <stdlib.h>
+#include <omp.h>
 #include <cublas_v2.h>
 #include <cuda_runtime.h>
 
-
+#define NGPUS 2
 #define BLOCK_SIZE 16
 
 void dgemm_(char*, char*, int*, int*, int*, double*, double*, int*, double*, int*,
@@ -65,134 +66,175 @@ void compute_r2_space_chol_gpu(const int nO, const int nV, const int cholesky_mo
                 double* cc_space_v_vv_chol,
                 double* r2)
 {
-    int m,n,k, lda, ldb, ldc;
-    double alpha, beta;
-    double* A;
-    double* B;
-    double* C;
+   int m,n,k, lda, ldb, ldc;
+   double alpha, beta;
+   double* A;
+   double* B;
+   double* C;
 
-    cublasHandle_t handle;
-    cublasCreate(&handle);
-
-    double* d_tau;
-    lda = nO * nO;
-    cudaMalloc((void **)&d_tau, lda * nV * nV * sizeof(double));
-    cublasSetMatrix(nO*nO, nV*nV, sizeof(double), tau, lda, d_tau, lda);
-
-    double* d_r2;
-    lda = nO * nO;
-    cudaMalloc((void **)&d_r2, lda * nV * nV * sizeof(double));
-
-    double* d_cc_space_v_vv_chol;
-    lda = cholesky_mo_num * nV;
-    cudaMalloc((void **)&d_cc_space_v_vv_chol, lda * nV * sizeof(double));
-    cublasSetMatrix(cholesky_mo_num*nV, nV, sizeof(double), cc_space_v_vv_chol, lda, d_cc_space_v_vv_chol, lda);
-
-    double* d_cc_space_v_vo_chol;
-    lda = cholesky_mo_num * nV;
-    cudaMalloc((void **)&d_cc_space_v_vo_chol, lda * nO * sizeof(double));
-    cublasSetMatrix(cholesky_mo_num*nV, nO, sizeof(double), cc_space_v_vo_chol, lda, d_cc_space_v_vo_chol, lda);
-
-    double* d_t1;
-    lda = nO;
-    cudaMalloc((void **)&d_t1, nO * nV * sizeof(double));
-    cublasSetMatrix(nO, nV, sizeof(double), t1, lda, d_t1, lda);
-
-    double* d_tmp_cc;
-    lda = cholesky_mo_num * nV;
-    cudaMalloc((void **)&d_tmp_cc, lda * nV * sizeof(double));
-
-    alpha=1.0; beta=0.0;
-    m=cholesky_mo_num*nV;  n=nV;  k=nO;
-    A = d_cc_space_v_vo_chol; B = d_t1; C = d_tmp_cc;
-    cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, m, B, k, &beta, C, m);
-    cublasDestroy(handle);
+   double* d_taus[NGPUS];
+   double* d_r2s[NGPUS];
+   double* d_cc_space_v_vv_chols[NGPUS];
+   double* d_cc_space_v_vo_chols[NGPUS];
+   double* d_t1s[NGPUS];
+   double* d_tmp_ccs[NGPUS];
 
+    cublasHandle_t handles[NGPUS];
     #pragma omp parallel
     {
-        cublasHandle_t handle;
-        cublasCreate(&handle);
 
-        double* d_tmp_cc2;
-        cudaMalloc((void **)&d_tmp_cc2, cholesky_mo_num*nV*sizeof(double));
+      int ithread = omp_get_thread_num();
+      int igpu = ithread % NGPUS;
+      cudaSetDevice(igpu);
 
-        double* d_B1;
-        cudaMalloc((void**)&d_B1, nV*nV*BLOCK_SIZE*sizeof(double));
+      if (ithread < NGPUS) {
+        cublasCreate(&handles[ithread]);
+      }
 
-        double* d_tmpB1;
-        cudaMalloc((void**)&d_tmpB1, nV*BLOCK_SIZE*nV*sizeof(double));
+      #pragma omp barrier
 
-        #pragma omp for
-        for (size_t gam=0 ; gam<nV ; ++gam)
+      cublasHandle_t handle = handles[igpu];
+
+      double* d_tau;
+      double* d_r2;
+      double* d_cc_space_v_vv_chol;
+      double* d_cc_space_v_vo_chol;
+      double* d_t1;
+      double* d_tmp_cc;
+
+      if (ithread < NGPUS) {
+        lda = nO * nO;
+        cudaMalloc((void **)&d_tau, lda * nV * nV * sizeof(double));
+        cublasSetMatrix(nO*nO, nV*nV, sizeof(double), tau, lda, d_tau, lda);
+        d_taus[igpu] = d_tau;
+
+        lda = nO * nO;
+        cudaMalloc((void **)&d_r2, lda * nV * nV * sizeof(double));
+        d_r2s[igpu] = d_r2;
+
+        lda = cholesky_mo_num * nV;
+        cudaMalloc((void **)&d_cc_space_v_vv_chol, lda * nV * sizeof(double));
+        cublasSetMatrix(cholesky_mo_num*nV, nV, sizeof(double), cc_space_v_vv_chol, lda, d_cc_space_v_vv_chol, lda);
+        d_cc_space_v_vv_chols[igpu] = d_cc_space_v_vv_chol;
+
+        lda = cholesky_mo_num * nV;
+        cudaMalloc((void **)&d_cc_space_v_vo_chol, lda * nO * sizeof(double));
+        cublasSetMatrix(cholesky_mo_num*nV, nO, sizeof(double), cc_space_v_vo_chol, lda, d_cc_space_v_vo_chol, lda);
+        d_cc_space_v_vo_chols[igpu] = d_cc_space_v_vo_chol;
+
+        lda = nO;
+        cudaMalloc((void **)&d_t1, nO * nV * sizeof(double));
+        cublasSetMatrix(nO, nV, sizeof(double), t1, lda, d_t1, lda);
+        d_t1s[igpu] = d_t1;
+
+        lda = cholesky_mo_num * nV;
+        cudaMalloc((void **)&d_tmp_cc, lda * nV * sizeof(double));
+        d_tmp_ccs[igpu] = d_tmp_cc;
+
+        alpha=1.0; beta=0.0;
+        m=cholesky_mo_num*nV;  n=nV;  k=nO;
+        A = d_cc_space_v_vo_chol; B = d_t1; C = d_tmp_cc;
+        cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, m, B, k, &beta, C, m);
+      }
+
+      #pragma omp barrier
+
+      d_tau = d_taus[igpu] ;
+      d_r2 = d_r2s[igpu] ;
+      d_cc_space_v_vv_chol = d_cc_space_v_vv_chols[igpu] ;
+      d_cc_space_v_vo_chol = d_cc_space_v_vo_chols[igpu] ;
+      d_t1 = d_t1s[igpu] ;
+      d_tmp_cc = d_tmp_ccs[igpu] ;
+
+      double* d_tmp_cc2;
+      cudaMalloc((void **)&d_tmp_cc2, cholesky_mo_num*nV*sizeof(double));
+
+      double* d_B1;
+      cudaMalloc((void**)&d_B1, nV*nV*BLOCK_SIZE*sizeof(double));
+
+      double* d_tmpB1;
+      cudaMalloc((void**)&d_tmpB1, nV*BLOCK_SIZE*nV*sizeof(double));
+
+      #pragma omp for
+      for (size_t gam=0 ; gam<nV ; ++gam)
+      {
+         double* d_tmp_cc_ = &(d_tmp_cc[gam*nV*cholesky_mo_num]);
+         double* d_cc_space_v_vv_chol_ = &(d_cc_space_v_vv_chol[gam*nV*cholesky_mo_num]);
+
+         alpha = 1.0;
+         beta = -1.0;
+         A = d_cc_space_v_vv_chol_; lda = cholesky_mo_num;
+         B = d_tmp_cc_; ldb = cholesky_mo_num;
+         C = d_tmp_cc2 ; ldc = cholesky_mo_num;
+         cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, cholesky_mo_num, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
+
+         for (size_t iblock=0 ; iblock<nV ; iblock += BLOCK_SIZE)
+         {
+              const size_t mbs = BLOCK_SIZE < nV-iblock ? BLOCK_SIZE : nV-iblock;
+
+              alpha=-1.0; beta=0.0;
+              m=nV*mbs;  n=nV;  k=cholesky_mo_num;
+
+              A=&(d_tmp_cc[iblock*cholesky_mo_num*nV]); lda=cholesky_mo_num;
+              B=d_cc_space_v_vv_chol_; ldb=cholesky_mo_num;
+              C=d_tmpB1 ; ldc=nV*BLOCK_SIZE;
+              cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, lda, &beta, C, ldc);
+
+              alpha=1.0; beta=1.0;
+              m=nV*mbs;  n=nV;  k=cholesky_mo_num;
+
+              A=&(d_cc_space_v_vv_chol[iblock*cholesky_mo_num*nV]); lda=cholesky_mo_num;
+              B=d_tmp_cc2; ldb=cholesky_mo_num;
+              C=d_tmpB1 ; ldc=nV*BLOCK_SIZE;
+              cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, lda, &beta, C, ldc);
+
+              for (size_t bet=iblock ; bet<(nV < iblock+BLOCK_SIZE ? nV : iblock+BLOCK_SIZE) ; ++bet)
+              {
+
+                alpha = 1.0;
+                beta = 0.0;
+                A = &(d_tmpB1[nV*(bet-iblock)]); lda = nV*BLOCK_SIZE;
+                B = d_tmpB1; ldb = nV;
+                C = &(d_B1[nV*nV*(bet-iblock)]) ; ldc = nV;
+                cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nV, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
+             }
+
+             alpha=1.0; beta=1.0;
+             m=nO*nO;  n=mbs;  k=nV*nV;
+
+             A=d_tau; lda=nO*nO;
+             B=d_B1 ; ldb=nV*nV;
+             C=&(d_r2[nO*nO*(iblock + nV*gam)]); ldc=nO*nO;
+             cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
+
+         }
+      }
+      cudaFree(d_tmpB1);
+      cudaFree(d_B1);
+      cudaFree(d_tmp_cc2);
+
+      if (igpu < NGPUS) {
+        cudaFree(d_cc_space_v_vo_chol);
+        cudaFree(d_cc_space_v_vv_chol);
+        cudaFree(d_tau);
+        cudaFree(d_t1);
+        cudaFree(d_tmp_cc);
+        double * r2_tmp = malloc(nO*nO*nV*nV*sizeof(double));
+        lda=nO*nO;
+        cublasGetMatrix(nO*nO, nV*nV, sizeof(double), d_r2, lda, r2_tmp, lda);
+        #pragma omp critical
         {
-           double* d_tmp_cc_ = &(d_tmp_cc[gam*nV*cholesky_mo_num]);
-           double* d_cc_space_v_vv_chol_ = &(d_cc_space_v_vv_chol[gam*nV*cholesky_mo_num]);
-
-           alpha = 1.0;
-           beta = -1.0;
-           A = d_cc_space_v_vv_chol_; lda = cholesky_mo_num;
-           B = d_tmp_cc_; ldb = cholesky_mo_num;
-           C = d_tmp_cc2 ; ldc = cholesky_mo_num;
-           cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, cholesky_mo_num, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
-
-           for (size_t iblock=0 ; iblock<nV ; iblock += BLOCK_SIZE)
-           {
-                const size_t mbs = BLOCK_SIZE < nV-iblock ? BLOCK_SIZE : nV-iblock;
-
-                alpha=-1.0; beta=0.0;
-                m=nV*mbs;  n=nV;  k=cholesky_mo_num;
-
-                A=&(d_tmp_cc[iblock*cholesky_mo_num*nV]); lda=cholesky_mo_num;
-                B=d_cc_space_v_vv_chol_; ldb=cholesky_mo_num;
-                C=d_tmpB1 ; ldc=nV*BLOCK_SIZE;
-                cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, lda, &beta, C, ldc);
-
-                alpha=1.0; beta=1.0;
-                m=nV*mbs;  n=nV;  k=cholesky_mo_num;
-
-                A=&(d_cc_space_v_vv_chol[iblock*cholesky_mo_num*nV]); lda=cholesky_mo_num;
-                B=d_tmp_cc2; ldb=cholesky_mo_num;
-                C=d_tmpB1 ; ldc=nV*BLOCK_SIZE;
-                cublasDgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, lda, &beta, C, ldc);
-
-                for (size_t bet=iblock ; bet<(nV < iblock+BLOCK_SIZE ? nV : iblock+BLOCK_SIZE) ; ++bet)
-                {
-
-                  alpha = 1.0;
-                  beta = 0.0;
-                  A = &(d_tmpB1[nV*(bet-iblock)]); lda = nV*BLOCK_SIZE;
-                  B = d_tmpB1; ldb = nV;
-                  C = &(d_B1[nV*nV*(bet-iblock)]) ; ldc = nV;
-                  cublasDgeam(handle, CUBLAS_OP_N, CUBLAS_OP_N, nV, nV, &alpha, A, lda, &beta, B, ldb, C, ldc);
-               }
-
-               alpha=1.0; beta=1.0;
-               m=nO*nO;  n=mbs;  k=nV*nV;
-
-               A=d_tau; lda=nO*nO;
-               B=d_B1 ; ldb=nV*nV;
-               C=&(d_r2[nO*nO*(iblock + nV*gam)]); ldc=nO*nO;
-               cublasDgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
-
+           for (size_t i=0 ; i<nO*nO*nV*nV ; ++i) {
+             r2[i] += r2_tmp[i];
            }
         }
+        free(r2_tmp);
+
+        cudaFree(d_r2);
         cublasDestroy(handle);
     }
-    cudaFree(d_cc_space_v_vo_chol);
-    cudaFree(d_cc_space_v_vv_chol);
-    cudaFree(d_tau);
-    cudaFree(d_t1);
-    cudaFree(d_tmp_cc);
 
-    double * r2_tmp = malloc(nO*nO*nV*nV*sizeof(double));
-    lda=nO*nO;
-    cublasGetMatrix(nO*nO, nV*nV, sizeof(double), d_r2, lda, r2_tmp, lda);
-    for (size_t i=0 ; i<nO*nO*nV*nV ; ++i) {
-      r2[i] += r2_tmp[i];
-    }
-    free(r2_tmp);
-
-    cudaFree(d_r2);
+  }
 }