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

Accelerate CUDA transpose (#1102)

Browse source 
For NCHW <-> NHWC and 2D case
This commit is contained in:
KeDengMS 2019-05-23 21:10:50 -07:00 committed by GitHub
parent f4a9ccae99
commit cd9b9e6102
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 59 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

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

@ -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<T>()), \
Transpose<T>);
// special case acceleration using cublas matrix tranpose
std::tuple<int, int> TryTransposeWithCublas(const std::vector<size_t>& 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<int>(input_shape[1]);
N = gsl::narrow<int>(input_shape[2] * input_shape[3]);
} else {
M = gsl::narrow<int>(input_shape[1] * input_shape[2]);
N = gsl::narrow<int>(input_shape[3]);
}
}
} else if (perm.size() == 2 && perm[1] == 0 && perm[0] == 1) {
// 2D matrix transpose
M = gsl::narrow<int>(input_shape[0]);
N = gsl::narrow<int>(input_shape[1]);
}
return std::make_tuple(M, N);
}
template <typename T>
Status Transpose<T>::ComputeInternal(OpKernelContext* ctx) const {
const Tensor* X_ptr = ctx->Input<Tensor>(0);
@ -37,6 +64,34 @@ Status Transpose<T>::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<T>::MappedType CudaT;
CudaT one = ToCudaType<T>::FromFloat(1.0f);
CudaT zero = ToCudaType<T>::FromFloat(0.0f);
const CudaT* input_data = reinterpret_cast<const CudaT*>(X.template Data<T>());
CudaT* output_data = reinterpret_cast<CudaT*>(Y->template MutableData<T>());
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<int64_t> input_strides(this, device_id, rank);
CudaAsyncBuffer<size_t> perm(this, device_id, *p_perm);