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 <wezhan@microsoft.com>

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<size_t>& permutation, const Tensor& input,
                  Tensor& output, const TensorShape* input_shape_override, void* einsum_cuda_assets) {
-  return cuda::Transpose::DoTranspose(static_cast<EinsumCudaAssets*>(einsum_cuda_assets)->cublas_handle_,
+  return cuda::Transpose::DoTranspose(static_cast<EinsumCudaAssets*>(einsum_cuda_assets)->cuda_ep_->GetDeviceProp(),
+                                      static_cast<EinsumCudaAssets*>(einsum_cuda_assets)->cublas_handle_,
                                       permutation, input, output, input_shape_override);
 }
 

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<size_t>& 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<size_t>& 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<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());
-  TensorPitches original_input_strides(input_dims);
-  TensorPitches original_output_strides(output_dims);
 
-  TArray<int64_t> input_strides(rank);
-  for (auto i = 0; i < rank; i++) {
-    input_strides[i] = original_input_strides[permutations[i]];
-  }
-  TArray<fast_divmod> output_strides(rank);
-  for (auto i = 0; i < rank; i++) {
-    output_strides[i] = fast_divmod(gsl::narrow_cast<int>(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<size_t> new_permutations(permutations);
+  std::vector<int64_t> new_input_dims(input_dims);
+  std::vector<int64_t> 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<int32_t>(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<int64_t> input_shape(new_input_dims);
+  TArray<int64_t> 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<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]];
+    }
+    return Transpose4DImpl(element_size, input_shape, tmp_input_strides, input.DataRaw(),
+                           tmp_output_strides, output.MutableDataRaw(), output.Shape().Size());
+  }
+
+  // General cases
+  TArray<int64_t> input_strides(new_rank);
+  for (auto i = 0; i < new_rank; i++) {
+    input_strides[i] = new_input_strides[new_permutations[i]];
+  }
+
+  TArray<fast_divmod> output_strides(new_rank);
+  for (auto i = 0; i < new_rank; i++) {
+    output_strides[i] = fast_divmod(gsl::narrow_cast<int>(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

onnxruntime/core/providers/cuda/tensor/transpose.h

@@ -23,7 +23,8 @@ class Transpose final : public CudaKernel, public TransposeBase {
                             const std::vector<size_t>& 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<size_t>& permutations,
                             const Tensor& input, Tensor& output, const TensorShape* input_shape_override = nullptr);
 };

onnxruntime/core/providers/cuda/tensor/transpose_impl.cu

@@ -7,9 +7,165 @@
 namespace onnxruntime {
 namespace cuda {
 
+constexpr int TILE_DIM = 16;
+
 template <typename T>
-__global__ void _TransposeKernel(int32_t shape_rank, const TArray<int64_t> input_strides,
-                                 const T* input_data, const TArray<fast_divmod> output_strides, T* output_data, CUDA_LONG N) {
+__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)];
+
+  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<int64_t>& input_dims,
+                      const std::vector<size_t>& 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<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(input_shape[2]/TILE_DIM, input_shape[1]/TILE_DIM, input_shape[0]);
+
+  switch (element_size) {
+    case sizeof(int8_t):
+      Transpose3DKernel<int8_t><<<grid_size, block_size, 0>>>(
+          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>>>(
+          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>>>(
+          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>>>(
+          input_shape, input_strides,
+          reinterpret_cast<const ToCudaType<int64_t>::MappedType*>(input_data),
+          reinterpret_cast<ToCudaType<int64_t>::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 <int element_size>
+__global__ void Transpose4DKernel(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];
+
+  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<const int4*>(input_data);
+  int4* v_output = reinterpret_cast<int4*>(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<int64_t>& input_dims,
+                      const std::vector<size_t>& 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<int64_t>& input_shape, const TArray<int64_t>& input_strides, const void* input_data,
+                       const TArray<int64_t>& 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<sizeof(int8_t)><<<grid_size, block_size, 0>>>(
+          input_strides, input_data,
+          output_strides, output_data, N/num_elements_per_thread);
+      break;
+    case sizeof(int16_t):
+      Transpose4DKernel<sizeof(int16_t)><<<grid_size, block_size, 0>>>(
+          input_strides, input_data,
+          output_strides, output_data, N/num_elements_per_thread);
+      break;
+    case sizeof(int32_t):
+      Transpose4DKernel<sizeof(int32_t)><<<grid_size, block_size, 0>>>(
+          input_strides, input_data,
+          output_strides, output_data, N/num_elements_per_thread);
+      break;
+    case sizeof(int64_t):
+      Transpose4DKernel<sizeof(int64_t)><<<grid_size, block_size, 0>>>(
+          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 <typename T>
+__global__ void TransposeKernel(int32_t shape_rank, const TArray<int64_t> input_strides,
+                                const T* input_data, const TArray<fast_divmod> 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<int64
   int blocksPerGrid = (int)(ceil(static_cast<float>(N) / GridDim::maxThreadsPerBlock));
   switch (element_size) {
     case sizeof(int8_t):
-      _TransposeKernel<int8_t><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+      TransposeKernel<int8_t><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
           shape_rank, input_strides,
           reinterpret_cast<const ToCudaType<int8_t>::MappedType*>(input_data),
           fdm_output_strides,
@@ -40,7 +196,7 @@ Status TransposeImpl(size_t element_size, int32_t shape_rank, const TArray<int64
           N);
       break;
     case sizeof(int16_t):
-      _TransposeKernel<int16_t><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+      TransposeKernel<int16_t><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
           shape_rank, input_strides,
           reinterpret_cast<const ToCudaType<int16_t>::MappedType*>(input_data),
           fdm_output_strides,
@@ -48,7 +204,7 @@ Status TransposeImpl(size_t element_size, int32_t shape_rank, const TArray<int64
           N);
       break;
     case sizeof(int32_t):
-      _TransposeKernel<int32_t><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+      TransposeKernel<int32_t><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
           shape_rank, input_strides,
           reinterpret_cast<const ToCudaType<int32_t>::MappedType*>(input_data),
           fdm_output_strides,
@@ -56,7 +212,7 @@ Status TransposeImpl(size_t element_size, int32_t shape_rank, const TArray<int64
           N);
       break;
     case sizeof(int64_t):
-      _TransposeKernel<int64_t><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+      TransposeKernel<int64_t><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
           shape_rank, input_strides,
           reinterpret_cast<const ToCudaType<int64_t>::MappedType*>(input_data),
           fdm_output_strides,

onnxruntime/core/providers/cuda/tensor/transpose_impl.h

@@ -8,8 +8,17 @@
 namespace onnxruntime {
 namespace cuda {
 
+bool CanDoTranspose3D(int32_t rank, const std::vector<int64_t>& input_dims, const std::vector<size_t>& permutations);
+Status Transpose3DImpl(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);
+bool CanDoTranspose4D(const cudaDeviceProp& prop,
+                      size_t element_size,
+                      int32_t rank,
+                      const std::vector<int64_t>& input_dims,
+                      const std::vector<size_t>& permutations);
+Status Transpose4DImpl(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, int64_t N);
 Status TransposeImpl(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, int64_t N);
-
 }  // namespace cuda
 }  // namespace onnxruntime

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<int64_t>& perm,
+    const std::vector<int64_t>& x_dims,
+    const std::vector<int64_t>& y_dims,
+    double error_tolerance = 1e-4) {
+  CompareOpTester test{"Transpose"};
+
+  RandomValueGenerator random{};
+  const auto X_data = random.Uniform<float>(x_dims, -10.0, 10.0);
+  const auto Y_data = FillZeros<float>(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<int64_t> X_dims{64, 128, 16, 64};
+  const std::vector<int64_t> perm{0, 2, 1, 3};
+  const std::vector<int64_t> Y_dims{64, 16, 128, 64};
+  TestTranspose(perm, X_dims, Y_dims);
+}
+
+TEST(TransposeOpTest, Transpose0231) {
+  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) {
+  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);
+}
+#endif
+
 }  // namespace test
 }  // namespace onnxruntime