diff --git a/Common/include/linear_algebra/CSysVector.hpp b/Common/include/linear_algebra/CSysVector.hpp index 2ffbb9a7967..1b57a6ffe94 100644 --- a/Common/include/linear_algebra/CSysVector.hpp +++ b/Common/include/linear_algebra/CSysVector.hpp @@ -270,6 +270,17 @@ class CSysVector : public VecExpr::CVecExpr, ScalarType> return *this; } + /*! + * \brief GPU helper for `dot`. + */ + ScalarType dotGPU(const CSysVector& other) const; + + /*! + * \brief GPU helper for multiDot. + */ + static su2matrix multiDotGPU(const std::vector>& V, size_t i0, size_t n, + const std::vector>& W, size_t m); + public: static constexpr bool StoreAsRef = true; /*! \brief Required by CVecExpr. */ @@ -394,21 +405,6 @@ class CSysVector : public VecExpr::CVecExpr, 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, 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]; } } diff --git a/Common/src/linear_algebra/CSysVector.cpp b/Common/src/linear_algebra/CSysVector.cpp index 54e47157d55..22449c3759d 100644 --- a/Common/src/linear_algebra/CSysVector.cpp +++ b/Common/src/linear_algebra/CSysVector.cpp @@ -76,66 +76,69 @@ const su2matrix& CSysVector::multiDot(const std::vector< const std::vector>& W, const size_t m) { SU2_ZONE_SCOPED - static constexpr size_t BLOCK_SIZE = 1024; + static su2matrix shared; if (n == 0 || m == 0) return shared; + su2matrix local; + + if (VecExpr::UseDeviceExpressions()) { #ifdef SU2_ENABLE_CUDA_KERNELS - if constexpr (su2_gpu_capable_v) { - 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) { + 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 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); diff --git a/Common/src/linear_algebra/CSysVectorGPU.cu b/Common/src/linear_algebra/CSysVectorGPU.cu index 01b146baf76..ba53b0cbc29 100644 --- a/Common/src/linear_algebra/CSysVectorGPU.cu +++ b/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::DtHTransfer(bool trigger) const { } template -ScalarType CSysVector::GPUDot(const CSysVector& other) const { +ScalarType CSysVector::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::GPUDot(const CSysVector& other) const { status = cublasDdot(handle, static_cast(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::GPUDot(const CSysVector& other) const { return global_dot; } +/*! + * \brief multi vector product with cublasgemmBatched + */ template -ScalarType CSysVector::GPUNorm() const { - SU2_ZONE_SCOPED - return sqrt(GPUDot(*this)); +su2matrix CSysVector::multiDotGPU(const std::vector>& V, const size_t i0, + const size_t n, const std::vector>& 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 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 h_A, h_B; + static std::vector 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) { + status = cublasSgemmBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, 1, 1, size, &alpha, ws.d_A, static_cast(size), + ws.d_B, static_cast(size), &beta, ws.d_C, 1, static_cast(batch)); + } else if constexpr (std::is_same_v) { + status = cublasDgemmBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, 1, 1, size, &alpha, ws.d_A, static_cast(size), + ws.d_B, static_cast(size), &beta, ws.d_C, 1, static_cast(batch)); + } else { + SU2_MPI::Error("Unsupported ScalarType in CSysVector::multiDotGPU.", CURRENT_FUNCTION); + return local; + } + + if (status != CUBLAS_STATUS_SUCCESS) { + SU2_MPI::Error("cuBLAS cublasgemmBatched 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::HtDTransfer(bool trigger) const; template void CSysVector::DtHTransfer(bool trigger) const; -template su2mixedfloat CSysVector::GPUDot(const CSysVector& other) const; -template su2mixedfloat CSysVector::GPUNorm() const; +template su2mixedfloat CSysVector::dotGPU(const CSysVector& other) const; +template su2matrix CSysVector::multiDotGPU( + const std::vector>& V, size_t i0, size_t n, + const std::vector>& W, size_t m); #if defined(USE_MIXED_PRECISION) && !defined(USE_SINGLE_PRECISION) template void CSysVector::HtDTransfer(bool trigger) const; template void CSysVector::DtHTransfer(bool trigger) const; -template passivedouble CSysVector::GPUDot(const CSysVector& other) const; -template passivedouble CSysVector::GPUNorm() const; +template passivedouble CSysVector::dotGPU(const CSysVector& other) const; +template su2matrix CSysVector::multiDotGPU( + const std::vector>& V, size_t i0, size_t n, + const std::vector>& W, size_t m); #endif