[GSoC] Feature: LU-SGS preconditioner to GPU

AUTHORS.md

@@ -69,6 +69,7 @@ Christian Bauer
 Clark Pederson
 Daumantas Kavolis
 Dave Taflin
+Davide Di Giusto
 Eduardo Molina
 Edwin van der Weide
 Eitan Aberman

Common/include/linear_algebra/CPreconditioner.hpp

@@ -260,13 +260,13 @@ class CLU_SGSPreconditioner final : public CPreconditioner<ScalarType> {
    * \param[out] v - CSysVector that is the result of the preconditioning.
    */
   inline void operator()(const CSysVector<ScalarType>& u, CSysVector<ScalarType>& v) const override {
-    ApplyPreconditionerOnHost(u, v, [&] { sparse_matrix.ComputeLU_SGSPreconditioner(u, v, geometry, config); });
+    sparse_matrix.ComputeLU_SGSPreconditioner(u, v, geometry, config);
   }
 
   /*!
    * \note Also serves Q_LU_SGS: quantizes the diagonal blocks, no-op for plain LU_SGS.
    */
-  inline void Build() override { sparse_matrix.QuantizeDiagonalBlocks(); }
+  inline void Build() override { sparse_matrix.BuildLU_SGSPreconditioner(); }
 };
 
 /*!

Common/include/linear_algebra/CSysMatrix.hpp

@@ -303,7 +303,7 @@ class CSysMatrix {
   LDU<ScalarType> gpu;          /*!< \brief Device matrix (all pointers to GPU memory). */
   LDU<ScalarType> ilu;          /*!< \brief ILU factorization, host (values owned; pattern from geometry). */
   LDU<ScalarType> gpu_ilu;      /*!< \brief ILU factorization, device (values and pattern in GPU memory). */
-  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 or LU-SGS preconditioner. */
 
   /*--- Quantized off-diagonal storage (used when quantized_mode == true). ---*/
   using QuantType = int8_t;
@@ -356,7 +356,7 @@ class CSysMatrix {
    * rows in level k only depend on rows in levels < k. The same table drives the forward
    * (increasing level) and backward (decreasing level) substitution, because the U pattern is
    * the transpose of the L pattern. Used directly by the host/OMP substitution, and flattened
-   * into ilu_level_ptr / d_ilu_level_idx below for the GPU triangular solves. */
+   * into ilu_level_ptr / d_precond_level_idx below for the GPU triangular solves. */
   CCompressedSparsePatternUL levels_ilu;
 
   /*!< \brief Coloring of the (domain-only) ILU dependency graph, used only by the GPU iterative
@@ -373,19 +373,22 @@ class CSysMatrix {
   vector<su2uint> ilu_color_ptr;      /*!< \brief Start of each color in d_ilu_color_idx, size nColors+1. */
   su2uint* d_ilu_color_idx = nullptr; /*!< \brief Row indices, grouped by color. */
 
-  vector<su2uint> ilu_level_ptr;      /*!< \brief Start of each level in d_ilu_level_idx, size nLevels+1. */
-  su2uint* d_ilu_level_idx = nullptr; /*!< \brief Row indices, grouped by level. */
+  vector<su2uint> precond_level_ptr;      /*!< \brief Start of each level in d_precond_level_idx, size nLevels+1. */
+  su2uint* d_precond_level_idx = nullptr; /*!< \brief Row indices, grouped by level. */
 
   /*--- The per-color (factorization) and per-level (triangular solves) kernel launch sequences
    * are identical on every call: same grid/block sizes, same device pointers (all fixed members,
    * allocated once). Each is captured once into a CUDA graph and replayed to remove
    * host-side launch overhead without changing the parallelization. ---*/
   mutable struct CUgraphExec_st* ilu_build_graph_exec = nullptr;
-  mutable struct CUgraphExec_st* ilu_apply_graph_exec = nullptr;
-  mutable const ScalarType* ilu_apply_graph_vec = nullptr; /*!< \brief Pointers the apply graph
-                                                            * was captured with, to detect when
-                                                            * it must be recaptured. */
-  mutable ScalarType* ilu_apply_graph_prod = nullptr;
+  mutable struct CUgraphExec_st* precond_fwd_graph_exec = nullptr;  // ILU or LU-SGS forward only
+  mutable struct CUgraphExec_st* precond_bwd_graph_exec = nullptr;  // LU-SGS backward only
+  mutable const ScalarType* precond_fwd_graph_vec = nullptr;        /*!< \brief Pointers the apply graph
+                                                                     * was captured with, to detect when
+                                                                     * it must be recaptured. */
+  mutable ScalarType* precond_fwd_graph_prod = nullptr;
+  mutable ScalarType* precond_bwd_graph_prod = nullptr;
+
   /*--- Non-default stream, needed for two mutually exclusive uses that never overlap on a given
    * matrix (quantized_mode and ILU are alternative preconditioner choices, decided once in
    * Initialize()): (1) the ILU build/apply CUDA graphs below, since the legacy default stream
@@ -657,6 +660,31 @@ class CSysMatrix {
    */
   void ComputeILUPreconditionerGPU(const CSysVector<ScalarType>& vec, CSysVector<ScalarType>& prod) const;
 
+  /*!
+   * \brief Build the LU-SGS preconditioner on the device
+   */
+  void BuildLU_SGSPreconditionerGPU();
+
+  /*!
+   * \brief Apply the LU-SGS preconditioner forward pass on the device
+   */
+  void ComputeLU_SGSForwardGPU(const CSysVector<ScalarType>& vec, CSysVector<ScalarType>& prod) const;
+
+  /*!
+   * \brief Apply the LU-SGS preconditioner backward pass on the device
+   */
+  void ComputeLU_SGSBackwardGPU(CSysVector<ScalarType>& prod) const;
+
+  /*!
+   * \brief Apply the forward pass of the LU-SGS preconditioner
+   */
+  void ComputeLU_SGSPreconditionerForward(const CSysVector<ScalarType>& vec, CSysVector<ScalarType>& prod) const;
+
+  /*!
+   * \brief Apply the backward pass of the LU-SGS preconditioner
+   */
+  void ComputeLU_SGSPreconditionerBackward(CSysVector<ScalarType>& prod) const;
+
  public:
   /*!
    * \brief Constructor of the class.
@@ -1230,6 +1258,11 @@ class CSysMatrix {
   void ComputeILUPreconditioner(const CSysVector<ScalarType>& vec, CSysVector<ScalarType>& prod, CGeometry* geometry,
                                 const CConfig* config) const;
 
+  /*!
+   * \brief Build the LU-SGS preconditioner.
+   */
+  void BuildLU_SGSPreconditioner();
+
   /*!
    * \brief Multiply CSysVector by the preconditioner
    * \param[in] vec - CSysVector to be multiplied by the preconditioner.

Common/src/linear_algebra/CSysMatrix.cpp

@@ -159,10 +159,11 @@ CSysMatrix<ScalarType>::~CSysMatrix() {
     freeLDU(d_q_blocks);
     GPUMemoryAllocation::gpu_free(d_invM);
     GPUMemoryAllocation::gpu_free(d_ilu_color_idx);
-    GPUMemoryAllocation::gpu_free(d_ilu_level_idx);
+    GPUMemoryAllocation::gpu_free(d_precond_level_idx);
 #ifdef SU2_ENABLE_CUDA_KERNELS
     if (ilu_build_graph_exec != nullptr) cudaGraphExecDestroy(ilu_build_graph_exec);
-    if (ilu_apply_graph_exec != nullptr) cudaGraphExecDestroy(ilu_apply_graph_exec);
+    if (precond_fwd_graph_exec != nullptr) cudaGraphExecDestroy(precond_fwd_graph_exec);
+    if (precond_bwd_graph_exec != nullptr) cudaGraphExecDestroy(precond_bwd_graph_exec);
     if (aux_stream != nullptr) cudaStreamDestroy(aux_stream);
     if (htd_event != nullptr) cudaEventDestroy(htd_event);
 #endif
@@ -217,6 +218,7 @@ void CSysMatrix<ScalarType>::Initialize(unsigned long npoint, unsigned long npoi
 
   const bool ilu_needed = (prec == ILU);
   const bool diag_needed = (prec == JACOBI) || (prec == Q_JACOBI) || (prec == LINELET);
+  const bool lu_sgs_on_device = useCuda && (prec == LU_SGS || prec == Q_LU_SGS);
 
   /*--- Linelet also builds the Jacobi preconditioner but reads the inverse diagonal blocks on
    * the host, so only plain (or quantized) Jacobi can keep them exclusively on the device. ---*/
@@ -408,62 +410,72 @@ void CSysMatrix<ScalarType>::Initialize(unsigned long npoint, unsigned long npoi
 
   if (diag_needed) allocAndInit(invM, nPointDomain * nVar * nEqn);
 
-  if (jacobi_on_device) {
+  const bool any_precond_on_device = useCuda && (jacobi_on_device || lu_sgs_on_device || ilu_needed);
+
+  if (any_precond_on_device) {
     if (nVar != nEqn) {
-      SU2_MPI::Error("CUDA Jacobi preconditioner requires square blocks.", CURRENT_FUNCTION);
+      SU2_MPI::Error("CUDA preconditioners require square blocks.", CURRENT_FUNCTION);
     }
     if (nVar * nVar > 1024) {
-      SU2_MPI::Error("CUDA Jacobi preconditioner uses one thread per block entry, nVar is too large.",
-                     CURRENT_FUNCTION);
+      SU2_MPI::Error("CUDA preconditioners use one thread per block entry, nVar is too large.", CURRENT_FUNCTION);
     }
-    d_invM = GPUMemoryAllocation::gpu_alloc<ScalarType, true>(nPointDomain * nVar * nEqn * sizeof(ScalarType));
-  }
 
-  if (useCuda && ilu_needed) {
-    if (nVar != nEqn) {
-      SU2_MPI::Error("CUDA ILU factorization requires square blocks.", CURRENT_FUNCTION);
-    }
-    if (nVar * nVar > 1024) {
-      SU2_MPI::Error("CUDA ILU factorization uses one thread per block entry, nVar is too large.", CURRENT_FUNCTION);
+    if (jacobi_on_device || lu_sgs_on_device) {
+      d_invM = GPUMemoryAllocation::gpu_alloc<ScalarType, true>(nPointDomain * nVar * nEqn * sizeof(ScalarType));
     }
-    /*--- The factors are built and used on the device, only the pattern and the level table
-     * are uploaded (once, here) because they do not change. ---*/
-    gpu_ilu.nnz_l = ilu.nnz_l;
-    gpu_ilu.nnz_u = ilu.nnz_u;
-    GPUAllocAndInit(gpu_ilu.d, nPointDomain * nVar * nEqn);
-    GPUAllocAndInit(gpu_ilu.l, ilu.nnz_l * nVar * nEqn);
-    GPUAllocAndInit(gpu_ilu.u, ilu.nnz_u * nVar * nEqn);
-    GPUAllocAndCopy(gpu_ilu.row_ptr_l, ilu.row_ptr_l, nPointDomain + 1);
-    GPUAllocAndCopy(gpu_ilu.col_ind_l, ilu.col_ind_l, ilu.nnz_l);
-    GPUAllocAndCopy(gpu_ilu.row_ptr_u, ilu.row_ptr_u, nPointDomain + 1);
-    GPUAllocAndCopy(gpu_ilu.col_ind_u, ilu.col_ind_u, ilu.nnz_u);
-
-    /*--- Flatten the coloring, the index type differs from the one of the pattern. It drives
-     * the factorization on the device. ---*/
-    std::vector<su2uint> color_idx;
-    color_idx.reserve(nPointDomain);
-    ilu_color_ptr.clear();
-    ilu_color_ptr.push_back(0);
-    for (auto color = 0ul; color < color_ilu.getOuterSize(); ++color) {
-      for (auto k = 0ul; k < color_ilu.getNumNonZeros(color); ++k) {
-        color_idx.push_back(static_cast<su2uint>(color_ilu.getInnerIdx(color, k)));
+
+    /*--- Flattens a grouped sparse pattern (levels, colors) into a host ptr and device index arrays.
+     * Used in ILU levels and colors, and LU-SGS levels---*/
+    auto FlattenGroupToDevice = [](const auto& grouped, std::vector<su2uint>& group_ptr, unsigned long reserveHint,
+                                   unsigned long bound = ~0ul) {
+      std::vector<su2uint> flat_idx;
+      flat_idx.reserve(reserveHint);
+      group_ptr.clear();
+      group_ptr.push_back(0);
+      for (auto group = 0ul; group < grouped.getOuterSize(); ++group) {
+        for (auto k = 0ul; k < grouped.getNumNonZeros(group); ++k) {
+          auto idx = grouped.getInnerIdx(group, k);
+          if (static_cast<unsigned long>(idx) >= bound)
+            continue;  // prevent out of bounds in LU-SGS kernels if more than 1 mpi task
+          flat_idx.push_back(static_cast<su2uint>(idx));
+        }
+        group_ptr.push_back(static_cast<su2uint>(flat_idx.size()));
       }
-      ilu_color_ptr.push_back(static_cast<su2uint>(color_idx.size()));
+      return GPUMemoryAllocation::gpu_alloc_cpy(flat_idx.data(), flat_idx.size() * sizeof(su2uint));
+    };
+
+    if (lu_sgs_on_device) {
+      // get the zero-filled sparse pattern for the LU-SGS
+      const auto& pat_lusgs = geometry->GetSparsePattern(type, 0);
+
+      /*--- Compute the levels using the lower pattern for the forward pass and
+       * reverse the levels for the backward pass. This works if L and U are symmetric, to be verified ---*/
+      auto levels_lusgs = computeLevels(pat_lusgs.l);
+
+      /*--- Flatten levels_lusgs. It drives both triangular solves on the device. ---*/
+      d_precond_level_idx = FlattenGroupToDevice(levels_lusgs, precond_level_ptr, nPointDomain, nPointDomain);
     }
-    d_ilu_color_idx = GPUMemoryAllocation::gpu_alloc_cpy(color_idx.data(), color_idx.size() * sizeof(su2uint));
 
-    /*--- Flatten levels_ilu the same way. It drives both triangular solves on the device. ---*/
-    std::vector<su2uint> level_idx;
-    level_idx.reserve(nPointDomain);
-    ilu_level_ptr.clear();
-    ilu_level_ptr.push_back(0);
-    for (auto level = 0ul; level < levels_ilu.getOuterSize(); ++level) {
-      for (auto k = 0ul; k < levels_ilu.getNumNonZeros(level); ++k) {
-        level_idx.push_back(static_cast<su2uint>(levels_ilu.getInnerIdx(level, k)));
-      }
-      ilu_level_ptr.push_back(static_cast<su2uint>(level_idx.size()));
+    if (ilu_needed) {
+      /*--- The factors are built and used on the device, only the pattern and the level table
+       * are uploaded (once, here) because they do not change. ---*/
+      gpu_ilu.nnz_l = ilu.nnz_l;
+      gpu_ilu.nnz_u = ilu.nnz_u;
+      GPUAllocAndInit(gpu_ilu.d, nPointDomain * nVar * nEqn);
+      GPUAllocAndInit(gpu_ilu.l, ilu.nnz_l * nVar * nEqn);
+      GPUAllocAndInit(gpu_ilu.u, ilu.nnz_u * nVar * nEqn);
+      GPUAllocAndCopy(gpu_ilu.row_ptr_l, ilu.row_ptr_l, nPointDomain + 1);
+      GPUAllocAndCopy(gpu_ilu.col_ind_l, ilu.col_ind_l, ilu.nnz_l);
+      GPUAllocAndCopy(gpu_ilu.row_ptr_u, ilu.row_ptr_u, nPointDomain + 1);
+      GPUAllocAndCopy(gpu_ilu.col_ind_u, ilu.col_ind_u, ilu.nnz_u);
+
+      /*--- Flatten the coloring, the index type differs from the one of the pattern. It drives
+       * the factorization on the device. ---*/
+      d_ilu_color_idx = FlattenGroupToDevice(color_ilu, ilu_color_ptr, nPointDomain);
+
+      /*--- Flatten levels_ilu the same way. It drives both triangular solves on the device. ---*/
+      d_precond_level_idx = FlattenGroupToDevice(levels_ilu, precond_level_ptr, nPointDomain);
     }
-    d_ilu_level_idx = GPUMemoryAllocation::gpu_alloc_cpy(level_idx.data(), level_idx.size() * sizeof(su2uint));
   }
 
   /*--- Thread parallel initialization. ---*/
@@ -1284,11 +1296,69 @@ void CSysMatrix<ScalarType>::ComputeILUPreconditioner(const CSysVector<ScalarTyp
   CSysMatrixComms::Complete(prod, geometry, config);
 }
 
+template <class ScalarType>
+void CSysMatrix<ScalarType>::BuildLU_SGSPreconditioner() {
+  SU2_ZONE_SCOPED
+
+  /*--- Quantize diagonal blocks if mode is active ---*/
+  QuantizeDiagonalBlocks();
+
+  /*--- if on GPU, precompute the inverse of the diagonal D. Otherwise, this is a no-op ---*/
+  if (useCuda) {
+#ifdef SU2_ENABLE_CUDA_KERNELS
+    if constexpr (su2_gpu_capable_v<ScalarType>) {
+      SU2_DEVICE_REGION(BuildLU_SGSPreconditionerGPU();)
+      return;
+    } else {
+      GPUNotAvailable(CURRENT_FUNCTION);
+    }
+#else
+    GPUNotAvailable(CURRENT_FUNCTION);
+#endif
+  }
+}
+
 template <class ScalarType>
 void CSysMatrix<ScalarType>::ComputeLU_SGSPreconditioner(const CSysVector<ScalarType>& vec,
                                                          CSysVector<ScalarType>& prod, CGeometry* geometry,
                                                          const CConfig* config) const {
   SU2_ZONE_SCOPED
+
+  /*--- First part of the symmetric iteration: (D+L).x* = b ---*/
+  ComputeLU_SGSPreconditionerForward(vec, prod);
+
+  /*--- MPI Parallelization ---*/
+
+  CSysMatrixComms::Initiate(prod, geometry, config);
+  CSysMatrixComms::Complete(prod, geometry, config);
+
+  /*--- Second part of the symmetric iteration: (D+U).x_(1) = D.x* ---*/
+  ComputeLU_SGSPreconditionerBackward(prod);
+
+  /*--- MPI Parallelization ---*/
+
+  CSysMatrixComms::Initiate(prod, geometry, config);
+  CSysMatrixComms::Complete(prod, geometry, config);
+}
+
+template <class ScalarType>
+void CSysMatrix<ScalarType>::ComputeLU_SGSPreconditionerForward(const CSysVector<ScalarType>& vec,
+                                                                CSysVector<ScalarType>& prod) const {
+  SU2_ZONE_SCOPED
+
+  if (useCuda) {
+#ifdef SU2_ENABLE_CUDA_KERNELS
+    if constexpr (su2_gpu_capable_v<ScalarType>) {
+      SU2_DEVICE_REGION(ComputeLU_SGSForwardGPU(vec, prod);)
+      return;
+    } else {
+      GPUNotAvailable(CURRENT_FUNCTION);
+    }
+#else
+    GPUNotAvailable(CURRENT_FUNCTION);
+#endif
+  }
+
   /*--- First part of the symmetric iteration: (D+L).x* = b ---*/
 
   /*--- Coherent view of vectors. ---*/
@@ -1325,14 +1395,27 @@ void CSysMatrix<ScalarType>::ComputeLU_SGSPreconditioner(const CSysVector<Scalar
     }
   }
   END_SU2_OMP_FOR
+}
 
-  /*--- MPI Parallelization ---*/
-
-  CSysMatrixComms::Initiate(prod, geometry, config);
-  CSysMatrixComms::Complete(prod, geometry, config);
+template <class ScalarType>
+void CSysMatrix<ScalarType>::ComputeLU_SGSPreconditionerBackward(CSysVector<ScalarType>& prod) const {
+  SU2_ZONE_SCOPED
 
   /*--- Second part of the symmetric iteration: (D+U).x_(1) = D.x* ---*/
 
+  if (useCuda) {
+#ifdef SU2_ENABLE_CUDA_KERNELS
+    if constexpr (su2_gpu_capable_v<ScalarType>) {
+      SU2_DEVICE_REGION(ComputeLU_SGSBackwardGPU(prod);)
+      return;
+    } else {
+      GPUNotAvailable(CURRENT_FUNCTION);
+    }
+#else
+    GPUNotAvailable(CURRENT_FUNCTION);
+#endif
+  }
+
   /*--- OpenMP Parallelization ---*/
   SU2_OMP_FOR_STAT(1)
   for (unsigned long thread = 0; thread < omp_num_parts; ++thread) {
@@ -1363,11 +1446,6 @@ void CSysMatrix<ScalarType>::ComputeLU_SGSPreconditioner(const CSysVector<Scalar
     }
   }
   END_SU2_OMP_FOR
-
-  /*--- MPI Parallelization ---*/
-
-  CSysMatrixComms::Initiate(prod, geometry, config);
-  CSysMatrixComms::Complete(prod, geometry, config);
 }
 
 template <class ScalarType>