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

optimize Transpose3DKernel (#10891)

Browse source
This commit is contained in:
Vincent Wang 2022-03-17 18:09:04 +08:00 committed by GitHub
parent 463fac67a3
commit 6c0eff1ae4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 54 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -202,7 +202,7 @@ Status Transpose::DoTranspose(const cudaDeviceProp& prop,
// Transpose021 has a specialized Transpose3DImpl kernel
dim3 grid_size, block_size;
if (CanDoTranspose3D(prop, new_rank, new_input_dims, new_permutations, grid_size, block_size)) {
if (CanDoTranspose3D(prop, static_cast<size_t>(new_rank), new_input_dims, new_permutations, grid_size, block_size)) {
TensorPitches new_input_strides(new_input_dims);
return Transpose3DImpl(stream, element_size, ToConstSpan(new_input_dims), ToConstSpan(new_input_strides),
input.DataRaw(), output.MutableDataRaw(), output.Shape().Size(), grid_size, block_size);

102
onnxruntime/core/providers/cuda/tensor/transpose_impl.cu
View file

@ -7,44 +7,52 @@
namespace onnxruntime {
namespace cuda {
constexpr unsigned int TILE_DIM = 16;
constexpr unsigned int NUM_ELE_PER_THREAD = 4;
template <typename T>
__global__ void Transpose3DKernel(const TArray<int64_t> input_shape,
const TArray<int64_t> input_strides,
template <typename T, unsigned int TILE_DIM>
__global__ void Transpose3DKernel(const TArray<int64_t> input_shape, const TArray<int64_t> input_strides,
const T* input_data, T* output_data) {
__shared__ T tile[TILE_DIM * (TILE_DIM + 1)];
__shared__ T tile[TILE_DIM][TILE_DIM + 1];
int x = blockIdx.x * TILE_DIM + threadIdx.x;
int y = blockIdx.y * TILE_DIM + threadIdx.y;
tile[threadIdx.y * TILE_DIM + threadIdx.x] = input_data[blockIdx.z * input_strides[0] + y * input_shape[2] + x];
#pragma unroll
for (unsigned int i = 0; i < TILE_DIM; i += (TILE_DIM / NUM_ELE_PER_THREAD)) {
tile[threadIdx.y + i][threadIdx.x] = input_data[blockIdx.z * input_strides[0] + (y + i) * input_shape[2] + x];
}
__syncthreads();
x = blockIdx.y * TILE_DIM + threadIdx.x;
y = blockIdx.x * TILE_DIM + threadIdx.y;
output_data[blockIdx.z * input_strides[0] + y * input_shape[1] + x] = tile[threadIdx.x * TILE_DIM + threadIdx.y];
#pragma unroll
for (unsigned int i = 0; i < TILE_DIM; i += (TILE_DIM / NUM_ELE_PER_THREAD)) {
output_data[blockIdx.z * input_strides[0] + (y + i) * input_shape[1] + x] = tile[threadIdx.x][threadIdx.y + i];
}
}
bool CanDoTranspose3D(const cudaDeviceProp& prop,
int32_t rank,
const gsl::span<const int64_t>& input_dims,
const gsl::span<const size_t>& permutations,
dim3& grid_size, dim3& block_size) {
if (rank == 3 &&
// permutation is done in the last two dimensions.
permutations[rank - 2] == (rank - 1) && permutations[rank - 1] == (rank - 2) &&
// the last two dimensions are aligned with TILE_DIM.
input_dims[rank - 2] % TILE_DIM == 0 && input_dims[rank - 1] % TILE_DIM == 0) {
int grid_size_x = static_cast<int>(input_dims[2] / TILE_DIM);
int grid_size_y = static_cast<int>(input_dims[1] / TILE_DIM);
bool CanDoTranspose3D(const cudaDeviceProp& prop, size_t rank, const gsl::span<const int64_t>& input_dims,
const gsl::span<const size_t>& permutations, dim3& grid_size, dim3& block_size) {
// Permutation is done in the last two dimensions and the last two dimensions are aligned with TILE_DIM.
if (rank == 3 && permutations[rank - 2] == (rank - 1) && permutations[rank - 1] == (rank - 2)) {
unsigned int tile_dim = 0;
if (input_dims[rank - 2] % 32 == 0 && input_dims[rank - 1] % 32 == 0) {
tile_dim = 32;
} else if (input_dims[rank - 2] % 16 == 0 && input_dims[rank - 1] % 16 == 0) {
tile_dim = 16;
} else {
return false;
}
int grid_size_x = static_cast<int>(input_dims[2] / tile_dim);
int grid_size_y = static_cast<int>(input_dims[1] / tile_dim);
int grid_size_z = static_cast<int>(input_dims[0]);
if (grid_size_x <= prop.maxGridSize[0] && grid_size_y <= prop.maxGridSize[1] && grid_size_z <= prop.maxGridSize[2]) {
block_size = dim3(TILE_DIM, TILE_DIM);
grid_size = dim3(static_cast<unsigned int>(grid_size_x),
static_cast<unsigned int>(grid_size_y),
if (grid_size_x <= prop.maxGridSize[0] && grid_size_y <= prop.maxGridSize[1] &&
grid_size_z <= prop.maxGridSize[2]) {
block_size = dim3(tile_dim, tile_dim / NUM_ELE_PER_THREAD);
grid_size = dim3(static_cast<unsigned int>(grid_size_x), static_cast<unsigned int>(grid_size_y),
static_cast<unsigned int>(grid_size_z));
return true;
} else {
@ -54,34 +62,28 @@ bool CanDoTranspose3D(const cudaDeviceProp& prop,
return false;
}
Status Transpose3DImpl(cudaStream_t stream, size_t element_size,
const TArray<int64_t>& input_shape, const TArray<int64_t>& input_strides,
const void* input_data, void* output_data, int64_t N, const dim3& grid_size, const dim3& block_size) {
#define CALL_TRANSPOSE_3D(type, tile_dim) \
Transpose3DKernel<type, tile_dim><<<grid_size, block_size, 0, stream>>>( \
input_shape, input_strides, reinterpret_cast<const ToCudaType<type>::MappedType*>(input_data), \
reinterpret_cast<ToCudaType<type>::MappedType*>(output_data))
#define HANDLE_TRANSPOSE_3D_TILE_DIM(type) \
case sizeof(type): { \
if (block_size.x == 32) { \
CALL_TRANSPOSE_3D(type, 32); \
} else { \
CALL_TRANSPOSE_3D(type, 16); \
} \
} break
Status Transpose3DImpl(cudaStream_t stream, size_t element_size, const TArray<int64_t>& input_shape,
const TArray<int64_t>& input_strides, const void* input_data, void* output_data, int64_t N,
const dim3& grid_size, const dim3& block_size) {
switch (element_size) {
case sizeof(int8_t):
Transpose3DKernel<int8_t><<<grid_size, block_size, 0, stream>>>(
input_shape, input_strides,
reinterpret_cast<const ToCudaType<int8_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<int8_t>::MappedType*>(output_data));
break;
case sizeof(int16_t):
Transpose3DKernel<int16_t><<<grid_size, block_size, 0, stream>>>(
input_shape, input_strides,
reinterpret_cast<const ToCudaType<int16_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<int16_t>::MappedType*>(output_data));
break;
case sizeof(int32_t):
Transpose3DKernel<int32_t><<<grid_size, block_size, 0, stream>>>(
input_shape, input_strides,
reinterpret_cast<const ToCudaType<int32_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<int32_t>::MappedType*>(output_data));
break;
case sizeof(int64_t):
Transpose3DKernel<int64_t><<<grid_size, block_size, 0, stream>>>(
input_shape, input_strides,
reinterpret_cast<const ToCudaType<int64_t>::MappedType*>(input_data),
reinterpret_cast<ToCudaType<int64_t>::MappedType*>(output_data));
break;
HANDLE_TRANSPOSE_3D_TILE_DIM(int8_t);
HANDLE_TRANSPOSE_3D_TILE_DIM(int16_t);
HANDLE_TRANSPOSE_3D_TILE_DIM(int32_t);
HANDLE_TRANSPOSE_3D_TILE_DIM(int64_t);
default:
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for transpose on CUDA. Element size was ",
element_size);

2
onnxruntime/core/providers/cuda/tensor/transpose_impl.h
View file

@ -9,7 +9,7 @@ namespace onnxruntime {
namespace cuda {
bool CanDoTranspose3D(const cudaDeviceProp& prop,
int32_t rank, const gsl::span<const int64_t>& input_dims, const gsl::span<const size_t>& permutations,
size_t rank, const gsl::span<const int64_t>& input_dims, const gsl::span<const size_t>& permutations,
dim3& grid_size, dim3& block_size);
Status Transpose3DImpl(cudaStream_t stream, size_t element_size, const TArray<int64_t>& input_shape, const TArray<int64_t>& input_strides, const void* input_data,
void* output_data, int64_t N,