saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-18 18:52:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix cuda Transpose bug 16039 (#16042)

Browse source 
### 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
This commit is contained in:
Ryan Hill 2023-05-23 17:19:30 -07:00 committed by GitHub
parent 96ee72d7f8
commit b9d39e3405
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 42 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
onnxruntime/core/providers/cuda/fpgeneric.cu
View file

@ -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<<<dimGrid, dimBlock, 0, stream>>>(C, A, n, m);
} else {
return CUBLAS_STATUS_NOT_SUPPORTED;

1
onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h
View file

@ -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

2
onnxruntime/core/providers/cuda/tensor/transpose.cc
View file

@ -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<float>()) {
return TransposeWithCublas<float>(stream, cublas_handle, input, output, M, N);
} else if (element_type == utils::GetONNXTensorElementDataType<double>()) {

2
onnxruntime/core/providers/rocm/shared_inc/fpgeneric.h
View file

@ -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

26
onnxruntime/test/providers/cpu/tensor/transpose_test.cc
View file

@ -755,6 +755,32 @@ TEST(TransposeOpTest, Transpose3DImpl) {
}
}
static void TestTransposeMLFloat16(
const std::vector<int64_t>& perm,
const std::vector<int64_t>& x_dims,
const std::vector<int64_t>& y_dims,
double error_tolerance = 1e-4) {
CompareOpTester test{"Transpose"};
RandomValueGenerator random{};
const auto X_data = random.Uniform<float>(x_dims, -10.0, 10.0);
std::vector<MLFloat16> X_data_f16 = FloatsToMLFloat16s(X_data);
const auto Y_data = FillZeros<float>(y_dims);
std::vector<MLFloat16> 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<int64_t> X_dims{1, 1449, 1449, 3};
const std::vector<int64_t> perm{0, 3, 1, 2};
const std::vector<int64_t> Y_dims{1, 1449, 1449, 3};
TestTransposeMLFloat16(perm, X_dims, Y_dims);
}
#endif
} // namespace test