From 9174cbe3d5e40eb457a574c82e1968a50c0995a1 Mon Sep 17 00:00:00 2001 From: Sherlock Date: Mon, 13 Sep 2021 23:00:53 -0700 Subject: [PATCH] Optimize CUDA Kernel for 3D and 4D Transpose (#8928) * Optimize Transpose120 and Transpose102 * Generalize Transpose0123 for more input shapes * Add Transpose3D test cases * update rocm kernel --- .../core/providers/cuda/tensor/transpose.cc | 80 +++++--- .../providers/cuda/tensor/transpose_impl.cu | 182 ++++++++++++------ .../providers/cuda/tensor/transpose_impl.h | 25 ++- .../core/providers/rocm/tensor/transpose.cc | 79 +++++--- .../providers/cpu/tensor/transpose_test.cc | 34 +++- .../python/orttraining_test_ortmodule_api.py | 135 +++++++++++++ 6 files changed, 399 insertions(+), 136 deletions(-) diff --git a/onnxruntime/core/providers/cuda/tensor/transpose.cc b/onnxruntime/core/providers/cuda/tensor/transpose.cc index e429a1eefb..b929889982 100644 --- a/onnxruntime/core/providers/cuda/tensor/transpose.cc +++ b/onnxruntime/core/providers/cuda/tensor/transpose.cc @@ -60,6 +60,7 @@ Status TransposeWithCublas(cudaStream_t stream, cublasHandle_t cublas_handle, co CudaT zero = ToCudaType::FromFloat(0.0f); const CudaT* input_data = reinterpret_cast(input.Data()); CudaT* output_data = reinterpret_cast(output.MutableData()); + 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()) { - return TransposeWithCublas(stream, cublas_handle, input, output, M, N); - } else if (element_type == utils::GetONNXTensorElementDataType()) { - return TransposeWithCublas(stream, cublas_handle, input, output, M, N); - } else { - return TransposeWithCublas(stream, cublas_handle, input, output, M, N); - } - } - } - const std::vector& input_dims = input_shape_override ? input_shape_override->GetDims() : input.Shape().GetDims(); const std::vector& output_dims = output.Shape().GetDims(); auto rank = static_cast(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()) { + return TransposeWithCublas(stream, cublas_handle, input, output, M, N); + } else if (element_type == utils::GetONNXTensorElementDataType()) { + return TransposeWithCublas(stream, cublas_handle, input, output, M, N); + } else { + return TransposeWithCublas(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 input_shape(new_input_dims); TArray 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 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(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(output.Shape().Size())); + tmp_output_strides, output.MutableDataRaw(), gsl::narrow(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 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(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(output.Shape().Size())); + tmp_output_strides, output.MutableDataRaw(), gsl::narrow(output.Shape().Size()), + grid_size, block_size); } // General cases diff --git a/onnxruntime/core/providers/cuda/tensor/transpose_impl.cu b/onnxruntime/core/providers/cuda/tensor/transpose_impl.cu index 006dce292f..df155f948d 100644 --- a/onnxruntime/core/providers/cuda/tensor/transpose_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/transpose_impl.cu @@ -27,25 +27,36 @@ __global__ void Transpose3DKernel(const TArray 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& input_dims, - const std::vector& permutations) { + const std::vector& 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(input_dims[2] / TILE_DIM); + int grid_size_y = static_cast(input_dims[1] / TILE_DIM); + int grid_size_z = static_cast(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(grid_size_x), + static_cast(grid_size_y), + static_cast(grid_size_z)); + return true; + } else { + return false; + } } return false; } Status Transpose3DImpl(cudaStream_t stream, size_t element_size, const TArray& input_shape, const TArray& input_strides, - const void* input_data, void* output_data, int64_t N) { - dim3 block_size(TILE_DIM, TILE_DIM); - dim3 grid_size(static_cast(input_shape[2] / TILE_DIM), static_cast(input_shape[1] / TILE_DIM), static_cast(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<<>>( @@ -83,24 +94,29 @@ template __global__ void Transpose4DKernelParallelizeMultipleElementsPerThreadInInnermostDim( const TArray input_strides, const void* input_data, const TArray 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(input_data); int4* v_output = reinterpret_cast(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& input_dims, - const std::vector& permutations) { + const std::vector& 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(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(block_size_x), static_cast(block_size_y)); + grid_size = dim3(static_cast(num_block_ext), + static_cast(input_dims[1]), + static_cast(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& input_shape, const TArray& input_strides, const void* input_data, const TArray& 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(element_size); // int4 is used in the kernel to access data. - dim3 block_size(static_cast(input_shape[3] / num_elements_per_thread), static_cast(input_shape[2])); - dim3 grid_size(static_cast(input_shape[1]), static_cast(input_shape[0])); switch (element_size) { case sizeof(int8_t): Transpose4DKernelParallelizeMultipleElementsPerThreadInInnermostDim <<>>( 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 <<>>( 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 <<>>( 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 <<>>( 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 input_strides, const int8_t* input_data, const TArray 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& input_dims, - const std::vector& permutations) { + const std::vector& 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(block_size_x), static_cast(block_size_y)); + grid_size = dim3(static_cast(num_block_ext), + static_cast(input_dims[1]), + static_cast(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& input_shape, const TArray& input_strides, const void* input_data, const TArray& 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(input_shape[3]), static_cast(input_shape[2])); - dim3 grid_size(static_cast(input_shape[1]), static_cast(input_shape[0])); - Transpose4DKernelParallelizeOneElementPerThread<<>>( input_strides, reinterpret_cast(input_data), output_strides, reinterpret_cast(output_data), - element_size, N); + element_size, input_shape[2], N); return Status::OK(); } diff --git a/onnxruntime/core/providers/cuda/tensor/transpose_impl.h b/onnxruntime/core/providers/cuda/tensor/transpose_impl.h index a9184d2a16..a2b5038096 100644 --- a/onnxruntime/core/providers/cuda/tensor/transpose_impl.h +++ b/onnxruntime/core/providers/cuda/tensor/transpose_impl.h @@ -8,27 +8,36 @@ namespace onnxruntime { namespace cuda { -bool CanDoTranspose3D(int32_t rank, const std::vector& input_dims, const std::vector& permutations); +bool CanDoTranspose3D(const cudaDeviceProp& prop, + int32_t rank, const std::vector& input_dims, const std::vector& permutations, + dim3& grid_size, dim3& block_size); Status Transpose3DImpl(cudaStream_t stream, size_t element_size, const TArray& input_shape, const TArray& 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& input_dims, - const std::vector& permutations); -Status Transpose4DParallelizeMultipleElementsPerThreadInInnermostDim(cudaStream_t stream, size_t element_size, const TArray& input_shape, + const std::vector& permutations, + dim3& grid_size, dim3& block_size); +Status Transpose4DParallelizeMultipleElementsPerThreadInInnermostDim(cudaStream_t stream, + size_t element_size, const TArray& input_shape, const TArray& input_strides, const void* input_data, - const TArray& output_strides, void* output_data, int N); + const TArray& 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& input_dims, - const std::vector& permutations); -Status Transpose4DParallelizeOneElementPerThread(cudaStream_t stream, size_t element_size, const TArray& input_shape, + const std::vector& permutations, + dim3& grid_size, dim3& block_size); +Status Transpose4DParallelizeOneElementPerThread(cudaStream_t stream, + size_t element_size, const TArray& input_shape, const TArray& input_strides, const void* input_data, - const TArray& output_strides, void* output_data, int N); + const TArray& 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& input_strides, const void* input_data, const TArray& fdm_output_strides, void* output_data, int N); diff --git a/onnxruntime/core/providers/rocm/tensor/transpose.cc b/onnxruntime/core/providers/rocm/tensor/transpose.cc index 61e1147abe..dc4de86db2 100644 --- a/onnxruntime/core/providers/rocm/tensor/transpose.cc +++ b/onnxruntime/core/providers/rocm/tensor/transpose.cc @@ -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() || - element_type == utils::GetONNXTensorElementDataType() || - element_type == utils::GetONNXTensorElementDataType()) { - 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()) { - return TransposeWithRocblas(stream, rocblas_handle, input, output, M, N); - } else if (element_type == utils::GetONNXTensorElementDataType()) { - return TransposeWithRocblas(stream, rocblas_handle, input, output, M, N); - } else { - return TransposeWithRocblas(stream, rocblas_handle, input, output, M, N); - } - } - } - const std::vector& input_dims = input_shape_override ? input_shape_override->GetDims() : input.Shape().GetDims(); const std::vector& output_dims = output.Shape().GetDims(); auto rank = static_cast(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()) { + return TransposeWithRocblas(stream, rocblas_handle, input, output, M, N); + } else if (element_type == utils::GetONNXTensorElementDataType()) { + return TransposeWithRocblas(stream, rocblas_handle, input, output, M, N); + } else { + return TransposeWithRocblas(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 input_shape(new_input_dims); TArray 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 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(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(output.Shape().Size())); + tmp_output_strides, output.MutableDataRaw(), gsl::narrow(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 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(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(output.Shape().Size())); + tmp_output_strides, output.MutableDataRaw(), gsl::narrow(output.Shape().Size()), + grid_size, block_size); } // General cases diff --git a/onnxruntime/test/providers/cpu/tensor/transpose_test.cc b/onnxruntime/test/providers/cpu/tensor/transpose_test.cc index 515fa120c6..e3e51e78f2 100644 --- a/onnxruntime/test/providers/cpu/tensor/transpose_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/transpose_test.cc @@ -188,7 +188,7 @@ TEST(TransposeOpTest, TwoDimStr) { } // Test 3 dimensional transpose, with permutation attribute specified -TEST(TransposeOpTest, ThreeDim) { +TEST(TransposeOpTest, Transpose021) { std::vector input_shape({4, 2, 3}); std::vector 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 input_shape({4, 2, 3}); + std::vector 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 perm = {1, 2, 0}; + std::vector 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 input_shape({4, 2, 3}); std::vector input_vals = { 1.0f, 2.0f, 3.0f, diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index f1ea028261..4256f58a04 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -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):