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

Common/include/linear_algebra/CSysVector.hpp

@@ -270,6 +270,17 @@ class CSysVector : public VecExpr::CVecExpr<CSysVector<ScalarType>, ScalarType>
     return *this;
   }
 
+  /*!
+   * \brief GPU helper for `dot`.
+   */
+  ScalarType dotGPU(const CSysVector& other) const;
+
+  /*!
+   * \brief GPU helper for multiDot.
+   */
+  static su2matrix<ScalarType> multiDotGPU(const std::vector<CSysVector<ScalarType>>& V, size_t i0, size_t n,
+                                           const std::vector<CSysVector<ScalarType>>& W, size_t m);
+
  public:
   static constexpr bool StoreAsRef = true; /*! \brief Required by CVecExpr. */
 
@@ -394,21 +405,6 @@ class CSysVector : public VecExpr::CVecExpr<CSysVector<ScalarType>, ScalarType>
    */
   void DtHTransfer(bool trigger = true) const;
 
-  /*!
-   * \brief Dot product between this vector and another vector on the device.
-   * \note Explicit GPU helper for solver-side reductions.
-   * \param[in] other - Input vector.
-   * \return Dot product result.
-   */
-  ScalarType GPUDot(const CSysVector& other) const;
-
-  /*!
-   * \brief L2 norm of this vector on the device.
-   * \note Explicit GPU helper for solver-side reductions.
-   * \return L2 norm result.
-   */
-  ScalarType GPUNorm() const;
-
   /*!
    * \brief return device pointer that points to the CSysVector values in GPU memory
    */
@@ -537,9 +533,9 @@ class CSysVector : public VecExpr::CVecExpr<CSysVector<ScalarType>, ScalarType>
                     "On the device the dot product is a cuBLAS call, so it only takes vectors. "
                     "Assign the expression to a vector first.");
       if (VecExpr::UseDeviceExpressions()) {
-        /*--- GPUDot reduces over MPI, which has to happen once for the team, so the result
+        /*--- dotGPU reduces over MPI, which has to happen once for the team, so the result
          * is published through the same scratch slot the host reduction below uses. ---*/
-        SU2_DEVICE_REGION(dot_scratch[0] = GPUDot(expr.derived());)
+        SU2_DEVICE_REGION(dot_scratch[0] = dotGPU(expr.derived());)
         return dot_scratch[0];
       }
     }

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
@@ -79,7 +79,7 @@ void CSysVector<ScalarType>::DtHTransfer(bool trigger) const {
 }
 
 template <class ScalarType>
-ScalarType CSysVector<ScalarType>::GPUDot(const CSysVector& other) const {
+ScalarType CSysVector<ScalarType>::dotGPU(const CSysVector& other) const {
   SU2_ZONE_SCOPED
   /*--- Both operands are already on the device, the caller owns the transfers. This
    * reduces over MPI, so it must be called by a single thread (see SU2_DEVICE_REGION). ---*/
@@ -95,12 +95,12 @@ ScalarType CSysVector<ScalarType>::GPUDot(const CSysVector& other) const {
     status = cublasDdot(handle, static_cast<int>(nElmDomain), GetDevicePointer(), 1, other.GetDevicePointer(), 1,
                         &local_dot);
   } else {
-    SU2_MPI::Error("Unsupported ScalarType in CSysVector::GPUDot.", CURRENT_FUNCTION);
+    SU2_MPI::Error("Unsupported ScalarType in CSysVector::dotGPU.", CURRENT_FUNCTION);
     return ScalarType(0);
   }
 
   if (status != CUBLAS_STATUS_SUCCESS) {
-    SU2_MPI::Error("cuBLAS dot failed in CSysVector::GPUDot.", CURRENT_FUNCTION);
+    SU2_MPI::Error("cuBLAS dot failed in CSysVector::dotGPU.", CURRENT_FUNCTION);
     return ScalarType(0);
   }
 
@@ -111,10 +111,105 @@ ScalarType CSysVector<ScalarType>::GPUDot(const CSysVector& other) const {
   return global_dot;
 }
 
+/*!
+ * \brief multi vector product with cublas<t>gemmBatched
+ */
 template <class ScalarType>
-ScalarType CSysVector<ScalarType>::GPUNorm() const {
-  SU2_ZONE_SCOPED
-  return sqrt(GPUDot(*this));
+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;
+
+  /*--- Persistent device workspace, cached across calls and freed automatically
+   * when the program exits (static local destruction), instead of leaking. ---*/
+  struct Workspace {
+    ScalarType* d_local = nullptr;
+    const ScalarType** d_A = nullptr;
+    const ScalarType** d_B = nullptr;
+    ScalarType** d_C = nullptr;
+    size_t capacity = 0;
+
+    void EnsureCapacity(size_t batch) {
+      if (batch <= capacity) return;
+      cudaFree(d_local);
+      cudaFree(d_A);
+      cudaFree(d_B);
+      cudaFree(d_C);
+      gpuErrChk(cudaMalloc(&d_local, batch * sizeof(ScalarType)));
+      gpuErrChk(cudaMalloc(&d_A, batch * sizeof(ScalarType*)));
+      gpuErrChk(cudaMalloc(&d_B, batch * sizeof(ScalarType*)));
+      gpuErrChk(cudaMalloc(&d_C, batch * sizeof(ScalarType*)));
+      capacity = batch;
+    }
+
+    ~Workspace() {
+      cudaFree(d_local);
+      cudaFree(d_A);
+      cudaFree(d_B);
+      cudaFree(d_C);
+    }
+  };
+  static Workspace ws;
+
+  // allocate persistent result buffer local on host and device, is resized if needed
+  su2matrix<ScalarType> local;
+  local.resize(n, m);
+  ws.EnsureCapacity(batch);
+
+  // zero out the result buffer
+  gpuErrChk(cudaMemset(ws.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] = ws.d_local + idx; // C maps to d_local to store the coefficients in the 2D array
+    }
+  }
+
+  // copy pointers to device
+  gpuErrChk(cudaMemcpy(ws.d_A, h_A.data(), batch * sizeof(ScalarType*), cudaMemcpyHostToDevice));
+  gpuErrChk(cudaMemcpy(ws.d_B, h_B.data(), batch * sizeof(ScalarType*), cudaMemcpyHostToDevice));
+  gpuErrChk(cudaMemcpy(ws.d_C, h_C.data(), batch * sizeof(ScalarType*), cudaMemcpyHostToDevice));
+
+  // define alpha = 1.0 and beta = 0.0
+  const auto alpha = ScalarType(1.0);
+  const auto 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, ws.d_A, static_cast<int>(size),
+                                ws.d_B, static_cast<int>(size), &beta, ws.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, ws.d_A, static_cast<int>(size),
+                                ws.d_B, static_cast<int>(size), &beta, ws.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(), ws.d_local, batch * sizeof(ScalarType), cudaMemcpyDeviceToHost));
+
+  return local;
 }
 
 /*--- Every expression the solvers assign to a CSysVector needs its assignment kernel
@@ -209,14 +304,19 @@ 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 su2mixedfloat CSysVector<su2mixedfloat>::dotGPU(const CSysVector<su2mixedfloat>& other) const;
+template su2matrix<su2mixedfloat> CSysVector<su2mixedfloat>::multiDotGPU(
+    const std::vector<CSysVector<su2mixedfloat>>& V, size_t i0, size_t n,
+    const std::vector<CSysVector<su2mixedfloat>>& W, 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 passivedouble CSysVector<passivedouble>::dotGPU(const CSysVector<passivedouble>& other) const;
+template su2matrix<passivedouble> CSysVector<passivedouble>::multiDotGPU(
+    const std::vector<CSysVector<passivedouble>>& V, size_t i0, size_t n,
+    const std::vector<CSysVector<passivedouble>>& W, size_t m);
 #endif