From cd9b9e6102757277e9761421ce9efa04753cd78e Mon Sep 17 00:00:00 2001 From: KeDengMS Date: Thu, 23 May 2019 21:10:50 -0700 Subject: [PATCH] Accelerate CUDA transpose (#1102) For NCHW <-> NHWC and 2D case --- onnxruntime/core/providers/cuda/fpgeneric.cu | 2 +- .../providers/cuda/shared_inc/fpgeneric.h | 6 +- .../core/providers/cuda/tensor/transpose.cc | 55 +++++++++++++++++++ 3 files changed, 59 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/cuda/fpgeneric.cu b/onnxruntime/core/providers/cuda/fpgeneric.cu index f00992ff50..5f8314d325 100644 --- a/onnxruntime/core/providers/cuda/fpgeneric.cu +++ b/onnxruntime/core/providers/cuda/fpgeneric.cu @@ -113,7 +113,7 @@ __global__ void CopyVectorHalf(const half* x, int incx, half* y, int incy, int n } // namespace -cublasStatus_t cublasTransposeHelper(cublasHandle_t, cublasOperation_t, cublasOperation_t, int m, int n, half*, half* A, int, half*, half*, int, half* C, int) { +cublasStatus_t cublasTransposeHelper(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 dimBlock(TRANS_TILE_DIM, BLOCK_ROWS, 1); diff --git a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h index afbea16f95..bb0d7ba56f 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h +++ b/onnxruntime/core/providers/cuda/shared_inc/fpgeneric.h @@ -59,13 +59,13 @@ inline cublasStatus_t cublasAxpyHelper(cublasHandle_t handle, int n, const half* } // transpose using geam -inline cublasStatus_t cublasTransposeHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, float* alpha, float* A, int lda, float* beta, float* B, int ldb, float* C, int ldc) { +inline cublasStatus_t cublasTransposeHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, const float* alpha, const float* A, int lda, const float* beta, const float* B, int ldb, float* C, int ldc) { return cublasSgeam(handle, transa, transb, m, n, alpha, A, lda, beta, B, ldb, C, ldc); } -inline cublasStatus_t cublasTransposeHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n, double* alpha, double* A, int lda, double* beta, double* B, int ldb, double* C, int ldc) { +inline cublasStatus_t cublasTransposeHelper(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t 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 cublasDgeam(handle, transa, transb, m, n, alpha, A, lda, beta, B, ldb, C, ldc); } -cublasStatus_t cublasTransposeHelper(cublasHandle_t, cublasOperation_t, cublasOperation_t, int m, int n, half*, half* A, int, half*, half*, int, half* C, int); +cublasStatus_t cublasTransposeHelper(cublasHandle_t, cublasOperation_t, cublasOperation_t, int m, int n, const half*, const half* A, int, const half*, const half*, int, half* C, int); // asum inline cublasStatus_t cublasAsumHelper(cublasHandle_t handle, int n, const float* x, int incx, float* result) { diff --git a/onnxruntime/core/providers/cuda/tensor/transpose.cc b/onnxruntime/core/providers/cuda/tensor/transpose.cc index b87d155d7e..04e69bbc8c 100644 --- a/onnxruntime/core/providers/cuda/tensor/transpose.cc +++ b/onnxruntime/core/providers/cuda/tensor/transpose.cc @@ -4,6 +4,7 @@ #include "transpose.h" #include "transpose_impl.h" #include "core/providers/cpu/tensor/utils.h" +#include "core/providers/cuda/shared_inc/fpgeneric.h" namespace onnxruntime { namespace cuda { @@ -19,6 +20,32 @@ namespace cuda { .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ Transpose); +// special case acceleration using cublas matrix tranpose +std::tuple TryTransposeWithCublas(const std::vector& perm, const TensorShape& input_shape) { + int M = 0; + int N = 0; + + if (perm.size() == 4 && input_shape[0] == 1 && perm[0] == 0) { + // NCHW < ->NHWC when N == 1 + if ((perm[1] == 2 && perm[2] == 3 && perm[3] == 1) || + (perm[1] == 3 && perm[2] == 1 && perm[3] == 2)) { + if (perm[1] == 2) { + M = gsl::narrow(input_shape[1]); + N = gsl::narrow(input_shape[2] * input_shape[3]); + } else { + M = gsl::narrow(input_shape[1] * input_shape[2]); + N = gsl::narrow(input_shape[3]); + } + } + } else if (perm.size() == 2 && perm[1] == 0 && perm[0] == 1) { + // 2D matrix transpose + M = gsl::narrow(input_shape[0]); + N = gsl::narrow(input_shape[1]); + } + + return std::make_tuple(M, N); +} + template Status Transpose::ComputeInternal(OpKernelContext* ctx) const { const Tensor* X_ptr = ctx->Input(0); @@ -37,6 +64,34 @@ Status Transpose::ComputeInternal(OpKernelContext* ctx) const { TensorShape output_shape{output_dims}; Tensor* Y = ctx->Output(0, output_shape); + + auto mn = TryTransposeWithCublas(*p_perm, input_shape); + int M = std::get<0>(mn); + int N = std::get<1>(mn); + if (M != 0 && N != 0) { + typedef typename ToCudaType::MappedType CudaT; + CudaT one = ToCudaType::FromFloat(1.0f); + CudaT zero = ToCudaType::FromFloat(0.0f); + const CudaT* input_data = reinterpret_cast(X.template Data()); + CudaT* output_data = reinterpret_cast(Y->template MutableData()); + CUBLAS_RETURN_IF_ERROR( + cublasTransposeHelper( + CublasHandle(), + CUBLAS_OP_T, + CUBLAS_OP_T, + M, + N, + &one, + input_data, + N, + &zero, + input_data, + N, + output_data, + M)); + return Status::OK(); + } + int device_id = 0; CudaAsyncBuffer input_strides(this, device_id, rank); CudaAsyncBuffer perm(this, device_id, *p_perm);