From 33e06be4ac13c6a58d3365f4cc6e33ef5ed410b4 Mon Sep 17 00:00:00 2001 From: Weixing Zhang Date: Thu, 2 Jul 2020 22:05:32 -0700 Subject: [PATCH] optimize transpose CUDA kernel (#4233) * optimize transpose * optimize for the case when the tensor is 3D and the permutation is done in last two dimension. BERT-L throughput is improved ~1.4% from transpose optimization * fix UT MegatronSelfAttentionPartitionCorrectnessTest * polish code. * add test and change tile size to 16x16 for better perf. * fix UT * fix test of mask_rcnn * address code review comments. Co-authored-by: Weixing Zhang --- .../math/einsum_utils/einsum_auxiliary_ops.cc | 3 +- .../core/providers/cuda/tensor/transpose.cc | 91 ++++++++-- .../core/providers/cuda/tensor/transpose.h | 3 +- .../providers/cuda/tensor/transpose_impl.cu | 168 +++++++++++++++++- .../providers/cuda/tensor/transpose_impl.h | 11 +- .../providers/cpu/tensor/transpose_test.cc | 43 +++++ 6 files changed, 296 insertions(+), 23 deletions(-) diff --git a/onnxruntime/core/providers/cuda/math/einsum_utils/einsum_auxiliary_ops.cc b/onnxruntime/core/providers/cuda/math/einsum_utils/einsum_auxiliary_ops.cc index e00d74a3e9..64981689eb 100644 --- a/onnxruntime/core/providers/cuda/math/einsum_utils/einsum_auxiliary_ops.cc +++ b/onnxruntime/core/providers/cuda/math/einsum_utils/einsum_auxiliary_ops.cc @@ -27,7 +27,8 @@ Status DataCopy(const Tensor& input, Tensor& output) { // CUDA EP specific Transpose helper Status Transpose(const std::vector& permutation, const Tensor& input, Tensor& output, const TensorShape* input_shape_override, void* einsum_cuda_assets) { - return cuda::Transpose::DoTranspose(static_cast(einsum_cuda_assets)->cublas_handle_, + return cuda::Transpose::DoTranspose(static_cast(einsum_cuda_assets)->cuda_ep_->GetDeviceProp(), + static_cast(einsum_cuda_assets)->cublas_handle_, permutation, input, output, input_shape_override); } diff --git a/onnxruntime/core/providers/cuda/tensor/transpose.cc b/onnxruntime/core/providers/cuda/tensor/transpose.cc index c9678f9b36..f711fc0131 100644 --- a/onnxruntime/core/providers/cuda/tensor/transpose.cc +++ b/onnxruntime/core/providers/cuda/tensor/transpose.cc @@ -66,10 +66,11 @@ Status TransposeWithCublas(cublasHandle_t cublas_handle, const Tensor& input, Te Status Transpose::DoTranspose(const Transpose& transpose_kernel, const std::vector& permutations, const Tensor& input, Tensor& output) { - return Transpose::DoTranspose(transpose_kernel.CublasHandle(), permutations, input, output); + return Transpose::DoTranspose(transpose_kernel.GetDeviceProp(), transpose_kernel.CublasHandle(), permutations, input, output); } -Status Transpose::DoTranspose(const cublasHandle_t cublas_handle, +Status Transpose::DoTranspose(const cudaDeviceProp& prop, + const cublasHandle_t cublas_handle, const std::vector& permutations, const Tensor& input, Tensor& output, const TensorShape* input_shape_override) { // special case when there is a dim value of 0 in the shape. @@ -96,22 +97,84 @@ Status Transpose::DoTranspose(const cublasHandle_t cublas_handle, 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()); - TensorPitches original_input_strides(input_dims); - TensorPitches original_output_strides(output_dims); - TArray input_strides(rank); - for (auto i = 0; i < rank; i++) { - input_strides[i] = original_input_strides[permutations[i]]; - } - TArray output_strides(rank); - for (auto i = 0; i < rank; i++) { - output_strides[i] = fast_divmod(gsl::narrow_cast(original_output_strides[i])); + // flatten the adjacent dimensions which are contiguous + // for example: permutations[0, 2, 3, 1] -> [0, 2, 1], permutations[0, 3, 1, 2] -> [0, 2, 1] + auto new_rank = rank; + std::vector new_permutations(permutations); + std::vector new_input_dims(input_dims); + std::vector new_output_dims(output_dims); + + for (auto i = rank - 1; i > 0; i--) { + auto curr = new_permutations[i]; + auto prev = new_permutations[i - 1]; + if (prev + 1 == curr) { + // all dims bigger than curr need to be reduced by 1 due to the merging. + for (auto j = 0; j < new_rank; j++) { + if (new_permutations[j] > curr) { + new_permutations[j] -= 1; + } + } + for (auto j = i+1; j < new_rank; j++) { + new_permutations[j-1] = new_permutations[j]; + } + + // update input dims + new_input_dims[prev] *= new_input_dims[curr]; + new_input_dims[curr] = 1; + for (auto j = static_cast(curr+1); j < new_rank; j++) { + new_input_dims[j-1] = new_input_dims[j]; + } + new_input_dims[new_rank-1] = 1; + + // update output dims + new_output_dims[i-1] *= new_output_dims[i]; + new_output_dims[i] = 1; + for (auto j = i+1; j < new_rank; j++) { + new_output_dims[j-1] = new_output_dims[j]; + } + new_output_dims[new_rank-1] = 1; + + new_rank--; + } } + new_permutations.resize(new_rank); + new_input_dims.resize(new_rank); + new_output_dims.resize(new_rank); + + 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(); - auto status = TransposeImpl(element_size, rank, input_strides, input.DataRaw(), + if (CanDoTranspose3D(new_rank, new_input_dims, new_permutations)) { + return Transpose3DImpl(element_size, input_shape, tmp_input_strides, + input.DataRaw(), output.MutableDataRaw(), output.Shape().Size()); + } else if (CanDoTranspose4D(prop, element_size, new_rank, new_input_dims, new_permutations)) { + TArray tmp_output_strides(new_rank); + for (auto i = 0; i < new_rank; i++) { + tmp_output_strides[i] = new_output_strides[new_permutations[i]]; + } + return Transpose4DImpl(element_size, input_shape, tmp_input_strides, input.DataRaw(), + tmp_output_strides, output.MutableDataRaw(), output.Shape().Size()); + } + + // General cases + TArray input_strides(new_rank); + for (auto i = 0; i < new_rank; i++) { + input_strides[i] = new_input_strides[new_permutations[i]]; + } + + TArray output_strides(new_rank); + for (auto i = 0; i < new_rank; i++) { + output_strides[i] = fast_divmod(gsl::narrow_cast(new_output_strides[i])); + } + + auto status = TransposeImpl(element_size, new_rank, input_strides, input.DataRaw(), output_strides, output.MutableDataRaw(), output.Shape().Size()); return status; @@ -135,7 +198,7 @@ Status Transpose::ComputeInternal(OpKernelContext* ctx) const { TensorShape output_shape{output_dims}; Tensor* Y = ctx->Output(0, output_shape); - return DoTranspose(this->CublasHandle(), *p_perm, X, *Y); + return DoTranspose(this->GetDeviceProp(), this->CublasHandle(), *p_perm, X, *Y); } } // namespace cuda diff --git a/onnxruntime/core/providers/cuda/tensor/transpose.h b/onnxruntime/core/providers/cuda/tensor/transpose.h index 475420da70..518801a57c 100644 --- a/onnxruntime/core/providers/cuda/tensor/transpose.h +++ b/onnxruntime/core/providers/cuda/tensor/transpose.h @@ -23,7 +23,8 @@ class Transpose final : public CudaKernel, public TransposeBase { const std::vector& permutations, const Tensor& input, Tensor& output); // `input_shape_override` (if provided) overrides the shape of `input` for compute purposes - static Status DoTranspose(const cublasHandle_t cublas_handle, + static Status DoTranspose(const cudaDeviceProp& prop, + const cublasHandle_t cublas_handle, const std::vector& permutations, const Tensor& input, Tensor& output, const TensorShape* input_shape_override = nullptr); }; diff --git a/onnxruntime/core/providers/cuda/tensor/transpose_impl.cu b/onnxruntime/core/providers/cuda/tensor/transpose_impl.cu index 8e8067cd06..1c776f252f 100644 --- a/onnxruntime/core/providers/cuda/tensor/transpose_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/transpose_impl.cu @@ -7,9 +7,165 @@ namespace onnxruntime { namespace cuda { +constexpr int TILE_DIM = 16; + template -__global__ void _TransposeKernel(int32_t shape_rank, const TArray input_strides, - const T* input_data, const TArray output_strides, T* output_data, CUDA_LONG N) { +__global__ void Transpose3DKernel(const TArray input_shape, + const TArray input_strides, + const T* input_data, T* output_data) { + __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]; + __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]; +} + +bool CanDoTranspose3D(int32_t rank, + const std::vector& input_dims, + const std::vector& permutations) { + 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; + } + return false; +} + +Status Transpose3DImpl(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(input_shape[2]/TILE_DIM, input_shape[1]/TILE_DIM, input_shape[0]); + + switch (element_size) { + case sizeof(int8_t): + Transpose3DKernel<<>>( + input_shape, input_strides, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data)); + break; + case sizeof(int16_t): + Transpose3DKernel<<>>( + input_shape, input_strides, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data)); + break; + case sizeof(int32_t): + Transpose3DKernel<<>>( + input_shape, input_strides, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data)); + break; + case sizeof(int64_t): + Transpose3DKernel<<>>( + input_shape, input_strides, + reinterpret_cast::MappedType*>(input_data), + reinterpret_cast::MappedType*>(output_data)); + break; + default: + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for transpose on CUDA. Element size was ", + element_size); + } + + return Status::OK(); +} + +template +__global__ void Transpose4DKernel(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]; + + CUDA_LONG output_index = (blockIdx.y * output_strides[0] + + blockIdx.x * output_strides[1] + + threadIdx.y * output_strides[2]) / (4 * sizeof(int) / element_size) + + threadIdx.x * 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) { + v_output[output_index] = v_input[input_index]; + } +} + +bool CanDoTranspose4D(const cudaDeviceProp& prop, + size_t element_size, + int32_t rank, + const std::vector& input_dims, + const std::vector& permutations) { + if (rank == 4 && + // the permutations is not on the last dimension. + permutations[rank-1] == (rank - 1)) { + + // The block size will be set based on the last two dimensions of 4D tensor. + // the number threads per block will be calculated as below. + int num_elements_per_thread = 4 * sizeof(int) / element_size; // int4 is used in the kernel to access data. + int64_t num_elements_in_last_two_dimensions = input_dims[rank-2] * input_dims[rank-1]; + 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 aligned with warp size: 32 + ((num_threads_per_block & (prop.warpSize - 1)) == 0)) { + return true; + } + } + return false; +} + +Status Transpose4DImpl(size_t element_size, const TArray& input_shape, const TArray& input_strides, const void* input_data, + const TArray& output_strides, void* output_data, int64_t N) { + int num_elements_per_thread = 4 * sizeof(int) / element_size; // int4 is used in the kernel to access data. + dim3 block_size(input_shape[3]/num_elements_per_thread, input_shape[2]); + dim3 grid_size(input_shape[1], input_shape[0]); + + switch (element_size) { + case sizeof(int8_t): + Transpose4DKernel<<>>( + input_strides, input_data, + output_strides, output_data, N/num_elements_per_thread); + break; + case sizeof(int16_t): + Transpose4DKernel<<>>( + input_strides, input_data, + output_strides, output_data, N/num_elements_per_thread); + break; + case sizeof(int32_t): + Transpose4DKernel<<>>( + input_strides, input_data, + output_strides, output_data, N/num_elements_per_thread); + break; + case sizeof(int64_t): + Transpose4DKernel<<>>( + input_strides, input_data, + output_strides, output_data, N/num_elements_per_thread); + break; + default: + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for transpose on CUDA. Element size was ", + element_size); + } + + return Status::OK(); +} + +template +__global__ void TransposeKernel(int32_t shape_rank, const TArray input_strides, + const T* input_data, const TArray output_strides, T* output_data, CUDA_LONG N) { CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); CUDA_LONG input_index = 0; CUDA_LONG output_index = id; @@ -32,7 +188,7 @@ Status TransposeImpl(size_t element_size, int32_t shape_rank, const TArray(N) / GridDim::maxThreadsPerBlock)); switch (element_size) { case sizeof(int8_t): - _TransposeKernel<<>>( + TransposeKernel<<>>( shape_rank, input_strides, reinterpret_cast::MappedType*>(input_data), fdm_output_strides, @@ -40,7 +196,7 @@ Status TransposeImpl(size_t element_size, int32_t shape_rank, const TArray<<>>( + TransposeKernel<<>>( shape_rank, input_strides, reinterpret_cast::MappedType*>(input_data), fdm_output_strides, @@ -48,7 +204,7 @@ Status TransposeImpl(size_t element_size, int32_t shape_rank, const TArray<<>>( + TransposeKernel<<>>( shape_rank, input_strides, reinterpret_cast::MappedType*>(input_data), fdm_output_strides, @@ -56,7 +212,7 @@ Status TransposeImpl(size_t element_size, int32_t shape_rank, const TArray<<>>( + TransposeKernel<<>>( shape_rank, input_strides, reinterpret_cast::MappedType*>(input_data), fdm_output_strides, diff --git a/onnxruntime/core/providers/cuda/tensor/transpose_impl.h b/onnxruntime/core/providers/cuda/tensor/transpose_impl.h index b510424045..0067995320 100644 --- a/onnxruntime/core/providers/cuda/tensor/transpose_impl.h +++ b/onnxruntime/core/providers/cuda/tensor/transpose_impl.h @@ -8,8 +8,17 @@ namespace onnxruntime { namespace cuda { +bool CanDoTranspose3D(int32_t rank, const std::vector& input_dims, const std::vector& permutations); +Status Transpose3DImpl(size_t element_size, const TArray& input_shape, const TArray& input_strides, const void* input_data, + void* output_data, int64_t N); +bool CanDoTranspose4D(const cudaDeviceProp& prop, + size_t element_size, + int32_t rank, + const std::vector& input_dims, + const std::vector& permutations); +Status Transpose4DImpl(size_t element_size, const TArray& input_shape, const TArray& input_strides, const void* input_data, + const TArray& output_strides, void* output_data, int64_t N); Status TransposeImpl(size_t element_size, int32_t shape_rank, const TArray& input_strides, const void* input_data, const TArray& fdm_output_strides, void* output_data, int64_t N); - } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/tensor/transpose_test.cc b/onnxruntime/test/providers/cpu/tensor/transpose_test.cc index d1420d8fe0..8ff1c8f2b5 100644 --- a/onnxruntime/test/providers/cpu/tensor/transpose_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/transpose_test.cc @@ -3,6 +3,7 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "test/providers/compare_provider_test_utils.h" namespace onnxruntime { namespace test { @@ -409,5 +410,47 @@ TEST(TransposeOpTest, SingleAxisMovingInwardsBlockCopy) { TransposeTest(input_shape, input_vals, &perm, expected_shape, expected_vals, false); } + +#ifdef USE_CUDA +static void TestTranspose( + const std::vector& perm, + const std::vector& x_dims, + const std::vector& y_dims, + double error_tolerance = 1e-4) { + CompareOpTester test{"Transpose"}; + + RandomValueGenerator random{}; + const auto X_data = random.Uniform(x_dims, -10.0, 10.0); + const auto Y_data = FillZeros(y_dims); + + test.AddAttribute("perm", perm); + test.AddInput("X", x_dims, X_data); + test.AddOutput("Y", y_dims, Y_data); + + test.CompareWithCPU(kCudaExecutionProvider, error_tolerance); +} + +TEST(TransposeOpTest, Transpose0213) { + const std::vector X_dims{64, 128, 16, 64}; + const std::vector perm{0, 2, 1, 3}; + const std::vector Y_dims{64, 16, 128, 64}; + TestTranspose(perm, X_dims, Y_dims); +} + +TEST(TransposeOpTest, Transpose0231) { + const std::vector X_dims{64, 128, 16, 64}; + const std::vector perm{0, 2, 3, 1}; + const std::vector Y_dims{64, 16, 64, 128}; + TestTranspose(perm, X_dims, Y_dims); +} + +TEST(TransposeOpTest, Transpose0312) { + const std::vector X_dims{64, 16, 64, 128}; + const std::vector perm{0, 3, 1, 2}; + const std::vector Y_dims{64, 128, 16, 64}; + TestTranspose(perm, X_dims, Y_dims); +} +#endif + } // namespace test } // namespace onnxruntime