Make elementwise op run 4 items per thread (#2335)

Description: Describe your changes. Make elementwise op run 4 items per thread unroll for loop to leverage ILP remove unnessary N==0 check inside elementwise GPU kernel Motivation and Context Why is this change required? What problem does it solve? It can improve the performance of GPU elementwise ops. ~2% performance gain on popular NLP bert model. If it fixes an open issue, please link to the issue here.
2026-07-30 20:18:08 +00:00 · 2019-11-06 17:15:25 -08:00 · 2019-11-06 17:15:25 -08:00 · 6651d2f662
commit 6651d2f662
parent ba0e7daf20
4 changed files with 167 additions and 123 deletions
--- a/onnxruntime/core/providers/cuda/cu_inc/binary_elementwise_impl.cuh
+++ b/onnxruntime/core/providers/cuda/cu_inc/binary_elementwise_impl.cuh
@ -10,7 +10,7 @@ namespace onnxruntime {
 namespace cuda {

 // broadcast by computing output coordinate from offset, using fast_divmod
-template <typename T, typename FuncT, bool lhs_need_compute, bool rhs_need_compute>
+template <typename T, typename FuncT, bool lhs_need_compute, bool rhs_need_compute, int NumThreadsPerBlock, int NumElementsPerThread>
 __global__ void _BinaryElementWise(
    size_t output_rank,
    const int64_t* lhs_padded_strides,
@ -21,50 +21,58 @@ __global__ void _BinaryElementWise(
    T* output_data,
    const FuncT& functor,
    CUDA_LONG N) {
-  if (N == 0)  // special case where there's a dim value of 0 in the output shape
-    return;
+   CUDA_LONG id = NumElementsPerThread * blockDim.x * blockIdx.x + threadIdx.x;

-  CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
-  CUDA_LONG lhs_index = (lhs_need_compute ? 0 : id);
-  CUDA_LONG rhs_index = (rhs_need_compute ? 0 : id);
-  // compute indexes with broadcasting rules: https://github.com/onnx/onnx/blob/master/docs/Broadcasting.md
-  CUDA_LONG offset = id;
-  for (int dim = 0; dim < output_rank; dim++) {
-    int q, r;
-    fdm_output_strides[dim].divmod(offset, q, r);
-    // compute index increase based on stride and broadcast
-    // note that stride[i-1] == stride[i] means dim[i] is 1 (broadcasting)
-    if (lhs_need_compute) {
-      if (lhs_padded_strides[dim] != lhs_padded_strides[dim + 1])
-        lhs_index += static_cast<int>(lhs_padded_strides[dim + 1]) * q;
-    }
+#pragma unroll
+  for (int i = 0; i < NumElementsPerThread; i++) {
+    if (id < N) {
+      CUDA_LONG lhs_index = (lhs_need_compute ? 0 : id);
+      CUDA_LONG rhs_index = (rhs_need_compute ? 0 : id);
+      // compute indexes with broadcasting rules: https://github.com/onnx/onnx/blob/master/docs/Broadcasting.md
+      CUDA_LONG offset = id;
+      for (int dim = 0; dim < output_rank; dim++) {
+        int q, r;
+        fdm_output_strides[dim].divmod(offset, q, r);
+        // compute index increase based on stride and broadcast
+        // note that stride[i-1] == stride[i] means dim[i] is 1 (broadcasting)
+        if (lhs_need_compute) {
+          if (lhs_padded_strides[dim] != lhs_padded_strides[dim + 1])
+            lhs_index += static_cast<int>(lhs_padded_strides[dim + 1]) * q;
+        }

-    if (rhs_need_compute) {
-      if (rhs_padded_strides[dim] != rhs_padded_strides[dim + 1])
-        rhs_index += static_cast<int>(rhs_padded_strides[dim + 1]) * q;
+        if (rhs_need_compute) {
+          if (rhs_padded_strides[dim] != rhs_padded_strides[dim + 1])
+            rhs_index += static_cast<int>(rhs_padded_strides[dim + 1]) * q;
+        }
+        offset = r;
+      }
+      output_data[id] = functor(lhs_data[lhs_index], rhs_data[rhs_index]);
+      id += NumThreadsPerBlock;
    }
-    offset = r;
  }
-  output_data[id] = functor(lhs_data[lhs_index], rhs_data[rhs_index]);
 }

 // for scalar broadcast or non-broadcast case
-template <bool IncL, bool IncR, typename T, typename FuncT>
+template <bool IncL, bool IncR, typename T, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
 __global__ void _BinaryElementWiseSimple(
    const T* lhs_data,
    const T* rhs_data,
    T* output_data,
    const FuncT& func,
    CUDA_LONG N) {
-  if (N == 0)  // special case where there's a dim value of 0 in the output shape
-    return;
+   CUDA_LONG id = NumElementsPerThread * blockDim.x * blockIdx.x + threadIdx.x;

-  CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
-  output_data[id] = func(lhs_data[IncL ? id : 0], rhs_data[IncR ? id : 0]);
+#pragma unroll
+  for (int i = 0; i < NumElementsPerThread; i++) {
+    if (id < N) {
+      output_data[id] = func(lhs_data[IncL ? id : 0], rhs_data[IncR ? id : 0]);
+      id += NumThreadsPerBlock;
+    }
+  }
 }

 // for rhs per-channel broadcast case
-template <typename T, typename FuncT>
+template <typename T, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
 __global__ void _BinaryElementWiseRhsPerChannelBatch1(
    const T* lhs_data,
    const T* rhs_data,
@ -72,12 +80,19 @@ __global__ void _BinaryElementWiseRhsPerChannelBatch1(
    T* output_data,
    FuncT func,
    CUDA_LONG N) {
-  CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
-  CUDA_LONG rhs_id = fdm_H.div(id);
-  output_data[id] = func(lhs_data[id], rhs_data[rhs_id]);
+   CUDA_LONG id = NumElementsPerThread * blockDim.x * blockIdx.x + threadIdx.x;
+
+#pragma unroll
+  for (int i = 0; i < NumElementsPerThread; i++) {
+    if (id < N) {
+      CUDA_LONG rhs_id = fdm_H.div(id);
+      output_data[id] = func(lhs_data[id], rhs_data[rhs_id]);
+      id += NumThreadsPerBlock;
+    }
+  }
 }

-template <typename T, typename FuncT>
+template <typename T, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
 __global__ void _BinaryElementWiseRhsPerChannelBatchN(
    const T* lhs_data,
    const T* rhs_data,
@ -86,12 +101,19 @@ __global__ void _BinaryElementWiseRhsPerChannelBatchN(
    T* output_data,
    FuncT func,
    CUDA_LONG N) {
-  CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
-  CUDA_LONG rhs_id = fdm_H.div(id);
-  int q, r;
-  fdm_C.divmod(rhs_id, q, r);
-  rhs_id = r;
-  output_data[id] = func(lhs_data[id], rhs_data[rhs_id]);
+   CUDA_LONG id = NumElementsPerThread * blockDim.x * blockIdx.x + threadIdx.x;
+
+#pragma unroll
+  for (int i = 0; i < NumElementsPerThread; i++) {
+    if (id < N) {
+      CUDA_LONG rhs_id = fdm_H.div(id);
+      int q, r;
+      fdm_C.divmod(rhs_id, q, r);
+      rhs_id = r;
+      output_data[id] = func(lhs_data[id], rhs_data[rhs_id]);
+      id += NumThreadsPerBlock;
+    }
+  }
 }

 template <typename T, typename FuncT>
@ -104,14 +126,15 @@ void BinaryElementWiseNoBroadcastImpl(
  if (count == 0)  // special case where there's a dim value of 0 in the output shape
    return;

-  int blocksPerGrid = (int)(ceil(static_cast<float>(count) / GridDim::maxThreadsPerBlock));
+  int blocksPerGrid = static_cast<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
  CUDA_LONG N = static_cast<CUDA_LONG>(count);
-  _BinaryElementWiseSimple<true, true, T, FuncT><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-      lhs_data,
-      rhs_data,
-      output_data,
-      func,
-      N);
+  _BinaryElementWiseSimple<true, true, T, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
+      <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+          lhs_data,
+          rhs_data,
+          output_data,
+          func,
+          N);
 }

 template <typename T, typename FuncT>
@ -130,80 +153,88 @@ void BinaryElementWiseImpl(
  if (count == 0)  // special case where there's a dim value of 0 in the output shape
    return;

-  int blocksPerGrid = (int)(ceil(static_cast<float>(count) / GridDim::maxThreadsPerBlock));
+  int blocksPerGrid = static_cast<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
  CUDA_LONG N = static_cast<CUDA_LONG>(count);
  if (output_rank_or_simple_broadcast == static_cast<size_t>(SimpleBroadcast::NoBroadcast)) {
-    _BinaryElementWiseSimple<true, true, T, FuncT><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-        lhs_data,
-        rhs_data,
-        output_data,
-        func,
-        N);
+    _BinaryElementWiseSimple<true, true, T, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
+        <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+            lhs_data,
+            rhs_data,
+            output_data,
+            func,
+            N);
  } else if (output_rank_or_simple_broadcast == static_cast<size_t>(SimpleBroadcast::LeftScalar)) {
-    _BinaryElementWiseSimple<false, true, T, FuncT><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-        lhs_data,
-        rhs_data,
-        output_data,
-        func,
-        N);
+    _BinaryElementWiseSimple<false, true, T, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
+        <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+            lhs_data,
+            rhs_data,
+            output_data,
+            func,
+            N);
  } else if (output_rank_or_simple_broadcast == static_cast<size_t>(SimpleBroadcast::RightScalar)) {
-    _BinaryElementWiseSimple<true, false, T, FuncT><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+    _BinaryElementWiseSimple<true, false, T, FuncT, GridDim::maxThreadsPerBlock,
+                             GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
        lhs_data,
        rhs_data,
        output_data,
        func,
        N);
  } else if (output_rank_or_simple_broadcast == static_cast<size_t>(SimpleBroadcast::RightPerChannelBatch1)) {
-    _BinaryElementWiseRhsPerChannelBatch1<T, FuncT><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-        lhs_data,
-        rhs_data,
-        fdm_H,
-        output_data,
-        func,
-        N);
+    _BinaryElementWiseRhsPerChannelBatch1<T, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
+        <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+            lhs_data,
+            rhs_data,
+            fdm_H,
+            output_data,
+            func,
+            N);
  } else if (output_rank_or_simple_broadcast == static_cast<size_t>(SimpleBroadcast::RightPerChannelBatchN)) {
-    _BinaryElementWiseRhsPerChannelBatchN<T, FuncT><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-        lhs_data,
-        rhs_data,
-        fdm_H,
-        fdm_C,
-        output_data,
-        func,
-        N);
+    _BinaryElementWiseRhsPerChannelBatchN<T, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
+        <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+            lhs_data,
+            rhs_data,
+            fdm_H,
+            fdm_C,
+            output_data,
+            func,
+            N);
  } else {
    if (lhs_padded_strides && rhs_padded_strides)
-      _BinaryElementWise<T, FuncT, true, true><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-          output_rank_or_simple_broadcast,
-          lhs_padded_strides,
-          lhs_data,
-          rhs_padded_strides,
-          rhs_data,
-          fdm_output_strides,
-          output_data,
-          func,
-          N);
+      _BinaryElementWise<T, FuncT, true, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
+          <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+              output_rank_or_simple_broadcast,
+              lhs_padded_strides,
+              lhs_data,
+              rhs_padded_strides,
+              rhs_data,
+              fdm_output_strides,
+              output_data,
+              func,
+              N);
    else if (lhs_padded_strides)
-      _BinaryElementWise<T, FuncT, true, false><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-          output_rank_or_simple_broadcast,
-          lhs_padded_strides,
-          lhs_data,
-          rhs_padded_strides,
-          rhs_data,
-          fdm_output_strides,
-          output_data,
-          func,
-          N);
+      _BinaryElementWise<T, FuncT, true, false, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
+          <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+              output_rank_or_simple_broadcast,
+              lhs_padded_strides,
+              lhs_data,
+              rhs_padded_strides,
+              rhs_data,
+              fdm_output_strides,
+              output_data,
+              func,
+              N);
    else
-      _BinaryElementWise<T, FuncT, false, true><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-          output_rank_or_simple_broadcast,
-          lhs_padded_strides,
-          lhs_data,
-          rhs_padded_strides,
-          rhs_data,
-          fdm_output_strides,
-          output_data,
-          func,
-          N);
+      _BinaryElementWise<T, FuncT, false, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
+          <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+              output_rank_or_simple_broadcast,
+              lhs_padded_strides,
+              lhs_data,
+              rhs_padded_strides,
+              rhs_data,
+              fdm_output_strides,
+              output_data,
+              func,
+              N);
  }
 }

--- a/onnxruntime/core/providers/cuda/cu_inc/common.cuh
+++ b/onnxruntime/core/providers/cuda/cu_inc/common.cuh
@ -208,8 +208,9 @@ static INT CeilDiv(INT a, INT2 b)  // ceil(a/b)

 struct GridDim {
  enum : CUDA_LONG {
-    maxThreadsPerBlock = 1024,  // use this many threads per block
-    maxWarpsPerBlock = 32,      // use this many warps per block. This means 1024 threads for warpSize=32
+    maxThreadsPerBlock = 256,  // max threads per block
+    maxWarpsPerBlock = 32,     // max warps per block
+    maxElementsPerThread = 4,  // max element processed per thread
  };

  // use these for launching
--- a/onnxruntime/core/providers/cuda/cu_inc/unary_elementwise_impl.cuh
+++ b/onnxruntime/core/providers/cuda/cu_inc/unary_elementwise_impl.cuh
@ -9,17 +9,21 @@
 namespace onnxruntime {
 namespace cuda {

-template <typename InT, typename OutT, typename FuncT>
+template <typename InT, typename OutT, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
 __global__ void _UnaryElementWise(
    const InT* input_data,
    OutT* output_data,
    const FuncT& functor,
    CUDA_LONG N) {
-  if (N == 0)  // special case where there's a dim value of 0 in the shape
-    return;
+  CUDA_LONG id = NumElementsPerThread * blockDim.x * blockIdx.x + threadIdx.x;

-  CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
-  output_data[id] = functor(input_data[id]);
+#pragma unroll
+  for (int i = 0; i < NumElementsPerThread; i++) {
+    if (id < N) {
+      output_data[id] = functor(input_data[id]);
+      id += NumThreadsPerBlock;
+    }
+  }
 }

 template <typename InT, typename OutT, typename FuncT>
@ -31,13 +35,14 @@ void UnaryElementWiseImpl(
  if (count == 0)  // special case where there's a dim value of 0 in the shape
    return;

-  int blocksPerGrid = (int)(ceil(static_cast<float>(count) / GridDim::maxThreadsPerBlock));
+  int blocksPerGrid = static_cast<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
  CUDA_LONG N = static_cast<CUDA_LONG>(count);
-  _UnaryElementWise<InT, OutT, FuncT><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
-      input_data,
-      output_data,
-      func,
-      N);
+  _UnaryElementWise<InT, OutT, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
+      <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+          input_data,
+          output_data,
+          func,
+          N);
 }

 }  // namespace cuda
--- a/onnxruntime/core/providers/cuda/cuda_utils.cu
+++ b/onnxruntime/core/providers/cuda/cuda_utils.cu
@ -10,22 +10,29 @@
 namespace onnxruntime {
 namespace cuda {

-template <typename T>
+template <typename T, int NumThreadsPerBlock, int NumElementsPerThread>
 __global__ void _Fill(
    T* output_data,
    T val,
    CUDA_LONG N) {
-  CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N);
-  output_data[id] = val;
+  CUDA_LONG id = NumElementsPerThread * blockDim.x * blockIdx.x + threadIdx.x;
+
+#pragma unroll
+  for (int i = 0; i < NumElementsPerThread; i++) {
+    if (id < N) {
+      output_data[id] = val;
+      id += NumThreadsPerBlock;
+    }
+  }
 }

 template <typename T>
 void Fill(T* output, T value, int64_t count) {
-  int blocksPerGrid = (int)(ceil(static_cast<float>(count) / GridDim::maxThreadsPerBlock));
+  int blocksPerGrid = static_cast<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
  CUDA_LONG N = static_cast<CUDA_LONG>(count);
-  _Fill<T><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(output, value, N);
+  _Fill<T, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>
+      <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(output, value, N);
 }
-
 template <typename T>
 class ConstantBufferImpl : public IConstantBuffer<T> {
 public: