[GSoC] Feature: MultiDot with cublas<t>gemm (Originally: Jacobi preconditioner on GPU and CUDA Unified preconditiong matrix)

Common/include/linear_algebra/CSysVector.hpp

@@ -585,6 +585,19 @@ class CSysVector : public VecExpr::CVecExpr<CSysVector<ScalarType>, ScalarType>
   static const su2matrix<ScalarType>& multiDot(const std::vector<CSysVector>& V, size_t i0, size_t n,
                                                const std::vector<CSysVector>& W, size_t m);
 
+  /*!
+   * \brief Computes the product of V^T W on the GPU, where V and W are tall matrices stored as vectors of CSysVector.
+   * \param[in] V - Tall matrix.
+   * \param[in] i0 - First column of V to consider.
+   * \param[in] n - Number of columns to consider from V starting at i0.
+   * \param[in] W - Tall matrix.
+   * \param[in] m - Number of columns to consider from W.
+   * \return n by m matrix with the result of the product.
+   */
+  static const su2matrix<ScalarType>& multiDotGPU(const std::vector<CSysVector<ScalarType>>& V, const size_t i0,
+                                                  const size_t n, const std::vector<CSysVector<ScalarType>>& W,
+                                                  const size_t m);
+
   /*!
    * \brief Squared L2 norm of the vector (via dot with self).
    * \return Squared L2 norm.

Common/src/linear_algebra/CSysVector.cpp

@@ -76,66 +76,69 @@ const su2matrix<ScalarType>& CSysVector<ScalarType>::multiDot(const std::vector<
                                                               const std::vector<CSysVector<ScalarType>>& W,
                                                               const size_t m) {
   SU2_ZONE_SCOPED
-  static constexpr size_t BLOCK_SIZE = 1024;
+
   static su2matrix<ScalarType> shared;
 
   if (n == 0 || m == 0) return shared;
 
+  su2matrix<ScalarType> local;
+
+  if (VecExpr::UseDeviceExpressions()) {
 #ifdef SU2_ENABLE_CUDA_KERNELS
-  if constexpr (su2_gpu_capable_v<ScalarType>) {
-    if (VecExpr::UseDeviceExpressions()) {
-      BEGIN_SU2_DEVICE_REGION {
-        shared.resize(n, m);
-        for (size_t i = 0; i < n; ++i) {
-          for (size_t j = 0; j < m; ++j) {
-            shared(i, j) = V[i0 + i].GPUDot(W[j]);
+    if constexpr (su2_gpu_capable_v<ScalarType>) {
+      BEGIN_SU2_DEVICE_REGION
+      local = multiDotGPU(V, i0, n, W, m);
+      END_SU2_DEVICE_REGION
+    } else {
+      SU2_MPI::Error("GPU acceleration is not supported for AD scalar types.", CURRENT_FUNCTION);
+    }
+#else
+    SU2_MPI::Error(
+        "\nError in multiDot\nENABLE_CUDA is set to YES\nPlease compile with CUDA options "
+        "enabled in Meson to access GPU Functions",
+        CURRENT_FUNCTION);
+#endif
+  } else {
+    static constexpr size_t BLOCK_SIZE = 1024;
+
+    SU2_OMP_BARRIER
+    const size_t size = V[0].nElmDomain;
+
+    local.resize(n, m);
+    local.setConstant(0);
+
+    SU2_OMP_FOR_(schedule(static) SU2_NOWAIT)
+    for (size_t offset = 0; offset < size; offset += BLOCK_SIZE) {
+      const auto limit = std::min(offset + BLOCK_SIZE, size);
+      for (size_t i = 0; i < n; ++i) {
+        const auto& vi = V[i0 + i];
+        for (size_t j = 0; j < m; ++j) {
+          const auto& wj = W[j];
+          ScalarType sum = 0.0;
+          SU2_OMP_SIMD
+          for (auto k = offset; k < limit; ++k) {
+            sum += vi[k] * wj[k];
           }
+          local(i, j) += sum;
         }
       }
-      END_SU2_DEVICE_REGION
-      return shared;
     }
-  }
-#endif
-
-  SU2_OMP_BARRIER
-  const size_t size = V[0].nElmDomain;
-
-  su2matrix<ScalarType> local(n, m);
-  local.setConstant(0);
+    END_SU2_OMP_FOR
 
-  SU2_OMP_FOR_(schedule(static) SU2_NOWAIT)
-  for (size_t offset = 0; offset < size; offset += BLOCK_SIZE) {
-    const auto limit = std::min(offset + BLOCK_SIZE, size);
+    /*--- Reduce over all threads in an ordered way to ensure a deterministic result. ---*/
     for (size_t i = 0; i < n; ++i) {
-      const auto& vi = V[i0 + i];
       for (size_t j = 0; j < m; ++j) {
-        const auto& wj = W[j];
-        ScalarType sum = 0.0;
-        SU2_OMP_SIMD
-        for (auto k = offset; k < limit; ++k) {
-          sum += vi[k] * wj[k];
-        }
-        local(i, j) += sum;
+        W[j].dot_scratch[omp_get_thread_num()] = local(i, j);
       }
-    }
-  }
-  END_SU2_OMP_FOR
-
-  /*--- Reduce over all threads in an ordered way to ensure a deterministic result. ---*/
-  for (size_t i = 0; i < n; ++i) {
-    for (size_t j = 0; j < m; ++j) {
-      W[j].dot_scratch[omp_get_thread_num()] = local(i, j);
-    }
-    BEGIN_SU2_OMP_SAFE_GLOBAL_ACCESS
-    for (size_t j = 0; j < m; ++j) {
-      for (int t = 1; t < omp_get_num_threads(); ++t) {
-        local(i, j) += W[j].dot_scratch[t];
+      BEGIN_SU2_OMP_SAFE_GLOBAL_ACCESS
+      for (size_t j = 0; j < m; ++j) {
+        for (int t = 1; t < omp_get_num_threads(); ++t) {
+          local(i, j) += W[j].dot_scratch[t];
+        }
       }
+      END_SU2_OMP_SAFE_GLOBAL_ACCESS
     }
-    END_SU2_OMP_SAFE_GLOBAL_ACCESS
   }
-
   /*--- Single AllReduce of the result, only the master thread communicates. ---*/
   SU2_OMP_MASTER {
     shared.resize(n, m);

Common/src/linear_algebra/CSysVectorGPU.cu

@@ -1,7 +1,7 @@
 /*!
  * \file CSysVectorGPU.cu
  * \brief Implementations of Kernels and Functions for Vector Operations on the GPU
- * \author A. Raj
+ * \author A. Raj, D. Di giusto
  * \version 8.5.0 "Harrier"
  *
  * SU2 Project Website: https://su2code.github.io
@@ -117,6 +117,99 @@ ScalarType CSysVector<ScalarType>::GPUNorm() const {
   return sqrt(GPUDot(*this));
 }
 
+/*!
+ * \brief multi vector product with cublas<t>gemmBatched
+ */
+template <class ScalarType>
+const su2matrix<ScalarType>& CSysVector<ScalarType>::multiDotGPU(const std::vector<CSysVector<ScalarType>>& V,
+                                                                 const size_t i0, const size_t n,
+                                                                 const std::vector<CSysVector<ScalarType>>& W,
+                                                                 const size_t m) {
+  /*--- The multiDot product between n V[size] and m W[size] vectors is performed as
+   * a General Matrix Multiplication between two tall-skinny matrices:
+   * C = \alpha * A^T * B + \beta * C
+   * being A = V[ size * n ] and B = W[ size * m ] the batched vectors ---*/
+  cublasHandle_t handle = GetBlasHandle();
+  cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
+
+  const size_t size = V[0].nElmDomain;
+  const size_t batch = n * m;
+
+  // allocate persisten result buffer local on host and device, is resized if needed
+  static su2matrix<ScalarType> local;
+  local.resize(n,m);
+  static ScalarType* d_local = nullptr;
+  static size_t local_capacity = 0;
+
+  if (batch > local_capacity) { // if not enough capacity, enlarge by re-allocation on device
+    if (d_local) gpuErrChk(cudaFree(d_local));
+    gpuErrChk(cudaMalloc(&d_local, batch * sizeof(ScalarType)));
+    local_capacity = batch;
+  }
+  // zero out the result buffer
+  gpuErrChk(cudaMemset(d_local, 0, batch * sizeof(ScalarType)));
+
+  // prepare the arrays A,B,C on host
+  static std::vector<const ScalarType*> h_A, h_B;
+  static std::vector<ScalarType*> h_C;
+  h_A.resize(batch); h_B.resize(batch); h_C.resize(batch);
+
+  for (size_t i = 0; i < n; ++i) {
+    for (size_t j =0; j < m; ++j) {
+      const size_t idx = i * m + j;
+      h_A[idx] = V[i0 + i].GetDevicePointer();
+      h_B[idx] = W[j].GetDevicePointer();
+      h_C[idx] = d_local + idx; // C maps to d_local to store the coefficients in the 2D array
+    }
+  }
+
+  // prepare device arrays for A,B,C. Re-allocate if batch size has grown
+  static const ScalarType** d_A = nullptr;
+  static const ScalarType** d_B = nullptr;
+  static ScalarType** d_C = nullptr;
+  static size_t ptrs_capacity = 0; // current capacity
+
+  if (batch > ptrs_capacity) { // if not enough capacity, enlarge by re-allocation on device
+    if (d_A) gpuErrChk(cudaFree(d_A));
+    if (d_B) gpuErrChk(cudaFree(d_B));
+    if (d_C) gpuErrChk(cudaFree(d_C));
+    gpuErrChk(cudaMalloc(&d_A, batch * sizeof(ScalarType*)));
+    gpuErrChk(cudaMalloc(&d_B, batch * sizeof(ScalarType*)));
+    gpuErrChk(cudaMalloc(&d_C, batch * sizeof(ScalarType*)));
+    ptrs_capacity = batch; // update current capacity
+  }
+  // copy pointers to device
+  gpuErrChk(cudaMemcpy(d_A, h_A.data(), batch * sizeof(ScalarType*), cudaMemcpyHostToDevice));
+  gpuErrChk(cudaMemcpy(d_B, h_B.data(), batch * sizeof(ScalarType*), cudaMemcpyHostToDevice));
+  gpuErrChk(cudaMemcpy(d_C, h_C.data(), batch * sizeof(ScalarType*), cudaMemcpyHostToDevice));
+
+  // define alpha = 1.0 and beta = 0.0
+  const ScalarType alpha = ScalarType(1.0);
+  const ScalarType beta = ScalarType(0.0);
+
+  if constexpr (std::is_same_v<ScalarType, float>) {
+    status = cublasSgemmBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, 1, 1, size, &alpha, d_A, static_cast<int>(size),
+                                d_B, static_cast<int>(size), &beta, d_C, 1, static_cast<int>(batch));
+  } else if constexpr (std::is_same_v<ScalarType, double>) {
+    status = cublasDgemmBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, 1, 1, size, &alpha, d_A, static_cast<int>(size),
+                                d_B, static_cast<int>(size), &beta, d_C, 1, static_cast<int>(batch));
+  } else {
+    SU2_MPI::Error("Unsupported ScalarType in CSysVector::multiDotGPU.", CURRENT_FUNCTION);
+    return local;
+  }
+
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    SU2_MPI::Error("cuBLAS cublas<t>gemmBatched failed in CSysVector::multiDotGPU.", CURRENT_FUNCTION);
+    return local;
+  }
+
+  // copy result to host for MPI reduce
+  gpuErrChk(cudaMemcpy(local.data(), d_local, batch * sizeof(ScalarType), cudaMemcpyDeviceToHost));
+
+  return local;
+
+}
+
 /*--- Every expression the solvers assign to a CSysVector needs its assignment kernel
  * instantiated here; the host compiler cannot emit one. A shape that is missing shows up
  * as an undefined reference to VecExpr::AssignDeviceExpression at link time, and is fixed
@@ -209,14 +302,21 @@ DEVICE_EXPRESSION_SHAPES(passivedouble);
 #undef INSTANTIATE_DEVICE_ASSIGN_EXPR
 #undef INSTANTIATE_DEVICE_ASSIGN
 
+
 template void CSysVector<su2mixedfloat>::HtDTransfer(bool trigger) const;
 template void CSysVector<su2mixedfloat>::DtHTransfer(bool trigger) const;
 template su2mixedfloat CSysVector<su2mixedfloat>::GPUDot(const CSysVector<su2mixedfloat>& other) const;
 template su2mixedfloat CSysVector<su2mixedfloat>::GPUNorm() const;
+template const su2matrix<su2mixedfloat>& CSysVector<su2mixedfloat>::multiDotGPU(
+    const std::vector<CSysVector<su2mixedfloat>>& V, const size_t i0, const size_t n,
+    const std::vector<CSysVector<su2mixedfloat>>& W, const size_t m);
 
 #if defined(USE_MIXED_PRECISION) && !defined(USE_SINGLE_PRECISION)
 template void CSysVector<passivedouble>::HtDTransfer(bool trigger) const;
 template void CSysVector<passivedouble>::DtHTransfer(bool trigger) const;
 template passivedouble CSysVector<passivedouble>::GPUDot(const CSysVector<passivedouble>& other) const;
 template passivedouble CSysVector<passivedouble>::GPUNorm() const;
+template const su2matrix<passivedouble>& CSysVector<passivedouble>::multiDotGPU(
+    const std::vector<CSysVector<passivedouble>>& V, const size_t i0, const size_t n,
+    const std::vector<CSysVector<passivedouble>>& W, const size_t m);
 #endif