diff --git a/Common/include/linear_algebra/CSysVector.hpp b/Common/include/linear_algebra/CSysVector.hpp index 2ffbb9a7967..0422761981d 100644 --- a/Common/include/linear_algebra/CSysVector.hpp +++ b/Common/include/linear_algebra/CSysVector.hpp @@ -585,6 +585,19 @@ class CSysVector : public VecExpr::CVecExpr, ScalarType> static const su2matrix& multiDot(const std::vector& V, size_t i0, size_t n, const std::vector& 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& multiDotGPU(const std::vector>& V, const size_t i0, + const size_t n, const std::vector>& W, + const size_t m); + /*! * \brief Squared L2 norm of the vector (via dot with self). * \return Squared L2 norm. 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..31ba7c0ffab 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 @@ -117,6 +117,99 @@ ScalarType CSysVector::GPUNorm() const { return sqrt(GPUDot(*this)); } +/*! + * \brief multi vector product with cublasgemmBatched + */ +template +const 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; + + // allocate persisten result buffer local on host and device, is resized if needed + static su2matrix local; + local.resize(n,m); + static ScalarType* d_local = nullptr; + static size_t local_capacity = 0; + + if (batch > local_capacity) { // if not enough capacity, enlarge by re-allocation on device + if (d_local) gpuErrChk(cudaFree(d_local)); + gpuErrChk(cudaMalloc(&d_local, batch * sizeof(ScalarType))); + local_capacity = batch; + } + // zero out the result buffer + gpuErrChk(cudaMemset(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] = d_local + idx; // C maps to d_local to store the coefficients in the 2D array + } + } + + // prepare device arrays for A,B,C. Re-allocate if batch size has grown + static const ScalarType** d_A = nullptr; + static const ScalarType** d_B = nullptr; + static ScalarType** d_C = nullptr; + static size_t ptrs_capacity = 0; // current capacity + + if (batch > ptrs_capacity) { // if not enough capacity, enlarge by re-allocation on device + if (d_A) gpuErrChk(cudaFree(d_A)); + if (d_B) gpuErrChk(cudaFree(d_B)); + if (d_C) gpuErrChk(cudaFree(d_C)); + gpuErrChk(cudaMalloc(&d_A, batch * sizeof(ScalarType*))); + gpuErrChk(cudaMalloc(&d_B, batch * sizeof(ScalarType*))); + gpuErrChk(cudaMalloc(&d_C, batch * sizeof(ScalarType*))); + ptrs_capacity = batch; // update current capacity + } + // copy pointers to device + gpuErrChk(cudaMemcpy(d_A, h_A.data(), batch * sizeof(ScalarType*), cudaMemcpyHostToDevice)); + gpuErrChk(cudaMemcpy(d_B, h_B.data(), batch * sizeof(ScalarType*), cudaMemcpyHostToDevice)); + gpuErrChk(cudaMemcpy(d_C, h_C.data(), batch * sizeof(ScalarType*), cudaMemcpyHostToDevice)); + + // define alpha = 1.0 and beta = 0.0 + const ScalarType alpha = ScalarType(1.0); + const ScalarType beta = ScalarType(0.0); + + if constexpr (std::is_same_v) { + status = cublasSgemmBatched(handle, CUBLAS_OP_T, CUBLAS_OP_N, 1, 1, size, &alpha, d_A, static_cast(size), + d_B, static_cast(size), &beta, 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, d_A, static_cast(size), + d_B, static_cast(size), &beta, 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(), d_local, batch * sizeof(ScalarType), cudaMemcpyDeviceToHost)); + + return local; + +} + /*--- 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 @@ -209,14 +302,21 @@ 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 const su2matrix& CSysVector::multiDotGPU( + const std::vector>& V, const size_t i0, const size_t n, + const std::vector>& W, const 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 const su2matrix& CSysVector::multiDotGPU( + const std::vector>& V, const size_t i0, const size_t n, + const std::vector>& W, const size_t m); #endif