[GSoC] Feature: MultiDot and Linear Combination CUDA kernels (Originally: Jacobi preconditioner on GPU and CUDA Unified preconditiong matrix)

Common/include/basic_types/datatype_structure.hpp

@@ -135,10 +135,12 @@ template <>
 struct Passive<su2double> {
   FORCEINLINE static passivedouble Value(const su2double& val) { return GetValue(val); }
 };
+#ifndef SWIG
 template <class T>
 FORCEINLINE auto PassiveValue(const T& val) {
   return Passive<T>::Value(val);
 }
+#endif
 
 /*!
  * \brief Casts the primitive value to int (uses GetValue, already implemented for each type).

Common/include/linear_algebra/CSysMatrix.hpp

@@ -250,7 +250,7 @@ class CSysMatrix {
   LDU mat;                      /*!< \brief Host matrix (values owned via aligned_alloc; pattern from geometry). */
   LDU gpu;                      /*!< \brief Device matrix (all pointers to GPU memory). */
   LDU ilu;                      /*!< \brief ILU factorization, host (values owned; pattern from geometry). */
-  ScalarType* d_invM = nullptr; /*!< \brief Device inverse diagonal blocks for the Jacobi preconditioner. */
+  ScalarType* d_invM = nullptr; /*!< \brief Device inverse diagonal blocks for the (Jacobi) preconditioner. */
 
   /*--- Quantized off-diagonal storage (used when quantized_mode == true). ---*/
   using QuantType = int8_t;

Common/include/linear_algebra/CSysVector.hpp

@@ -585,6 +585,30 @@ 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 Computes v = vs * ws or v += vs * ws with unrolling of up to 4 iterations on the GPU
+   * \param[in] n - number of vectors to consider
+   * \param[in] ws - array of scalar weights corresponding to the device pointers to vectors
+   * \param[in] vs_ptrs - array of device pointers
+   * \param[in] v - target vector
+   * \param[in] inc - If true, adds results to target vector. If false, overwrites
+   */
+  static void LinearCombinationGPU(const unsigned long n, const std::vector<CSysVector<ScalarType>>& vs,
+                                   const ScalarType* ws, CSysVector<ScalarType>& v, bool inc = false);
+
   /*!
    * \brief Squared L2 norm of the vector (via dot with self).
    * \return Squared L2 norm.

Common/src/linear_algebra/CSysPreconditionerGPU.cu

@@ -27,6 +27,7 @@
 
 #include "../../include/linear_algebra/CSysMatrix.inl"
 #include "../../include/linear_algebra/GPUComms.cuh"
+#include "../../include/linear_algebra/CSysMatrix.hpp"
 
 namespace {
 
@@ -61,7 +62,7 @@ void CSysMatrix<ScalarType>::ComputeJacobiPreconditionerGPU(const CSysVector<Sca
   SU2_ZONE_SCOPED
 
   if (d_invM == nullptr) {
-    SU2_MPI::Error("CUDA Jacobi preconditioner used before BuildJacobiPreconditionerGPU.", CURRENT_FUNCTION);
+    SU2_MPI::Error("CUDA Jacobi preconditioner used before BuildJacobiPreconditioner.", CURRENT_FUNCTION);
   }
 
   constexpr unsigned threadsPerBlock = 128;

Common/src/linear_algebra/CSysSolve.cpp

@@ -33,6 +33,7 @@
 #include "../../include/linear_algebra/CSysMatrix.hpp"
 #include "../../include/linear_algebra/CMatrixVectorProduct.hpp"
 #include "../../include/linear_algebra/CPreconditioner.hpp"
+#include "../../include/linear_algebra/CSysVector.hpp"
 
 SU2_IGNORE_WARNING("-Wmaybe-uninitialized")
 #include "Eigen/Eigenvalues"
@@ -107,6 +108,21 @@ void LinearCombinationImpl(const unsigned long n, const Vectors& vs, const Weigh
 template <class ScalarType, class Weights>
 void LinearCombinationImpl(const unsigned long n, const std::vector<CSysVector<ScalarType>>& vs, const Weights& ws,
                            CSysVector<ScalarType>& v, bool inc = false) {
+#ifdef SU2_ENABLE_CUDA_KERNELS
+  if constexpr (su2_gpu_capable_v<ScalarType>) {
+    std::vector<ScalarType> ws_host(n);  // collect weights into simple host array
+    for (unsigned long i = 0; i < n; ++i) {
+      ws_host[i] = static_cast<ScalarType>(ws(i));
+    }
+    BEGIN_SU2_DEVICE_REGION
+    CSysVector<ScalarType>::LinearCombinationGPU(n, vs, ws_host.data(), v, inc);
+    END_SU2_DEVICE_REGION
+    return;
+  } else {
+    SU2_MPI::Error("GPU acceleration is not supported for AD scalar types.", CURRENT_FUNCTION);
+  }
+#endif
+
   LinearCombinationImpl(
       n, [&vs](auto i) -> auto& { return vs[i]; }, ws, v, inc);
 }

Common/src/linear_algebra/CSysVector.cpp

@@ -76,27 +76,21 @@ 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;
 
 #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]);
-          }
-        }
-      }
-      END_SU2_DEVICE_REGION
-      return shared;
-    }
+    BEGIN_SU2_DEVICE_REGION
+    shared = 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);
   }
-#endif
+#else
+  static constexpr size_t BLOCK_SIZE = 1024;
 
   SU2_OMP_BARRIER
   const size_t size = V[0].nElmDomain;
@@ -148,7 +142,7 @@ const su2matrix<ScalarType>& CSysVector<ScalarType>::multiDot(const std::vector<
 
   /*--- All threads have the same view of the result. ---*/
   SU2_OMP_BARRIER
-
+#endif
   return shared;
 }
 

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
@@ -113,6 +113,196 @@ ScalarType CSysVector<ScalarType>::GPUNorm() const {
   return sqrt(GPUDot(*this));
 }
 
+/*!
+ * \brief multi vector product CUDA kernel one line of blocks per pair V[i0+i],W[j];
+ *        Configurable multiple blocks reducing over the size of the vectors
+ */
+template <class ScalarType>
+__global__ void GPUmultiDot(const ScalarType* const* __restrict__ d_V, const size_t n,
+                            const ScalarType* const* __restrict__ d_W, const size_t m, const size_t size,
+                            ScalarType* __restrict__ d_local)
+{
+  // Map each x,y block to the specific (i,j) dot product
+  const size_t pair_idx = blockIdx.y;
+  if (pair_idx >= n * m) return;
+
+  const size_t i = pair_idx / m;
+  const size_t j = pair_idx % m;
+
+  //get the corresponding vectors
+  const ScalarType* __restrict__ vi = d_V[i];
+  const ScalarType* __restrict__ wj = d_W[j];
+
+  // grid strided loop over the vector elements
+  ScalarType local_sum = 0.0;
+  const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const size_t stride = gridDim.x * blockDim.x;
+
+  for (size_t k = tid; k < size; k += stride)
+  {
+      local_sum += vi[k] * wj[k];
+  }
+
+  // shared memory reduction within the block
+  extern __shared__ char shared_mem[];
+  ScalarType* sdata = reinterpret_cast<ScalarType*>(shared_mem);
+
+  sdata[threadIdx.x] = local_sum;
+  __syncthreads();
+
+  // parallel reduction on the block
+  for (unsigned int s = blockDim.x / 2; s > 0; s >>= 1)
+  {
+      if (threadIdx.x < s)
+      {
+        sdata[threadIdx.x] += sdata[threadIdx.x + s];
+      }
+      __syncthreads();
+  }
+
+  // atomic add of each block partial sum to the output matrix, operated by thread 0 of each block
+  if (threadIdx.x == 0)
+  {
+      atomicAdd(&d_local[i * m + j], sdata[0]);
+  }
+}
+
+/*!
+ * \brief multi vector dot produt method for GPU
+ * \note this is a vectors-read only method that returns an array of scalars
+ */
+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) {
+
+  static su2matrix<ScalarType> shared;
+  if (n == 0 || m == 0) return shared;
+
+  const size_t size = V[0].nElmDomain;
+
+  // get all the device pointers for V and W in one array, resize if needed
+  static std::vector<const ScalarType*> h_V_W_ptrs;
+  h_V_W_ptrs.resize(n + m);
+
+  for (size_t i = 0; i < n; ++i){
+    h_V_W_ptrs[i] = V[i0 + i].GetDevicePointer();
+  }
+  for (size_t j = 0; j < m; ++j){
+    h_V_W_ptrs[j + n] = W[j].GetDevicePointer();
+  }
+
+  // persistent device pointer storing all vectors, resizes when needed
+  static const ScalarType** d_V_W_ptrs = nullptr;
+  static size_t ptrs_capacity = 0; // current capacity
+  const size_t ptrs_needed = n + m; // needed capacity for both V and W
+
+  if (ptrs_needed > ptrs_capacity) { // if not enough capacity, enlarge by re-allocation on device
+    if (d_V_W_ptrs) gpuErrChk(cudaFree(d_V_W_ptrs));
+    gpuErrChk(cudaMalloc(&d_V_W_ptrs, ptrs_needed * sizeof(ScalarType*)));
+    ptrs_capacity = ptrs_needed; // update current capacity
+  }
+  // copy pointers to device
+  gpuErrChk(cudaMemcpy(d_V_W_ptrs, h_V_W_ptrs.data(), ptrs_needed * sizeof(ScalarType*), cudaMemcpyHostToDevice));
+
+  // allocate persisten result buffer that grows if needed
+  static ScalarType* d_local = nullptr;
+  static size_t local_capacity = 0;
+  const size_t local_needed = n * m;
+
+  if (local_needed > local_capacity) { // if not enough capacity, enlarge by re-allocation on device
+    if (d_local) gpuErrChk(cudaFree(d_local));
+    gpuErrChk(cudaMalloc(&d_local, local_needed * sizeof(ScalarType)));
+    local_capacity = local_needed;
+  }
+  // zero out the result buffer
+  gpuErrChk(cudaMemset(d_local, 0, local_needed * sizeof(ScalarType)));
+
+
+  dim3 blockDim(KernelParameters::MVP_BLOCK_SIZE,1,1);
+  int numBlocksPerPair = KernelParameters::round_up_division(KernelParameters::MVP_BLOCK_SIZE, size);
+  dim3 gridDim(numBlocksPerPair, n * m, 1);
+
+  GPUmultiDot<<<gridDim, blockDim, KernelParameters::MVP_BLOCK_SIZE * sizeof(ScalarType)>>>(&d_V_W_ptrs[0], n, &d_V_W_ptrs[n], m, size, d_local);
+  gpuErrChk(cudaGetLastError());
+
+  // copy result to host for MPI reduce
+  su2matrix<ScalarType> local(n,m);
+  gpuErrChk(cudaMemcpy(local.data(), d_local, n * m * sizeof(ScalarType), cudaMemcpyDeviceToHost));
+
+  /*--- Single AllReduce of the result, only the master thread communicates. ---*/
+  // this is a duplicate. Ideally the cuda section should return local but that depends on the intended OpenMP/CUDA combined usage
+  SU2_OMP_MASTER {
+    shared.resize(n, m);
+
+    const auto mpi_type = (sizeof(ScalarType) < sizeof(double)) ? MPI_FLOAT : MPI_DOUBLE;
+    SelectMPIWrapper<ScalarType>::W::Allreduce(local.data(), shared.data(), n * m, mpi_type, MPI_SUM,
+                                               SU2_MPI::GetComm());
+  }
+  END_SU2_OMP_MASTER
+
+  /*--- All threads have the same view of the result. ---*/
+  SU2_OMP_BARRIER
+
+  return shared;
+}
+
+template<class ScalarType, int N>
+struct WeightedVecs {
+  const ScalarType* ptrs[N];
+  ScalarType weights[N];
+};
+
+/*!
+ * \brief linear combination kernel to calculate the next vector v from weights and vectors
+ */
+template<class ScalarType, int N>
+__global__ void LinearCombinationKernel(ScalarType* __restrict__ v, WeightedVecs<ScalarType, N> wv,
+                                        int n, unsigned long nElm, bool inc)
+{
+  const unsigned long k = blockIdx.x * blockDim.x + threadIdx.x;
+  if (k >= nElm) return;
+
+  //handle overwriting or combination with existing
+  ScalarType result = inc ? v[k] : ScalarType(0);
+
+  #pragma unroll
+  for (int i = 0; i < N; ++i) // N is known at compile time (4), this unrolls to: if (i < n) result += weight[i] * vector[i][k]; i<4
+    if (i < n) result += wv.weights[i] * wv.ptrs[i][k];
+  v[k] = result;
+}
+
+/*!
+ * \brief dispatcher for the linear combination kernel on GPU
+ */
+template<class ScalarType>
+void CSysVector<ScalarType>::LinearCombinationGPU(const unsigned long n, const std::vector<CSysVector<ScalarType>>& vs, const ScalarType* ws,
+                                                  CSysVector<ScalarType>& v, bool inc)
+{
+  const unsigned long nElm = v.nElmDomain;
+  dim3 blockDim(KernelParameters::MVP_BLOCK_SIZE,1,1);
+  int numBlocks = KernelParameters::round_up_division(KernelParameters::MVP_BLOCK_SIZE, nElm);
+  dim3 gridDim(numBlocks, 1, 1);
+
+  ScalarType* d_v = v.GetDevicePointer();
+
+  for (unsigned long i = 0; i < n; i += 4) {
+    const int rem = static_cast<int>(std::min(n - i, 4ul));
+    //prepare vectors pointers and corresponding weights, passing them by value
+    WeightedVecs<ScalarType, 4> vs_ws = {};
+    for (int j = 0; j < rem; ++j) {
+      vs_ws.ptrs[j] = vs[i + j].GetDevicePointer();  // get the pointer
+      vs_ws.weights[j] = ws[i + j];                  // plain array indexing, not ws(k)
+    }
+    //calculate the linear combination on GPU, handle more than 4 vectors through inc || i > 0
+    LinearCombinationKernel<ScalarType, 4><<<gridDim, blockDim>>>(d_v, vs_ws, rem, nElm, inc || i > 0);
+    gpuErrChk(cudaPeekAtLastError());
+  }
+
+}
+
+template class CSysVector<su2double>; //This is a temporary fix for invalid instantiations due to separating the member function from the header file the class is defined in. Will try to rectify it in coming commits.
 /*--- 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
@@ -205,10 +395,12 @@ DEVICE_EXPRESSION_SHAPES(passivedouble);
 #undef INSTANTIATE_DEVICE_ASSIGN_EXPR
 #undef INSTANTIATE_DEVICE_ASSIGN
 
+#if defined(USE_MIXED_PRECISION)
 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;
+#endif
 
 #if defined(USE_MIXED_PRECISION) && !defined(USE_SINGLE_PRECISION)
 template void CSysVector<passivedouble>::HtDTransfer(bool trigger) const;