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

Optimize CUDA Kernel for 3D and 4D Transpose (#8928)

Browse source 
* Optimize Transpose120 and Transpose102

* Generalize Transpose0123 for more input shapes

* Add Transpose3D test cases

* update rocm kernel
This commit is contained in:
Sherlock 2021-09-13 23:00:53 -07:00 committed by GitHub
parent 5969d576e5
commit 9174cbe3d5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 399 additions and 136 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -60,6 +60,7 @@ Status TransposeWithCublas(cudaStream_t stream, cublasHandle_t cublas_handle, co
CudaT zero = ToCudaType<T>::FromFloat(0.0f);
const CudaT* input_data = reinterpret_cast<const CudaT*>(input.Data<T>());
CudaT* output_data = reinterpret_cast<CudaT*>(output.MutableData<T>());
CUBLAS_RETURN_IF_ERROR(
cublasTransposeHelper(stream,
cublas_handle,
@ -89,24 +90,6 @@ Status Transpose::DoTranspose(const cudaDeviceProp& prop,
if (output.Shape().Size() == 0)
return Status::OK();
auto element_type = input.GetElementType();
if (element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT ||
element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE ||
element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16) {
auto mn = TryTransposeWithCublas(permutations, input_shape_override ? *input_shape_override : input.Shape());
int M = std::get<0>(mn);
int N = std::get<1>(mn);
if (M != 0 && N != 0) {
if (element_type == utils::GetONNXTensorElementDataType<float>()) {
return TransposeWithCublas<float>(stream, cublas_handle, input, output, M, N);
} else if (element_type == utils::GetONNXTensorElementDataType<double>()) {
return TransposeWithCublas<double>(stream, cublas_handle, input, output, M, N);
} else {
return TransposeWithCublas<MLFloat16>(stream, cublas_handle, input, output, M, N);
}
}
}
const std::vector<int64_t>& input_dims = input_shape_override ? input_shape_override->GetDims() : input.Shape().GetDims();
const std::vector<int64_t>& output_dims = output.Shape().GetDims();
auto rank = static_cast<int32_t>(input_dims.size());
@ -155,37 +138,72 @@ Status Transpose::DoTranspose(const cudaDeviceProp& prop,
new_input_dims.resize(new_rank);
new_output_dims.resize(new_rank);
auto element_type = input.GetElementType();
size_t element_size = input.DataType()->Size();
if (element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT ||
element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE ||
element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16) {
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 (element_type == utils::GetONNXTensorElementDataType<float>()) {
return TransposeWithCublas<float>(stream, cublas_handle, input, output, M, N);
} else if (element_type == utils::GetONNXTensorElementDataType<double>()) {
return TransposeWithCublas<double>(stream, cublas_handle, input, output, M, N);
} else {
return TransposeWithCublas<MLFloat16>(stream, cublas_handle, input, output, M, N);
}
}
}
// 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)) {
TensorPitches new_input_strides(new_input_dims);
return Transpose3DImpl(stream, element_size, new_input_dims, new_input_strides,
input.DataRaw(), output.MutableDataRaw(), output.Shape().Size(), grid_size, block_size);
}
// 3D-Transpose can treated as a special case of 4D-Transpose with first dimension being 1.
if (new_rank == 3) {
new_permutations[0]++;
new_permutations[1]++;
new_permutations[2]++;
new_permutations.insert(new_permutations.begin(), 0);
new_input_dims.insert(new_input_dims.begin(), 1);
new_output_dims.insert(new_output_dims.begin(), 1);
new_rank = 4;
}
TensorPitches new_input_strides(new_input_dims);
TensorPitches new_output_strides(new_output_dims);
// Optimize the permutation of 3D/4D tensor
TArray<int64_t> input_shape(new_input_dims);
TArray<int64_t> tmp_input_strides(new_input_strides);
size_t element_size = input.DataType()->Size();
if (CanDoTranspose3D(new_rank, new_input_dims, new_permutations)) {
return Transpose3DImpl(stream, element_size, input_shape, tmp_input_strides,
input.DataRaw(), output.MutableDataRaw(), output.Shape().Size());
} else if (CanDoTranspose4DParallelizeMultipleElementsPerThreadInInnermostDim(
prop, element_size, new_rank, new_input_dims, new_permutations)) {
if (CanDoTranspose4DParallelizeMultipleElementsPerThreadInInnermostDim(
prop, element_size, new_rank, new_input_dims, new_permutations,
grid_size, block_size)) {
TArray<int64_t> tmp_output_strides(new_rank);
for (auto i = 0; i < new_rank; i++) {
tmp_output_strides[i] = new_output_strides[new_permutations[i]];
tmp_output_strides[static_cast<int32_t>(new_permutations[i])] = new_output_strides[i];
}
return Transpose4DParallelizeMultipleElementsPerThreadInInnermostDim(
stream, element_size, input_shape, tmp_input_strides, input.DataRaw(),
tmp_output_strides, output.MutableDataRaw(), gsl::narrow<int>(output.Shape().Size()));
tmp_output_strides, output.MutableDataRaw(), gsl::narrow<int>(output.Shape().Size()),
grid_size, block_size);
} else if (CanDoTranspose4DParallelizeOneElementPerThread(
prop, element_size, new_rank, new_input_dims, new_permutations)) {
prop, element_size, new_rank, new_input_dims, new_permutations, grid_size, block_size)) {
// Trying to see if we can still do (best effort) more optimized transposing
// for the 4-D case before falling back to the generic case
TArray<int64_t> tmp_output_strides(new_rank);
for (auto i = 0; i < new_rank; i++) {
tmp_output_strides[i] = new_output_strides[new_permutations[i]];
tmp_output_strides[static_cast<int32_t>(new_permutations[i])] = new_output_strides[i];
}
return Transpose4DParallelizeOneElementPerThread(
stream, element_size, input_shape, tmp_input_strides, input.DataRaw(),
tmp_output_strides, output.MutableDataRaw(), gsl::narrow<int>(output.Shape().Size()));
tmp_output_strides, output.MutableDataRaw(), gsl::narrow<int>(output.Shape().Size()),
grid_size, block_size);
}
// General cases

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

@ -27,25 +27,36 @@ __global__ void Transpose3DKernel(const TArray<int64_t> input_shape,
output_data[blockIdx.z * input_strides[0] + y * input_shape[1] + x] = tile[threadIdx.x * TILE_DIM + threadIdx.y];
}
bool CanDoTranspose3D(int32_t rank,
bool CanDoTranspose3D(const cudaDeviceProp& prop,
int32_t rank,
const std::vector<int64_t>& input_dims,
const std::vector<size_t>& permutations) {
const std::vector<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) {
return true;
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),
static_cast<unsigned int>(grid_size_z));
return true;
} else {
return false;
}
}
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) {
dim3 block_size(TILE_DIM, TILE_DIM);
dim3 grid_size(static_cast<unsigned int>(input_shape[2] / TILE_DIM), static_cast<unsigned int>(input_shape[1] / TILE_DIM), static_cast<unsigned int>(input_shape[0]));
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>>>(
@ -83,24 +94,29 @@ template <int element_size>
__global__ void Transpose4DKernelParallelizeMultipleElementsPerThreadInInnermostDim(
const TArray<int64_t> input_strides, const void* input_data,
const TArray<int64_t> output_strides, void* output_data,
CUDA_LONG N) {
// output coordinates will be: blockIdx.y, blockIdx.x, threadIdx.y, threadIdx.x
CUDA_LONG input_index = (blockIdx.y * input_strides[0] +
blockIdx.x * input_strides[1] +
threadIdx.y * input_strides[2]) /
(4 * sizeof(int) / element_size) +
threadIdx.x * input_strides[3];
int64_t input_shape_2, CUDA_LONG N) {
// coordinates will be: [d0, d1, d2, d3]
CUDA_LONG d0 = blockIdx.z;
CUDA_LONG d1 = blockIdx.y;
CUDA_LONG d2 = threadIdx.y + blockIdx.x * blockDim.y;
CUDA_LONG d3 = threadIdx.x;
CUDA_LONG output_index = (blockIdx.y * output_strides[0] +
blockIdx.x * output_strides[1] +
threadIdx.y * output_strides[2]) /
CUDA_LONG input_index = (d0 * input_strides[0] +
d1 * input_strides[1] +
d2 * input_strides[2]) /
(4 * sizeof(int) / element_size) +
d3 * input_strides[3];
CUDA_LONG output_index = (d0 * output_strides[0] +
d1 * output_strides[1] +
d2 * output_strides[2]) /
(4 * sizeof(int) / element_size) +
threadIdx.x * output_strides[3];
d3 * output_strides[3];
const int4* v_input = reinterpret_cast<const int4*>(input_data);
int4* v_output = reinterpret_cast<int4*>(output_data);
if (input_index < N && output_index < N) {
if (input_index < N && output_index < N && d2 < input_shape_2) {
v_output[output_index] = v_input[input_index];
}
}
@ -109,24 +125,38 @@ bool CanDoTranspose4DParallelizeMultipleElementsPerThreadInInnermostDim(const cu
size_t element_size,
int32_t rank,
const std::vector<int64_t>& input_dims,
const std::vector<size_t>& permutations) {
const std::vector<size_t>& permutations,
dim3& grid_size, dim3& block_size) {
if (rank == 4 &&
// the permutations is not on the last dimension.
permutations[3] == 3) {
// The block size will be set based on the outer-most two dimensions of 4D tensor.
// the number threads per block will be calculated as below.
unsigned int num_elements_per_thread = 4 * sizeof(int) / static_cast<unsigned int>(element_size); // int4 is used in the kernel to access data.
int64_t num_elements_in_last_two_dimensions = input_dims[2] * input_dims[3];
int64_t num_threads_per_block = num_elements_in_last_two_dimensions / num_elements_per_thread;
if (((num_elements_in_last_two_dimensions & (num_elements_per_thread - 1)) == 0) &&
num_threads_per_block <= prop.maxThreadsPerBlock &&
num_threads_per_block >= prop.warpSize &&
// num_threads_per_block must be a multiple of warp size (32)
((num_threads_per_block & (prop.warpSize - 1)) == 0) &&
// input_dims[3] must be a multiple of `num_elements_per_thread`
((input_dims[3] % num_elements_per_thread) == 0)) {
return true;
// dims[3]: block.x
// dims[2]: block.y + grid.x
// dims[1]: grid.y
// dims[0]: grid.z
if (input_dims[3] / num_elements_per_thread <= prop.maxThreadsPerBlock &&
(input_dims[3] % num_elements_per_thread) == 0 &&
input_dims[1] <= prop.maxGridSize[1] &&
input_dims[0] <= prop.maxGridSize[2]) {
// There are 2 constrains when luanching the kernels
// 1. block_size_x * block_size_y <= prop.maxThreadsPerBlock
// 2. block_size_y * num_block_ext >= input_dims[2]
int64_t block_size_x = input_dims[3] / num_elements_per_thread;
int64_t max_block_size_y = prop.maxThreadsPerBlock / block_size_x;
int64_t block_size_y = min(input_dims[2], max_block_size_y);
int64_t num_block_ext = CeilDiv(input_dims[2], block_size_y);
if (num_block_ext <= prop.maxGridSize[0]) {
block_size = dim3(static_cast<unsigned int>(block_size_x), static_cast<unsigned int>(block_size_y));
grid_size = dim3(static_cast<unsigned int>(num_block_ext),
static_cast<unsigned int>(input_dims[1]),
static_cast<unsigned int>(input_dims[0]));
return true;
} else {
return false;
}
}
}
return false;
@ -136,35 +166,41 @@ Status Transpose4DParallelizeMultipleElementsPerThreadInInnermostDim(
cudaStream_t stream, size_t element_size,
const TArray<int64_t>& input_shape, const TArray<int64_t>& input_strides,
const void* input_data, const TArray<int64_t>& output_strides,
void* output_data, int N) {
void* output_data, int N, const dim3& grid_size, const dim3& block_size) {
unsigned int num_elements_per_thread = 4 * sizeof(int) / static_cast<unsigned int>(element_size); // int4 is used in the kernel to access data.
dim3 block_size(static_cast<unsigned int>(input_shape[3] / num_elements_per_thread), static_cast<unsigned int>(input_shape[2]));
dim3 grid_size(static_cast<unsigned int>(input_shape[1]), static_cast<unsigned int>(input_shape[0]));
switch (element_size) {
case sizeof(int8_t):
Transpose4DKernelParallelizeMultipleElementsPerThreadInInnermostDim<sizeof(int8_t)>
<<<grid_size, block_size, 0, stream>>>(
input_strides, input_data,
output_strides, output_data, N / num_elements_per_thread);
output_strides, output_data,
input_shape[2],
N / num_elements_per_thread);
break;
case sizeof(int16_t):
Transpose4DKernelParallelizeMultipleElementsPerThreadInInnermostDim<sizeof(int16_t)>
<<<grid_size, block_size, 0, stream>>>(
input_strides, input_data,
output_strides, output_data, N / num_elements_per_thread);
output_strides, output_data,
input_shape[2],
N / num_elements_per_thread);
break;
case sizeof(int32_t):
Transpose4DKernelParallelizeMultipleElementsPerThreadInInnermostDim<sizeof(int32_t)>
<<<grid_size, block_size, 0, stream>>>(
input_strides, input_data,
output_strides, output_data, N / num_elements_per_thread);
output_strides, output_data,
input_shape[2],
N / num_elements_per_thread);
break;
case sizeof(int64_t):
Transpose4DKernelParallelizeMultipleElementsPerThreadInInnermostDim<sizeof(int64_t)>
<<<grid_size, block_size, 0, stream>>>(
input_strides, input_data,
output_strides, output_data, N / num_elements_per_thread);
output_strides, output_data,
input_shape[2],
N / num_elements_per_thread);
break;
default:
// User will not hit this as this kernel is for fixed element size tensors only
@ -178,19 +214,24 @@ Status Transpose4DParallelizeMultipleElementsPerThreadInInnermostDim(
__global__ void Transpose4DKernelParallelizeOneElementPerThread(
const TArray<int64_t> input_strides, const int8_t* input_data,
const TArray<int64_t> output_strides, int8_t* output_data,
size_t element_size,
CUDA_LONG N) {
CUDA_LONG input_index = blockIdx.y * input_strides[0] +
blockIdx.x * input_strides[1] +
threadIdx.y * input_strides[2] +
threadIdx.x * input_strides[3];
size_t element_size, int64_t input_shape_2, CUDA_LONG N) {
// coordinates will be: [d0, d1, d2, d3]
CUDA_LONG d0 = blockIdx.z;
CUDA_LONG d1 = blockIdx.y;
CUDA_LONG d2 = threadIdx.y + blockIdx.x * blockDim.y;
CUDA_LONG d3 = threadIdx.x;
CUDA_LONG output_index = blockIdx.y * output_strides[0] +
blockIdx.x * output_strides[1] +
threadIdx.y * output_strides[2] +
threadIdx.x * output_strides[3];
CUDA_LONG input_index = d0 * input_strides[0] +
d1 * input_strides[1] +
d2 * input_strides[2] +
d3 * input_strides[3];
if (input_index < N && output_index < N) {
CUDA_LONG output_index = d0 * output_strides[0] +
d1 * output_strides[1] +
d2 * output_strides[2] +
d3 * output_strides[3];
if (input_index < N && output_index < N && d2 < input_shape_2) {
const int8_t* input_data_to_be_copied = input_data + (input_index * element_size);
int8_t* output_data_to_be_copied = output_data + (output_index * element_size);
@ -205,17 +246,33 @@ bool CanDoTranspose4DParallelizeOneElementPerThread(const cudaDeviceProp& prop,
size_t element_size,
int32_t rank,
const std::vector<int64_t>& input_dims,
const std::vector<size_t>& permutations) {
const std::vector<size_t>& permutations,
dim3& grid_size, dim3& block_size) {
if (rank == 4) {
// The block size will be set based on the outer-most two dimensions of 4D tensor.
// the number threads per block will be calculated as below.
int64_t number_of_threads_per_block = input_dims[2] * input_dims[3];
// dims[3]: block.x
// dims[2]: block.y + grid.x
// dims[1]: grid.y
// dims[0]: grid.z
if (input_dims[3] <= prop.maxThreadsPerBlock &&
input_dims[1] <= prop.maxGridSize[1] &&
input_dims[0] <= prop.maxGridSize[2]) {
// There are 2 constrains when luanching the kernels
// 1. block_size_x * block_size_y <= prop.maxThreadsPerBlock
// 2. block_size_y * num_block_ext >= input_dims[2]
int64_t block_size_x = input_dims[3];
int64_t max_block_size_y = prop.maxThreadsPerBlock / block_size_x;
int64_t block_size_y = min(input_dims[2], max_block_size_y);
int64_t num_block_ext = CeilDiv(input_dims[2], block_size_y);
if (number_of_threads_per_block <= prop.maxThreadsPerBlock &&
number_of_threads_per_block >= prop.warpSize &&
// num_threads_per_block must be a multiple of warp size (32)
((number_of_threads_per_block & (prop.warpSize - 1)) == 0)) {
return true;
if (num_block_ext <= prop.maxGridSize[0]) {
block_size = dim3(static_cast<unsigned int>(block_size_x), static_cast<unsigned int>(block_size_y));
grid_size = dim3(static_cast<unsigned int>(num_block_ext),
static_cast<unsigned int>(input_dims[1]),
static_cast<unsigned int>(input_dims[0]));
return true;
} else {
return false;
}
}
}
return false;
@ -225,7 +282,7 @@ Status Transpose4DParallelizeOneElementPerThread(
cudaStream_t stream, size_t element_size,
const TArray<int64_t>& input_shape, const TArray<int64_t>& input_strides,
const void* input_data, const TArray<int64_t>& output_strides,
void* output_data, int N) {
void* output_data, int N, const dim3& grid_size, const dim3& block_size) {
if (element_size != sizeof(int8_t) &&
element_size != sizeof(int16_t) &&
element_size != sizeof(int32_t) &&
@ -235,13 +292,10 @@ Status Transpose4DParallelizeOneElementPerThread(
element_size);
}
dim3 block_size(static_cast<unsigned int>(input_shape[3]), static_cast<unsigned int>(input_shape[2]));
dim3 grid_size(static_cast<unsigned int>(input_shape[1]), static_cast<unsigned int>(input_shape[0]));
Transpose4DKernelParallelizeOneElementPerThread<<<grid_size, block_size, 0, stream>>>(
input_strides, reinterpret_cast<const int8_t*>(input_data),
output_strides, reinterpret_cast<int8_t*>(output_data),
element_size, N);
element_size, input_shape[2], N);
return Status::OK();
}

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

@ -8,27 +8,36 @@
namespace onnxruntime {
namespace cuda {
bool CanDoTranspose3D(int32_t rank, const std::vector<int64_t>& input_dims, const std::vector<size_t>& permutations);
bool CanDoTranspose3D(const cudaDeviceProp& prop,
int32_t rank, const std::vector<int64_t>& input_dims, const std::vector<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);
void* output_data, int64_t N,
const dim3& grid_size, const dim3& block_size);
bool CanDoTranspose4DParallelizeMultipleElementsPerThreadInInnermostDim(const cudaDeviceProp& prop,
size_t element_size,
int32_t rank,
const std::vector<int64_t>& input_dims,
const std::vector<size_t>& permutations);
Status Transpose4DParallelizeMultipleElementsPerThreadInInnermostDim(cudaStream_t stream, size_t element_size, const TArray<int64_t>& input_shape,
const std::vector<size_t>& permutations,
dim3& grid_size, dim3& block_size);
Status Transpose4DParallelizeMultipleElementsPerThreadInInnermostDim(cudaStream_t stream,
size_t element_size, const TArray<int64_t>& input_shape,
const TArray<int64_t>& input_strides, const void* input_data,
const TArray<int64_t>& output_strides, void* output_data, int N);
const TArray<int64_t>& output_strides, void* output_data, int N,
const dim3& grid_size, const dim3& block_size);
bool CanDoTranspose4DParallelizeOneElementPerThread(const cudaDeviceProp& prop,
size_t element_size,
int32_t rank,
const std::vector<int64_t>& input_dims,
const std::vector<size_t>& permutations);
Status Transpose4DParallelizeOneElementPerThread(cudaStream_t stream, size_t element_size, const TArray<int64_t>& input_shape,
const std::vector<size_t>& permutations,
dim3& grid_size, dim3& block_size);
Status Transpose4DParallelizeOneElementPerThread(cudaStream_t stream,
size_t element_size, const TArray<int64_t>& input_shape,
const TArray<int64_t>& input_strides, const void* input_data,
const TArray<int64_t>& output_strides, void* output_data, int N);
const TArray<int64_t>& output_strides, void* output_data, int N,
const dim3& grid_size, const dim3& block_size);
Status TransposeImpl(cudaStream_t stream, size_t element_size, int32_t shape_rank, const TArray<int64_t>& input_strides,
const void* input_data, const TArray<fast_divmod>& fdm_output_strides, void* output_data, int N);

79
onnxruntime/core/providers/rocm/tensor/transpose.cc
View file

@ -89,24 +89,6 @@ Status Transpose::DoTranspose(const hipDeviceProp_t& prop,
if (output.Shape().Size() == 0)
return Status::OK();
auto element_type = input.GetElementType();
if (element_type == utils::GetONNXTensorElementDataType<float>() ||
element_type == utils::GetONNXTensorElementDataType<double>() ||
element_type == utils::GetONNXTensorElementDataType<MLFloat16>()) {
auto mn = TryTransposeWithRocblas(permutations, input_shape_override ? *input_shape_override : input.Shape());
int M = std::get<0>(mn);
int N = std::get<1>(mn);
if (M != 0 && N != 0) {
if (element_type == utils::GetONNXTensorElementDataType<float>()) {
return TransposeWithRocblas<float>(stream, rocblas_handle, input, output, M, N);
} else if (element_type == utils::GetONNXTensorElementDataType<double>()) {
return TransposeWithRocblas<double>(stream, rocblas_handle, input, output, M, N);
} else {
return TransposeWithRocblas<MLFloat16>(stream, rocblas_handle, input, output, M, N);
}
}
}
const std::vector<int64_t>& input_dims = input_shape_override ? input_shape_override->GetDims() : input.Shape().GetDims();
const std::vector<int64_t>& output_dims = output.Shape().GetDims();
auto rank = static_cast<int32_t>(input_dims.size());
@ -155,37 +137,72 @@ Status Transpose::DoTranspose(const hipDeviceProp_t& prop,
new_input_dims.resize(new_rank);
new_output_dims.resize(new_rank);
auto element_type = input.GetElementType();
size_t element_size = input.DataType()->Size();
if (element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT ||
element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE ||
element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16) {
auto mn = TryTransposeWithRocblas(new_permutations, new_input_dims);
int M = std::get<0>(mn);
int N = std::get<1>(mn);
if (M != 0 && N != 0) {
if (element_type == utils::GetONNXTensorElementDataType<float>()) {
return TransposeWithRocblas<float>(stream, rocblas_handle, input, output, M, N);
} else if (element_type == utils::GetONNXTensorElementDataType<double>()) {
return TransposeWithRocblas<double>(stream, rocblas_handle, input, output, M, N);
} else {
return TransposeWithRocblas<MLFloat16>(stream, rocblas_handle, input, output, M, N);
}
}
}
// 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)) {
TensorPitches new_input_strides(new_input_dims);
return Transpose3DImpl(stream, element_size, new_input_dims, new_input_strides,
input.DataRaw(), output.MutableDataRaw(), output.Shape().Size(), grid_size, block_size);
}
// 3D-Transpose can treated as a special case of 4D-Transpose with first dimension being 1.
if (new_rank == 3) {
new_permutations[0]++;
new_permutations[1]++;
new_permutations[2]++;
new_permutations.insert(new_permutations.begin(), 0);
new_input_dims.insert(new_input_dims.begin(), 1);
new_output_dims.insert(new_output_dims.begin(), 1);
new_rank = 4;
}
TensorPitches new_input_strides(new_input_dims);
TensorPitches new_output_strides(new_output_dims);
// Optimize the permutation of 3D/4D tensor
TArray<int64_t> input_shape(new_input_dims);
TArray<int64_t> tmp_input_strides(new_input_strides);
size_t element_size = input.DataType()->Size();
if (CanDoTranspose3D(new_rank, new_input_dims, new_permutations)) {
return Transpose3DImpl(stream, element_size, input_shape, tmp_input_strides,
input.DataRaw(), output.MutableDataRaw(), output.Shape().Size());
} else if (CanDoTranspose4DParallelizeMultipleElementsPerThreadInInnermostDim(
prop, element_size, new_rank, new_input_dims, new_permutations)) {
if (CanDoTranspose4DParallelizeMultipleElementsPerThreadInInnermostDim(
prop, element_size, new_rank, new_input_dims, new_permutations,
grid_size, block_size)) {
TArray<int64_t> tmp_output_strides(new_rank);
for (auto i = 0; i < new_rank; i++) {
tmp_output_strides[i] = new_output_strides[new_permutations[i]];
tmp_output_strides[static_cast<int32_t>(new_permutations[i])] = new_output_strides[i];
}
return Transpose4DParallelizeMultipleElementsPerThreadInInnermostDim(
stream, element_size, input_shape, tmp_input_strides, input.DataRaw(),
tmp_output_strides, output.MutableDataRaw(), gsl::narrow<int>(output.Shape().Size()));
tmp_output_strides, output.MutableDataRaw(), gsl::narrow<int>(output.Shape().Size()),
grid_size, block_size);
} else if (CanDoTranspose4DParallelizeOneElementPerThread(
prop, element_size, new_rank, new_input_dims, new_permutations)) {
prop, element_size, new_rank, new_input_dims, new_permutations, grid_size, block_size)) {
// Trying to see if we can still do (best effort) more optimized transposing
// for the 4-D case before falling back to the generic case
TArray<int64_t> tmp_output_strides(new_rank);
for (auto i = 0; i < new_rank; i++) {
tmp_output_strides[i] = new_output_strides[new_permutations[i]];
tmp_output_strides[static_cast<int32_t>(new_permutations[i])] = new_output_strides[i];
}
return Transpose4DParallelizeOneElementPerThread(
stream, element_size, input_shape, tmp_input_strides, input.DataRaw(),
tmp_output_strides, output.MutableDataRaw(), gsl::narrow<int>(output.Shape().Size()));
tmp_output_strides, output.MutableDataRaw(), gsl::narrow<int>(output.Shape().Size()),
grid_size, block_size);
}
// General cases

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

@ -188,7 +188,7 @@ TEST(TransposeOpTest, TwoDimStr) {
}
// Test 3 dimensional transpose, with permutation attribute specified
TEST(TransposeOpTest, ThreeDim) {
TEST(TransposeOpTest, Transpose021) {
std::vector<int64_t> input_shape({4, 2, 3});
std::vector<float> input_vals = {
1.0f, 2.0f, 3.0f,
@ -225,8 +225,38 @@ TEST(TransposeOpTest, ThreeDim) {
TransposeTest(input_shape, input_vals, &perm, expected_shape, expected_vals, false); //TensorRT: illegal error
}
TEST(TransposeOpTest, Transpose120) {
std::vector<int64_t> input_shape({4, 2, 3});
std::vector<float> input_vals = {
1.0f, 2.0f, 3.0f,
4.0f, 5.0f, 6.0f,
1.1f, 2.1f, 3.1f,
4.1f, 5.1f, 6.1f,
1.2f, 2.2f, 3.2f,
4.2f, 5.2f, 6.2f,
1.3f, 2.3f, 3.3f,
4.3f, 5.3f, 6.3f};
std::vector<int64_t> perm = {1, 2, 0};
std::vector<int64_t> expected_shape({2, 3, 4});
auto expected_vals = {
1.0f, 1.1f, 1.2f, 1.3f,
2.0f, 2.1f, 2.2f, 2.3f,
3.0f, 3.1f, 3.2f, 3.3f,
4.0f, 4.1f, 4.2f, 4.3f,
5.0f, 5.1f, 5.2f, 5.3f,
6.0f, 6.1f, 6.2f, 6.3f
};
TransposeTest(input_shape, input_vals, &perm, expected_shape, expected_vals, false); //TensorRT: illegal error
}
// test when the suffix size is > 1 (last dimension is not moved)
TEST(TransposeOpTest, ThreeDimSuffix) {
TEST(TransposeOpTest, Transpose102) {
std::vector<int64_t> input_shape({4, 2, 3});
std::vector<float> input_vals = {
1.0f, 2.0f, 3.0f,

135
orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
View file

@ -667,6 +667,141 @@ def test_gradient_correctness_conv1d(use_fp16, input_requires_grad):
_test_helpers.assert_values_are_close(ort_prediction, pt_prediction, atol=1e-5)
_test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model, rtol=5e-2, atol=4e-2)
def _run_gradient_correctness_transpose(perm, shape):
class NeuralNetTranspose(torch.nn.Module):
def __init__(self, perm):
super(NeuralNetTranspose, self).__init__()
self.perm = perm
def forward(self, input):
out = torch.sin(input.permute(*self.perm))
return out
device = 'cuda'
pt_model = NeuralNetTranspose(perm).to(device)
ort_model = ORTModule(copy.deepcopy(pt_model))
def run_step(model, x):
prediction = model(x)
loss = prediction.sum()
loss.backward()
return prediction
x = torch.randn(*shape, device=device, requires_grad=True)
pt_prediction = run_step(pt_model, x)
ort_prediction = run_step(ort_model, x)
_test_helpers.assert_values_are_close(ort_prediction, pt_prediction)
_test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model)
@pytest.mark.parametrize("perm", [
[0,1,2], # no-op
[0,2,1], # special handle by Transpose021
[1,0,2], # handled as [0,2,1,3]
[1,2,0], # coalesced to [1,0]
[2,0,1], # coalesced to [1,0]
[2,1,0], # handled as [0,3,2,1]
])
@pytest.mark.parametrize("shape", [
[245,1024,32],
[255,2272,32],
[246,2080,32],
[254,128,256],
[260,245,256],
[284,254,256],
[245,260,256],
[1024,1024,256],
[254,284,256],
[4,5,2944],
[4,28,3136],
[4,312,768],
[3,224,224],
[17,5,4],
[8,2080,32],
[8,2272,32],
[2,2,2],
[1024,245,32],
[2080,246,32],
[1024,254,32],
[2272,255,32],
[4,5,736],
[4,111,160],
[8,246,32],
[8,255,32],
[4,1,2],
[1,2,2],
[2,1,2],
[2,2,1],
[2,1,4],
[4,2,1],
])
def test_gradient_correctness_transpose3d(perm, shape):
_run_gradient_correctness_transpose(perm, shape)
@pytest.mark.parametrize("perm", [
[0,1,2,3],
[0,1,3,2],
[0,2,1,3],
[0,2,3,1],
[0,3,1,2],
[0,3,2,1],
[1,0,2,3],
[1,0,3,2],
[1,2,0,3],
[1,2,3,0],
[1,3,0,2],
[1,3,2,0],
[2,0,1,3],
[2,0,3,1],
[2,1,0,3],
[2,1,3,0],
[2,3,0,1],
[2,3,1,0],
[3,0,1,2],
[3,0,2,1],
[3,1,0,2],
[3,1,2,0],
[3,2,0,1],
[3,2,1,0],
])
@pytest.mark.parametrize("shape", [
[1,245,1024,32],
[1,255,2272,32],
[1,246,2080,32],
[1,254,128,256],
[1,260,245,256],
[1,284,254,256],
[1,245,260,256],
[1,1024,1024,256],
[1,254,284,256],
[1,4,5,2944],
[1,4,28,3136],
[1,4,312,768],
[1,3,224,224],
[1,17,5,4],
[260,8,2080,32],
[284,8,2272,32],
[1,2,2,2],
[1,1024,245,32],
[1,2080,246,32],
[1,1024,254,32],
[1,2272,255,32],
[1,4,5,736],
[1,4,111,160],
[260,8,246,32],
[284,8,255,32],
[4,1,2,1],
[1,1,2,2],
[1,2,1,2],
[1,2,2,1],
[2,1,4,1],
[2,2,2,1],
[2,1,2,1],
[1,4,2,1],
])
def test_gradient_correctness_transpose4d(perm, shape):
_run_gradient_correctness_transpose(perm, shape)
@pytest.mark.parametrize("device", ['cuda', 'cpu'])
@pytest.mark.parametrize("padding_idx", [None, 1])
def test_gradient_correctness_embedding(device, padding_idx):