From 7e955960f1edfdad1fecd9967042348d281184ac Mon Sep 17 00:00:00 2001
From: Vincent Wang <wangwchpku@outlook.com>
Date: Wed, 12 Aug 2020 16:36:03 +0800
Subject: [PATCH] Optimize Slice Kernel by Removing If-statement (#4753)

* Slice kernel optimization.

* remove space

Co-authored-by: Vincent Wang <weicwang@AiFramework2080ti2.corp.microsoft.com>
---
 .../core/providers/cuda/tensor/slice_impl.cu  | 124 ++++++++++--------
 1 file changed, 72 insertions(+), 52 deletions(-)

diff --git a/onnxruntime/core/providers/cuda/tensor/slice_impl.cu b/onnxruntime/core/providers/cuda/tensor/slice_impl.cu
index df51104194..5a74018852 100644
--- a/onnxruntime/core/providers/cuda/tensor/slice_impl.cu
+++ b/onnxruntime/core/providers/cuda/tensor/slice_impl.cu
@@ -8,36 +8,57 @@
 namespace onnxruntime {
 namespace cuda {
 
-template <bool is_grad, typename T>
-__global__ void _SliceKernel(const int32_t dimension_count,
-                             const TArray<int64_t> starts,
+template <bool is_grad, int DIMS, int NumThreadsPerBlock, int NumElementsPerThread, typename T>
+__global__ void _SliceKernel(const TArray<int64_t> starts,
                              const TArray<int64_t> steps,
                              const TArray<int64_t> input_strides,
                              const TArray<fast_divmod> output_strides,
                              const T* input_data,
                              T* output_data,
                              const CUDA_LONG N) {
-  CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
-  CUDA_LONG input_index = 0;
-  int div;
-  int mod = id;
-  int value = id;
-  int dim = 0;
-#pragma unroll
-  for (; dim < starts.Capacity(); ++dim) {
-    if (dim >= dimension_count - 1) {
-      break;
-    }
+  CUDA_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
+  CUDA_LONG input_indices[NumElementsPerThread];
 
-    output_strides[dim].divmod(value, div, mod);
-    input_index += (starts[dim] + div * steps[dim]) * input_strides[dim];
-    value = mod;
+  CUDA_LONG id = start;
+#pragma unroll
+  for (int i = 0; i < NumElementsPerThread; i++) {
+    if (id < N) {
+      CUDA_LONG input_index = 0;
+      int div;
+      int mod = id;
+      int value = id;
+      int dim = 0;
+#pragma unroll
+      for (; dim < DIMS - 1; ++dim) {
+        output_strides[dim].divmod(value, div, mod);
+        input_index += (starts[dim] + div * steps[dim]) * input_strides[dim];
+        value = mod;
+      }
+      input_index += starts[dim] + mod * steps[dim];
+      input_indices[i] = input_index;
+      id += NumThreadsPerBlock;
+    }
+  }
+
+  if (is_grad) {
+    id = start;
+#pragma unroll
+    for (int i = 0; i < NumElementsPerThread; i++) {
+      if (id < N) {
+        output_data[input_indices[i]] = input_data[id];
+        id += NumThreadsPerBlock;
+      }
+    }
+  } else {
+    id = start;
+#pragma unroll
+    for (int i = 0; i < NumElementsPerThread; i++) {
+      if (id < N) {
+        output_data[id] = input_data[input_indices[i]];
+        id += NumThreadsPerBlock;
+      }
+    }
   }
-  input_index += starts[dim] + mod * steps[dim];
-  if (is_grad)
-    output_data[input_index] = input_data[id];
-  else
-    output_data[id] = input_data[input_index];
 }
 
 Status SliceImpl(const size_t element_size,
@@ -66,6 +87,30 @@ Status SliceImplGrad(const size_t element_size,
                            output_data, N);
 }
 
+#define HANDLE_DIMS(ELEMENT_TYPE, DIMS)                                                     \
+  case DIMS: {                                                                              \
+    _SliceKernel<is_grad, DIMS, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread> \
+        <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(                                \
+            starts, steps, input_strides, output_strides,                                   \
+            reinterpret_cast<const ToCudaType<ELEMENT_TYPE>::MappedType*>(input_data),      \
+            reinterpret_cast<ToCudaType<ELEMENT_TYPE>::MappedType*>(output_data),           \
+            (CUDA_LONG)N);                                                                  \
+  } break
+
+#define HANDLE_ELEMENT_TYPE(ELEMENT_TYPE) \
+  case sizeof(ELEMENT_TYPE): {            \
+    switch (dimension_count) {            \
+      HANDLE_DIMS(ELEMENT_TYPE, 1);       \
+      HANDLE_DIMS(ELEMENT_TYPE, 2);       \
+      HANDLE_DIMS(ELEMENT_TYPE, 3);       \
+      HANDLE_DIMS(ELEMENT_TYPE, 4);       \
+      HANDLE_DIMS(ELEMENT_TYPE, 5);       \
+      HANDLE_DIMS(ELEMENT_TYPE, 6);       \
+      HANDLE_DIMS(ELEMENT_TYPE, 7);       \
+      HANDLE_DIMS(ELEMENT_TYPE, 8);       \
+    }                                     \
+  } break
+
 template <bool is_grad>
 Status SliceImplEx(const size_t element_size,
                    const int32_t dimension_count,
@@ -76,37 +121,12 @@ Status SliceImplEx(const size_t element_size,
                    const void* input_data,
                    void* output_data,
                    const size_t N) {
-  int blocksPerGrid = (int)(ceil(static_cast<float>(N) / GridDim::maxThreadsPerBlock));
-
+  int blocksPerGrid = static_cast<int>(CeilDiv(N, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
   switch (element_size) {
-    case sizeof(int8_t):
-      _SliceKernel<is_grad><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-          dimension_count, starts, steps, input_strides, output_strides,
-          reinterpret_cast<const ToCudaType<int8_t>::MappedType*>(input_data),
-          reinterpret_cast<ToCudaType<int8_t>::MappedType*>(output_data),
-          (CUDA_LONG)N);
-      break;
-    case sizeof(int16_t):
-      _SliceKernel<is_grad><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-          dimension_count, starts, steps, input_strides, output_strides,
-          reinterpret_cast<const ToCudaType<int16_t>::MappedType*>(input_data),
-          reinterpret_cast<ToCudaType<int16_t>::MappedType*>(output_data),
-          (CUDA_LONG)N);
-      break;
-    case sizeof(int32_t):
-      _SliceKernel<is_grad><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-          dimension_count, starts, steps, input_strides, output_strides,
-          reinterpret_cast<const ToCudaType<int32_t>::MappedType*>(input_data),
-          reinterpret_cast<ToCudaType<int32_t>::MappedType*>(output_data),
-          (CUDA_LONG)N);
-      break;
-    case sizeof(int64_t):
-      _SliceKernel<is_grad><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-          dimension_count, starts, steps, input_strides, output_strides,
-          reinterpret_cast<const ToCudaType<int64_t>::MappedType*>(input_data),
-          reinterpret_cast<ToCudaType<int64_t>::MappedType*>(output_data),
-          (CUDA_LONG)N);
-      break;
+    HANDLE_ELEMENT_TYPE(int8_t);
+    HANDLE_ELEMENT_TYPE(int16_t);
+    HANDLE_ELEMENT_TYPE(int32_t);
+    HANDLE_ELEMENT_TYPE(int64_t);
     default:
       return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for Slice operator");
   }