From b9d39e34053b6d595306e49ceb6f3b8bf3f4f050 Mon Sep 17 00:00:00 2001 From: Ryan Hill <38674843+RyanUnderhill@users.noreply.github.com> Date: Tue, 23 May 2023 17:19:30 -0700 Subject: [PATCH] Fix cuda Transpose bug 16039 (#16042) ### Description Transpose will fail in cuda for FLOAT16 for tensors larger than 1048x1048 due to our optimized case exceeding the cuda grid size of 65536. The fix is to just use our regular cuda transpose in these cases. ### Motivation and Context https://github.com/microsoft/onnxruntime/issues/16039 --- onnxruntime/core/providers/cuda/fpgeneric.cu | 13 +++++++++- .../providers/cuda/shared_inc/fpgeneric.h | 1 + .../core/providers/cuda/tensor/transpose.cc | 2 +- .../providers/rocm/shared_inc/fpgeneric.h | 2 ++ .../providers/cpu/tensor/transpose_test.cc | 26 +++++++++++++++++++ 5 files changed, 42 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/cuda/fpgeneric.cu b/onnxruntime/core/providers/cuda/fpgeneric.cu index 3ae0b9c8eb..ef9eda8a05 100644 --- a/onnxruntime/core/providers/cuda/fpgeneric.cu +++ b/onnxruntime/core/providers/cuda/fpgeneric.cu @@ -64,11 +64,22 @@ __global__ void CopyVectorBFloat16(const onnxruntime::BFloat16* x, int incx, onn } // namespace +dim3 cublasTransposeHelperDimGrid(int m, int n) { + return dim3((n + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM, (m + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM, 1); +} + +// cublasTransposeHelper can only be used if it won't overflow the 65536 grid y dimension size +__host__ bool CanUse_cublasTransposeHelper_MLFloat16(int m, int n) { + dim3 dimGrid = cublasTransposeHelperDimGrid(m, n); + return dimGrid.y < 65536; +} + cublasStatus_t cublasTransposeHelper(cudaStream_t stream, cublasHandle_t, cublasOperation_t, cublasOperation_t, int m, int n, const half*, const half* A, int, const half*, const half*, int, half* C, int) { if (C != A) { - dim3 dimGrid((n + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM, (m + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM, 1); + dim3 dimGrid = cublasTransposeHelperDimGrid(m, n); dim3 dimBlock(TRANS_TILE_DIM, BLOCK_ROWS, 1); + ORT_ENFORCE(dimGrid.y < 65536); // To prevent this, call CanUse_cublasTransposeHelper_MLFloat16 first transposeNoOverlap<<>>(C, A, n, m); } else { return CUBLAS_STATUS_NOT_SUPPORTED; diff --git a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h index 8bf429949e..b37ab3ec30 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h +++ b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h @@ -485,6 +485,7 @@ inline cublasStatus_t cublasTransposeHelper(cudaStream_t, cublasHandle_t handle, return cublasDgeam(handle, transa, transb, m, n, alpha, A, lda, beta, B, ldb, C, ldc); } +bool CanUse_cublasTransposeHelper_MLFloat16(int m, int n); cublasStatus_t cublasTransposeHelper(cudaStream_t, cublasHandle_t, cublasOperation_t, cublasOperation_t, int m, int n, const half*, const half* A, int, const half*, const half*, int, half* C, int); // copy diff --git a/onnxruntime/core/providers/cuda/tensor/transpose.cc b/onnxruntime/core/providers/cuda/tensor/transpose.cc index d056219a89..0e10ebcbef 100644 --- a/onnxruntime/core/providers/cuda/tensor/transpose.cc +++ b/onnxruntime/core/providers/cuda/tensor/transpose.cc @@ -196,7 +196,7 @@ Status Transpose::DoTranspose(const cudaDeviceProp& prop, auto mn = TryTransposeWithCublas(new_permutations, new_input_dims); int M = std::get<0>(mn); int N = std::get<1>(mn); - if (M != 0 && N != 0) { + if (M != 0 && N != 0 && (element_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 || CanUse_cublasTransposeHelper_MLFloat16(M, N))) { if (element_type == utils::GetONNXTensorElementDataType()) { return TransposeWithCublas(stream, cublas_handle, input, output, M, N); } else if (element_type == utils::GetONNXTensorElementDataType()) { diff --git a/onnxruntime/core/providers/rocm/shared_inc/fpgeneric.h b/onnxruntime/core/providers/rocm/shared_inc/fpgeneric.h index 657a11ccb8..7cbc37cb64 100644 --- a/onnxruntime/core/providers/rocm/shared_inc/fpgeneric.h +++ b/onnxruntime/core/providers/rocm/shared_inc/fpgeneric.h @@ -465,6 +465,8 @@ inline rocblas_status rocblasTransposeHelper(hipStream_t /*stream*/, rocblas_han inline rocblas_status rocblasTransposeHelper(hipStream_t /*stream*/, rocblas_handle handle, rocblas_operation transa, rocblas_operation transb, int m, int n, const double* alpha, const double* A, int lda, const double* beta, const double* B, int ldb, double* C, int ldc) { return rocblas_dgeam(handle, transa, transb, m, n, alpha, A, lda, beta, B, ldb, C, ldc); } + +inline bool CanUse_rocblasTransposeHelper_MLFloat16(int /*m*/, int /*n*/) { return true; } // CUDA has a limited grid size of 65536, ROCm has higher limits. rocblas_status rocblasTransposeHelper(hipStream_t stream, rocblas_handle, rocblas_operation, rocblas_operation, int m, int n, const half*, const half* A, int, const half*, const half*, int, half* C, int); // copy diff --git a/onnxruntime/test/providers/cpu/tensor/transpose_test.cc b/onnxruntime/test/providers/cpu/tensor/transpose_test.cc index f07d4f8874..5caa047cf2 100644 --- a/onnxruntime/test/providers/cpu/tensor/transpose_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/transpose_test.cc @@ -755,6 +755,32 @@ TEST(TransposeOpTest, Transpose3DImpl) { } } +static void TestTransposeMLFloat16( + const std::vector& perm, + const std::vector& x_dims, + const std::vector& y_dims, + double error_tolerance = 1e-4) { + CompareOpTester test{"Transpose"}; + + RandomValueGenerator random{}; + const auto X_data = random.Uniform(x_dims, -10.0, 10.0); + std::vector X_data_f16 = FloatsToMLFloat16s(X_data); + const auto Y_data = FillZeros(y_dims); + std::vector Y_data_f16 = FloatsToMLFloat16s(Y_data); + + test.AddAttribute("perm", perm); + test.AddInput("X", x_dims, X_data_f16); + test.AddOutput("Y", y_dims, Y_data_f16); + + test.CompareWithCPU(kGpuExecutionProvider, error_tolerance); +} + +TEST(TransposeOpTest, TransposeBigMLFloat16) { // Exercises CanUse_cublasTransposeHelper_MLFloat16 (cuda 65535 grid dimension limit) + const std::vector X_dims{1, 1449, 1449, 3}; + const std::vector perm{0, 3, 1, 2}; + const std::vector Y_dims{1, 1449, 1449, 3}; + TestTransposeMLFloat16(perm, X_dims, Y_dims); +} #endif } // namespace test