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

[CUDA] Transpose3DImpl Supporting more Cases (#13611)

Browse source 
CUDA's Transpose3DImpl is to transpose [batch, m, n] to [batch, n, m].
Currently it requires both m and n can be divided by 32 or 16. If it's
not this case, the compute will fallback to general implementation,
which is slow. This PR is to remove the limitation.

Profiling in V100 using below size of tensors, got the cycles number
from Nsight Compute:
  | Old | New
-- | -- | --
[3072,64,512] | 760793 | 727140
[3072,16,2048] | 854303 | 851146
[3072,2048,12] | 986924 | 737884
[3072,1024,24] | 1212427 | 495117

It shows that even we added extra IF statements to the kernel
implementation, it has nearly no impact to the old version (case 1 and
2). And for case 3 and 4 which will fallback to general implementation
before, it's much faster.

Above data was collected using FP16 tensors, similar results was
observed for float tensors.

This PR is to enhance the perf of ORT training of Huggingface's XLNet
model which has[8,1024,1024,12].permute(0,3,1,2).
This commit is contained in:
Vincent Wang 2022-11-24 09:40:48 +08:00 committed by GitHub
parent 87d5703b14
commit 47e7630378
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 71 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -7,51 +7,60 @@
namespace onnxruntime {
namespace cuda {
constexpr unsigned int NUM_ELE_PER_THREAD = 4;
constexpr unsigned int kNumElementsPerThread = 4;
constexpr unsigned int kTileSize = 32;
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];
// TileSize for current implementation is always 32, but still use template parameter to make it flexible for future.
// For each batch, transpose matrix [m, n] to [n, m].
template <typename T, unsigned int TileSize>
__global__ void Transpose3DKernel(const int64_t m, const int64_t n, const int64_t batch_stride, const T* input_data,
T* output_data) {
__shared__ T tile[TileSize][TileSize + 1];
int x = blockIdx.x * TILE_DIM + threadIdx.x;
int y = blockIdx.y * TILE_DIM + threadIdx.y;
int x = blockIdx.x * TileSize + threadIdx.x;
int y = blockIdx.y * TileSize + threadIdx.y;
if (x < n) {
#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];
for (unsigned int i = 0; i < TileSize; i += (TileSize / kNumElementsPerThread)) {
int y_idx = y + i;
if (y_idx < m) {
tile[threadIdx.y + i][threadIdx.x] = input_data[blockIdx.z * batch_stride + y_idx * n + x];
}
}
}
__syncthreads();
x = blockIdx.y * TILE_DIM + threadIdx.x;
y = blockIdx.x * TILE_DIM + threadIdx.y;
x = blockIdx.y * TileSize + threadIdx.x;
y = blockIdx.x * TileSize + threadIdx.y;
if (x < m) {
#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];
for (unsigned int i = 0; i < TileSize; i += (TileSize / kNumElementsPerThread)) {
int y_idx = y + i;
if (y_idx < n) {
output_data[blockIdx.z * batch_stride + y_idx * m + x] = tile[threadIdx.x][threadIdx.y + i];
}
}
}
}
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.
// Permutation is done in the last two dimensions.
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);
// Normally maxGridSize.x is a large number but maxGridSize.y and maxGridSize.z are limited. Ideally we can check
// the input sizes to see if a dimension is too large so that we can use grid.x for it to avoid returning false.
// But this requires different versions of kernel implementation with different index compute logics.
// Below code is good enough for most of the cases for now, and if we see any case that input_dims[0] or
// input_dims[1] is too large in the future, we will handle it accordingly.
int grid_size_x = CeilDiv(static_cast<int>(input_dims[2]), kTileSize);
int grid_size_y = CeilDiv(static_cast<int>(input_dims[1]), kTileSize);
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 / NUM_ELE_PER_THREAD);
block_size = dim3(kTileSize, kTileSize / kNumElementsPerThread);
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;
@ -62,18 +71,12 @@ bool CanDoTranspose3D(const cudaDeviceProp& prop, size_t rank, const gsl::span<c
return false;
}
#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); \
} \
#define HANDLE_TRANSPOSE_3D_TILE_DIM(type) \
case sizeof(type): { \
Transpose3DKernel<type, kTileSize> \
<<<grid_size, block_size, 0, stream>>>(input_shape[1], input_shape[2], input_strides[0], \
reinterpret_cast<const ToCudaType<type>::MappedType*>(input_data), \
reinterpret_cast<ToCudaType<type>::MappedType*>(output_data)); \
} break
Status Transpose3DImpl(cudaStream_t stream, size_t element_size, const TArray<int64_t>& input_shape,

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

@ -723,18 +723,37 @@ TEST(TransposeOpTest, Transpose0213_V2) { // Will trigger Transpose4DParalleliz
TestTranspose(perm, X_dims, Y_dims);
}
TEST(TransposeOpTest, Transpose0231) { // Will trigger Transpose3DImpl() because of "flattening" of dims 2 and 3 into one dim
const std::vector<int64_t> X_dims{64, 128, 16, 64};
const std::vector<int64_t> perm{0, 2, 3, 1};
const std::vector<int64_t> Y_dims{64, 16, 64, 128};
TestTranspose(perm, X_dims, Y_dims);
}
TEST(TransposeOpTest, Transpose3DImpl) {
// Flattening dims 2 and 3 into one dim.
{
const std::vector<int64_t> X_dims{64, 128, 16, 64};
const std::vector<int64_t> perm{0, 2, 3, 1};
const std::vector<int64_t> Y_dims{64, 16, 64, 128};
TestTranspose(perm, X_dims, Y_dims);
}
TEST(TransposeOpTest, Transpose0312) { // Will trigger Transpose3DImpl() because of "flattening" of dims 1 and 2 into one dim
const std::vector<int64_t> X_dims{64, 16, 64, 128};
const std::vector<int64_t> perm{0, 3, 1, 2};
const std::vector<int64_t> Y_dims{64, 128, 16, 64};
TestTranspose(perm, X_dims, Y_dims);
// Flattening dims 1 and 2 into one dim.
{
const std::vector<int64_t> X_dims{64, 16, 64, 128};
const std::vector<int64_t> perm{0, 3, 1, 2};
const std::vector<int64_t> Y_dims{64, 128, 16, 64};
TestTranspose(perm, X_dims, Y_dims);
}
// dim-1 or dim-2 is not power of 2.
{
const std::vector<int64_t> X_dims{64, 12, 128};
const std::vector<int64_t> perm{0, 2, 1};
const std::vector<int64_t> Y_dims{64, 128, 12};
TestTranspose(perm, X_dims, Y_dims);
}
{
const std::vector<int64_t> X_dims{64, 99, 24};
const std::vector<int64_t> perm{0, 2, 1};
const std::vector<int64_t> Y_dims{64, 24, 99};
TestTranspose(perm, X_dims, Y_dims);
}
}
#endif