Skip to content

Commit 44bf083

Browse files
LwhJessepcarruscag
andauthored
Add initial end-to-end CUDA FGMRES solver path (#2825)
## Proposed Changes This PR adds an end-to-end CUDA linear solve path: the Krylov solvers keep their host control flow, and `CSysVector` operations, the SpMV and the Jacobi preconditioner are dispatched to CUDA kernels when `ENABLE_CUDA=YES`. Transfers are explicit and owned by the object responsible for the data, with no coherency or dirty-flag tracking in `CSysVector` / `CSysMatrix`: - `CSysMatrixVectorProduct` uploads the matrix on construction - the preconditioner uploads its data in `Build()` - `CSysSolve` uploads `b` and `x` and downloads `x` in `HandleTemporariesIn/Out`, which is also where device evaluation is switched on and off - preconditioners without a device implementation (ILU, LU-SGS, Linelet, PaStiX) download the input, apply on the host, and upload the result A solve is therefore fully device resident for Identity and Jacobi preconditioners: **2 uploads and 1 download per linear system**, independent of the Krylov subspace size. Implementation notes: - `cuBLAS` for `dot` / `norm`, custom kernels for the block-LDU SpMV, the Jacobi apply, and `CSysVector` expression assignment - `CSysVector` operands are captured in expressions by value, so an arbitrary expression tree is trivially copyable into the assignment kernel; the required expression shapes are explicitly instantiated in `CSysVectorGPU.cu`, consistent with how `CSysMatrix` is instantiated - device work is issued by a single thread with the OpenMP team synchronized around it, so the GPU path is usable from inside the existing parallel regions; this is internal to the linear algebra layer - the CUDA translation units are only linked into the primal libraries, since they cannot be compiled with the CoDiPack defines; device dispatch is compiled out of the AD builds - adds `LINEAR_SOLVER_PREC= NONE` (identity) ## Related Work Follows the review direction in #2822 (show a working end-to-end GPU linear solve before splitting out infrastructure) and the implementation preferences in #2816. ## Validation - CPU vs GPU on inviscid NACA0012, 25 iterations, RTX 4070 Ti SUPER (sm_89), for FGMRES + `JACOBI`, FGMRES + `NONE`, FGMRES + `ILU` (host preconditioner path) and BCGSTAB. Results agree to the printed precision in double, and to ~6 significant figures in mixed precision. BCGSTAB agrees exactly once the linear system is converged (`LINEAR_SOLVER_ERROR=1e-10`); at loose tolerances the two paths diverge through BCGSTAB's own sensitivity, not a difference in the algebra. - Mixed, normal and single precision builds all compile and run; device dispatch confirmed live in each (kernel launch counts track the subspace size). - `OMP_NUM_THREADS=1` and `4` give bit-identical results on the GPU path. - Builds verified: primal, primal + AD + directdiff, `enable-cuda` + `with-omp`, mixed / normal / single precision. - Transfer counts measured per solve: 2 H2D + 1 D2H + 1 matrix upload for Jacobi and `NONE`, plus one download/upload pair per preconditioner application for ILU. Earlier validation of the original design (6 representative cases, `nsys` / `ncu` profiling) predates the rework of the transfer and dispatch model and should be repeated. ## PR Checklist - [x] I am submitting my contribution to the develop branch. - [x] My contribution generates no new compiler warnings (try with --warnlevel=3 when using meson). - [x] My contribution is commented and consistent with SU2 style (https://su2code.github.io/docs_v7/Style-Guide/). - [x] I used the pre-commit hook to prevent dirty commits and used `pre-commit run --all` to format old commits. - [ ] I have added a test case that demonstrates my contribution, if necessary. - [ ] I have updated appropriate documentation (Tutorials, Docs Page, config_template.cpp), if necessary. --------- Co-authored-by: Pedro Gomes <pcarruscag@gmail.com>
1 parent de8c501 commit 44bf083

21 files changed

Lines changed: 808 additions & 101 deletions

AUTHORS.md

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,18 +10,18 @@ Thomas D. Economon (formerly Stanford University)
1010
Juan J. Alonso (Stanford University)
1111
```
1212

13-
## Current Maintainer ##
13+
## Current Maintainers ##
1414

1515
The SU2 project is maintained by members of the SU2 Foundation (https://su2foundation.org)
1616

1717
```
18-
Thomas D. Economon - Executive Director - tom@su2foundation.org
19-
Tim Albring - Director - tim@su2foundation.org
20-
Juan J. Alonso - Director - juan@su2foundation.org
21-
Eran Arad - Director - eran@su2foundation.org
22-
Piero Colonna - Director - piero@su2foundation.org
23-
Pedro Gomes - Director - pedro@su2foundation.org
24-
Daniel Mayer - Director - daniel@su2foundation.org
18+
Thomas D. Economon - Chairperson
19+
Matteo Pini - Vice Chairperson
20+
Nijso Beishuizen - Treasurer
21+
Pedro Gomes - Development Officer
22+
Giulio Gori - Secretary
23+
Nitish Anand - Editorial Officer
24+
Edwin van der Weide - Events Officer
2525
```
2626

2727
in collaboration with the following main contributors and research teams:
@@ -93,6 +93,7 @@ Jairo Paes Cavalcante Filho
9393
Jason Howison
9494
Jayant Mukhopadhaya
9595
Jeffrey van Oostrom
96+
Jesse Li
9697
Jessie Lauzon
9798
João Loureiro
9899
Johannes Blühdorn

Common/include/code_config.hpp

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,16 @@ FORCEINLINE Out su2staticcast_p(In ptr) {
9696
#define HAVE_OMP
9797
#endif
9898

99+
/*--- Detect whether the CUDA kernels are part of this build. The .cu translation units
100+
* cannot be compiled with the CoDiPack defines (nvcc's device pass cannot parse the tape
101+
* machinery), and an object compiled with a different definition of su2double must not be
102+
* linked into an AD library. They are therefore only built into the primal libraries, and
103+
* all device dispatch has to be compiled out of the AD builds, which HAVE_CUDA alone does
104+
* not do because su2mixedfloat is a passive type there as well. ---*/
105+
#if defined(HAVE_CUDA) && !defined(CODI_REVERSE_TYPE) && !defined(CODI_FORWARD_TYPE)
106+
#define SU2_ENABLE_CUDA_KERNELS
107+
#endif
108+
99109
/*--- No full single precision for AD builds. ---*/
100110
#if (defined(CODI_REVERSE_TYPE) || defined(CODI_FORWARD_TYPE)) && defined(USE_SINGLE_PRECISION)
101111
#undef USE_SINGLE_PRECISION

Common/include/linear_algebra/CMatrixVectorProduct.hpp

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,6 @@ class CSysMatrixVectorProduct final : public CMatrixVectorProduct<ScalarType> {
7272
const CSysMatrix<ScalarType>& matrix; /*!< \brief pointer to matrix that defines the product. */
7373
CGeometry* geometry; /*!< \brief geometry associated with the matrix. */
7474
const CConfig* config; /*!< \brief config of the problem. */
75-
mutable bool matrix_uploaded = false; /*!< \brief Upload the matrix lazily on the first actual GPU matvec. */
7675

7776
public:
7877
/*!
@@ -83,7 +82,17 @@ class CSysMatrixVectorProduct final : public CMatrixVectorProduct<ScalarType> {
8382
*/
8483
inline CSysMatrixVectorProduct(const CSysMatrix<ScalarType>& matrix_ref, CGeometry* geometry_ref,
8584
const CConfig* config_ref)
86-
: matrix(matrix_ref), geometry(geometry_ref), config(config_ref) {}
85+
: matrix(matrix_ref), geometry(geometry_ref), config(config_ref) {
86+
/*--- The matrix does not change while this object lives, so it crosses the bus once,
87+
* here. The vectors are uploaded by CSysSolve, see HandleTemporariesIn. ---*/
88+
#ifdef SU2_ENABLE_CUDA_KERNELS
89+
if constexpr (su2_gpu_capable_v<ScalarType>) {
90+
if (config->GetCUDA()) {
91+
SU2_DEVICE_REGION(matrix.HtDTransfer();)
92+
}
93+
}
94+
#endif
95+
}
8796

8897
/*!
8998
* \note This class cannot be default constructed as that would leave us with invalid pointers.
@@ -97,12 +106,19 @@ class CSysMatrixVectorProduct final : public CMatrixVectorProduct<ScalarType> {
97106
*/
98107
inline void operator()(const CSysVector<ScalarType>& u, CSysVector<ScalarType>& v) const override {
99108
if (config->GetCUDA()) {
100-
#ifdef HAVE_CUDA
101-
if (!matrix_uploaded) {
102-
matrix.HtDTransfer();
103-
matrix_uploaded = true;
109+
#ifdef SU2_ENABLE_CUDA_KERNELS
110+
if constexpr (su2_gpu_capable_v<ScalarType>) {
111+
BEGIN_SU2_DEVICE_REGION
112+
matrix.GPUMatrixVectorProduct(u, v, geometry, config);
113+
END_SU2_DEVICE_REGION
114+
} else {
115+
SU2_MPI::Error("GPU acceleration is not supported for AD scalar types.", CURRENT_FUNCTION);
104116
}
105-
matrix.GPUMatrixVectorProduct(u, v, geometry, config);
117+
#elif defined(HAVE_CUDA)
118+
SU2_MPI::Error(
119+
"\nError in launching Matrix-Vector Product Function\nENABLE_CUDA is set to YES\nThe GPU kernels are not "
120+
"part of the AD libraries, use the primal build for GPU acceleration",
121+
CURRENT_FUNCTION);
106122
#else
107123
SU2_MPI::Error(
108124
"\nError in launching Matrix-Vector Product Function\nENABLE_CUDA is set to YES\nPlease compile with CUDA "

Common/include/linear_algebra/CPreconditioner.hpp

Lines changed: 47 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,33 @@
3737
/// \addtogroup SpLinSys
3838
/// @{
3939

40+
/*!
41+
* \brief Applies a preconditioner that only has a host implementation to vectors that live
42+
* on the device: bring the input down, apply, put the result back.
43+
* \note This is what keeps ILU, LU-SGS, Linelet and PaStiX usable on the GPU path. The
44+
* transfers are issued by one thread with the team synchronized around them, the apply
45+
* itself is the normal OpenMP parallel host code.
46+
*/
47+
template <class ScalarType, class Apply>
48+
inline void ApplyPreconditionerOnHost(const CSysVector<ScalarType>& u, CSysVector<ScalarType>& v, Apply&& apply) {
49+
#ifdef SU2_ENABLE_CUDA_KERNELS
50+
if constexpr (su2_gpu_capable_v<ScalarType>) {
51+
if (VecExpr::UseDeviceExpressions()) {
52+
/*--- The host code must not see the device pointers of any expression it builds, so
53+
* the switch is flipped for the duration of the apply. It is written inside the
54+
* regions, by one thread, and published to the team by the trailing barrier. ---*/
55+
SU2_DEVICE_REGION(u.DtHTransfer(); VecExpr::SetUseDeviceExpressions(false);)
56+
57+
apply();
58+
59+
SU2_DEVICE_REGION(VecExpr::SetUseDeviceExpressions(true); v.HtDTransfer();)
60+
return;
61+
}
62+
}
63+
#endif
64+
apply();
65+
}
66+
4067
/*!
4168
* \class CPreconditioner
4269
* \brief Abstract base class for defining a preconditioning operation.
@@ -77,6 +104,18 @@ class CPreconditioner {
77104
template <class ScalarType>
78105
CPreconditioner<ScalarType>::~CPreconditioner() {}
79106

107+
/*!
108+
* \class CIdentityPreconditioner
109+
* \brief No-op preconditioner used when Krylov solvers run without preconditioning.
110+
*/
111+
template <class ScalarType>
112+
class CIdentityPreconditioner final : public CPreconditioner<ScalarType> {
113+
public:
114+
inline void operator()(const CSysVector<ScalarType>& u, CSysVector<ScalarType>& v) const override { v = u; }
115+
116+
inline bool IsIdentity() const override { return true; }
117+
};
118+
80119
/*!
81120
* \class CJacobiPreconditioner
82121
* \brief Specialization of preconditioner that uses CSysMatrix class.
@@ -160,7 +199,7 @@ class CILUPreconditioner final : public CPreconditioner<ScalarType> {
160199
* \param[out] v - CSysVector that is the result of the preconditioning.
161200
*/
162201
inline void operator()(const CSysVector<ScalarType>& u, CSysVector<ScalarType>& v) const override {
163-
sparse_matrix.ComputeILUPreconditioner(u, v, geometry, config);
202+
ApplyPreconditionerOnHost(u, v, [&] { sparse_matrix.ComputeILUPreconditioner(u, v, geometry, config); });
164203
}
165204

166205
/*!
@@ -206,7 +245,7 @@ class CLU_SGSPreconditioner final : public CPreconditioner<ScalarType> {
206245
* \param[out] v - CSysVector that is the result of the preconditioning.
207246
*/
208247
inline void operator()(const CSysVector<ScalarType>& u, CSysVector<ScalarType>& v) const override {
209-
sparse_matrix.ComputeLU_SGSPreconditioner(u, v, geometry, config);
248+
ApplyPreconditionerOnHost(u, v, [&] { sparse_matrix.ComputeLU_SGSPreconditioner(u, v, geometry, config); });
210249
}
211250
};
212251

@@ -234,7 +273,7 @@ class CQuantizedLUSGSPreconditioner final : public CPreconditioner<ScalarType> {
234273
CQuantizedLUSGSPreconditioner() = delete;
235274

236275
inline void operator()(const CSysVector<ScalarType>& u, CSysVector<ScalarType>& v) const override {
237-
sparse_matrix.ComputeLU_SGSPreconditioner(u, v, geometry, config);
276+
ApplyPreconditionerOnHost(u, v, [&] { sparse_matrix.ComputeLU_SGSPreconditioner(u, v, geometry, config); });
238277
}
239278

240279
/*! \brief Quantize the diagonal blocks (off diagonals are quantized on the fly). */
@@ -278,7 +317,7 @@ class CLineletPreconditioner final : public CPreconditioner<ScalarType> {
278317
* \param[out] v - CSysVector that is the result of the preconditioning.
279318
*/
280319
inline void operator()(const CSysVector<ScalarType>& u, CSysVector<ScalarType>& v) const override {
281-
sparse_matrix.ComputeLineletPreconditioner(u, v, geometry, config);
320+
ApplyPreconditionerOnHost(u, v, [&] { sparse_matrix.ComputeLineletPreconditioner(u, v, geometry, config); });
282321
}
283322

284323
/*!
@@ -328,7 +367,7 @@ class CPastixPreconditioner final : public CPreconditioner<ScalarType> {
328367
* \param[out] v - CSysVector that is the result of the preconditioning.
329368
*/
330369
inline void operator()(const CSysVector<ScalarType>& u, CSysVector<ScalarType>& v) const override {
331-
sparse_matrix.ComputePastixPreconditioner(u, v, geometry, config);
370+
ApplyPreconditionerOnHost(u, v, [&] { sparse_matrix.ComputePastixPreconditioner(u, v, geometry, config); });
332371
}
333372

334373
/*!
@@ -363,6 +402,9 @@ CPreconditioner<ScalarType>* CPreconditioner<ScalarType>::Create(ENUM_LINEAR_SOL
363402
CPreconditioner<ScalarType>* prec = nullptr;
364403

365404
switch (kind) {
405+
case IDENTITY:
406+
prec = new CIdentityPreconditioner<ScalarType>();
407+
break;
366408
case JACOBI:
367409
prec = new CJacobiPreconditioner<ScalarType>(jacobian, geometry, config);
368410
break;

Common/include/linear_algebra/CSysMatrix.hpp

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -247,9 +247,10 @@ class CSysMatrix {
247247
unsigned long nnz_u = 0; /*!< \brief Number of U nonzeros. */
248248
};
249249

250-
LDU mat; /*!< \brief Host matrix (values owned via aligned_alloc; pattern from geometry). */
251-
LDU gpu; /*!< \brief Device matrix (all pointers to GPU memory). */
252-
LDU ilu; /*!< \brief ILU factorization, host (values owned; pattern from geometry). */
250+
LDU mat; /*!< \brief Host matrix (values owned via aligned_alloc; pattern from geometry). */
251+
LDU gpu; /*!< \brief Device matrix (all pointers to GPU memory). */
252+
LDU ilu; /*!< \brief ILU factorization, host (values owned; pattern from geometry). */
253+
ScalarType* d_invM = nullptr; /*!< \brief Device inverse diagonal blocks for the Jacobi preconditioner. */
253254

254255
/*--- Quantized off-diagonal storage (used when quantized_mode == true). ---*/
255256
using QuantType = int8_t;
@@ -1100,6 +1101,14 @@ class CSysMatrix {
11001101
void ComputeJacobiPreconditioner(const CSysVector<ScalarType>& vec, CSysVector<ScalarType>& prod, CGeometry* geometry,
11011102
const CConfig* config) const;
11021103

1104+
/*!
1105+
* \brief Apply the Jacobi preconditioner on the GPU/device side.
1106+
* \note This helper is intended as the implementation hook for GPU-resident Krylov solvers.
1107+
* The actual implementation belongs in CSysMatrixGPU.cu.
1108+
*/
1109+
void ComputeJacobiPreconditionerGPU(const CSysVector<ScalarType>& vec, CSysVector<ScalarType>& prod,
1110+
CGeometry* geometry, const CConfig* config) const;
1111+
11031112
/*!
11041113
* \brief Build the ILU preconditioner.
11051114
*/

Common/include/linear_algebra/CSysSolve.hpp

Lines changed: 45 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -240,13 +240,50 @@ class CSysSolve {
240240
*/
241241
void WriteWarning(ScalarType res_calc, ScalarType res_true, ScalarType tol) const;
242242

243+
/*!
244+
* \brief Moves the linear system to the device, if the GPU path is in use.
245+
* \note This and DownloadSolution are the only places where b and x cross the bus. The
246+
* work vectors of the solvers never do, they are allocated on the device and read back
247+
* only through the reductions.
248+
*/
249+
void UploadSystem(bool useCuda) const {
250+
#ifdef SU2_ENABLE_CUDA_KERNELS
251+
if constexpr (su2_gpu_capable_v<ScalarType>) {
252+
if (!useCuda) return;
253+
BEGIN_SU2_DEVICE_REGION {
254+
LinSysRes_ptr->HtDTransfer();
255+
LinSysSol_ptr->HtDTransfer();
256+
VecExpr::SetUseDeviceExpressions(true);
257+
}
258+
END_SU2_DEVICE_REGION
259+
}
260+
#endif
261+
}
262+
263+
/*!
264+
* \brief Brings the solution back from the device and returns to host evaluation.
265+
*/
266+
void DownloadSolution(bool useCuda) const {
267+
#ifdef SU2_ENABLE_CUDA_KERNELS
268+
if constexpr (su2_gpu_capable_v<ScalarType>) {
269+
if (!useCuda) return;
270+
BEGIN_SU2_DEVICE_REGION {
271+
LinSysSol_ptr->DtHTransfer();
272+
VecExpr::SetUseDeviceExpressions(false);
273+
}
274+
END_SU2_DEVICE_REGION
275+
}
276+
#endif
277+
}
278+
243279
/*!
244280
* \brief Used by Solve for compatibility between passive and active CSysVector.
245281
* \param[in] LinSysRes - Linear system residual
246282
* \param[in,out] LinSysSol - Linear system solution
283+
* \param[in] useCuda - Whether to move the system to the device for the solve.
247284
*/
248285
template <class OtherType>
249-
void HandleTemporariesIn(const CSysVector<OtherType>& LinSysRes, CSysVector<OtherType>& LinSysSol) {
286+
void HandleTemporariesIn(const CSysVector<OtherType>& LinSysRes, CSysVector<OtherType>& LinSysSol, bool useCuda) {
250287
SU2_ZONE_SCOPED
251288
if constexpr (std::is_same_v<ScalarType, OtherType>) {
252289
/*--- Same type specialization, temporary variables are not required. ---*/
@@ -255,6 +292,7 @@ class CSysSolve {
255292
LinSysSol_ptr = &LinSysSol;
256293
}
257294
END_SU2_OMP_SAFE_GLOBAL_ACCESS
295+
UploadSystem(useCuda);
258296
} else {
259297
/*--- Copy data, the solution is also copied as it serves as initial condition. ---*/
260298
LinSysRes_tmp.PassiveCopy(LinSysRes);
@@ -266,16 +304,19 @@ class CSysSolve {
266304
LinSysSol_ptr = &LinSysSol_tmp;
267305
}
268306
END_SU2_OMP_SAFE_GLOBAL_ACCESS
307+
UploadSystem(useCuda);
269308
}
270309
}
271310

272311
/*!
273312
* \brief Used by Solve for compatibility between passive and active CSysVector.
274313
* \param[out] LinSysSol - Linear system solution
314+
* \param[in] useCuda - Whether the system was solved on the device.
275315
*/
276316
template <class OtherType>
277-
void HandleTemporariesOut(CSysVector<OtherType>& LinSysSol) {
317+
void HandleTemporariesOut(CSysVector<OtherType>& LinSysSol, bool useCuda) {
278318
SU2_ZONE_SCOPED
319+
DownloadSolution(useCuda);
279320
if constexpr (std::is_same_v<ScalarType, OtherType>) {
280321
/*--- Same type specialization, temporary variables are not required. ---*/
281322
BEGIN_SU2_OMP_SAFE_GLOBAL_ACCESS {
@@ -432,9 +473,9 @@ class CSysSolve {
432473
template <class OtherType, su2enable_if<!std::is_same_v<ScalarType, OtherType>> = 0>
433474
unsigned long Solve_b(MatrixType& Jacobian, const CSysVector<OtherType>& LinSysRes, CSysVector<OtherType>& LinSysSol,
434475
CGeometry* geometry, const CConfig* config, bool directCall = true) {
435-
HandleTemporariesIn(LinSysRes, LinSysSol);
476+
HandleTemporariesIn(LinSysRes, LinSysSol, false);
436477
auto iter = Solve_b(Jacobian, *LinSysRes_ptr, *LinSysSol_ptr, geometry, config, directCall);
437-
HandleTemporariesOut(LinSysSol);
478+
HandleTemporariesOut(LinSysSol, false);
438479
return iter;
439480
}
440481

0 commit comments

Comments
 (0)