Add FastGelu Cuda Op for Gelu and Add bias fusion (#2293)

* Add FastGelu cuda op

* Add AddBiasGelu for experiment

* Revert "Add AddBiasGelu for experiment"

This reverts commit 5c1ee019858c657e6bb75887265cb85675626e5b.

* Add bias

* Add unit tests

* update comment

* update script

* fix build error

* update coding style

* update for CR feedback
Enable half2 optimization only when cuda arch >= 7.0

* move _Tanh to common.cuh

onnxruntime/contrib_ops/cuda/bert/fast_gelu.cc

@@ -0,0 +1,86 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/common.h"
+#include "core/providers/cuda/cudnn_common.h"
+#include "core/framework/tensorprotoutils.h"
+#include "fast_gelu.h"
+#include "fast_gelu_impl.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+#define REGISTER_KERNEL_TYPED(T)                                  \
+  ONNX_OPERATOR_TYPED_KERNEL_EX(                                  \
+      FastGelu,                                                   \
+      kMSDomain,                                                  \
+      1,                                                          \
+      T,                                                          \
+      kCudaExecutionProvider,                                     \
+      KernelDefBuilder()                                          \
+          .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
+      FastGelu<T>);
+
+REGISTER_KERNEL_TYPED(float)
+REGISTER_KERNEL_TYPED(MLFloat16)
+
+using namespace ONNX_NAMESPACE;
+
+template <typename T>
+FastGelu<T>::FastGelu(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info) {
+}
+
+template <typename T>
+Status FastGelu<T>::ComputeInternal(OpKernelContext* ctx) const {
+  const Tensor* input = ctx->Input<Tensor>(0);
+
+  const auto input_dims = input->Shape().GetDims();
+  if (input_dims.size() < 1) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "Input 0 is expected to have 1 or more dimensions, got ", input_dims.size());
+  }
+
+  size_t num_inputs = OpKernel::Node().InputDefs().size();
+  bool has_bias = (num_inputs == 2);
+
+  int input_length = 1;
+  for (size_t i = 0; i < input_dims.size(); i++) {
+    input_length *= static_cast<int>(input_dims[i]);
+  }
+
+  int bias_length = 0;
+  const Tensor* bias = nullptr;
+  if (has_bias) {
+    bias = ctx->Input<Tensor>(1);
+    const auto bias_dims = bias->Shape().GetDims();
+    if (bias_dims.size() != 1) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "Input 1 is expected to have 1 dimensions, got ", bias_dims.size());
+    }
+    if (bias_dims[0] != input_dims[input_dims.size() - 1]) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "Input 1 dimension 0 should have same length as the last dimension of input 0");
+    }
+    bias_length = static_cast<int>(bias_dims[0]);
+  }
+
+  Tensor* output = ctx->Output(0, input->Shape());
+
+  typedef typename ToCudaType<T>::MappedType CudaT;
+  if (!LaunchFastGeluKernel<CudaT>(nullptr,
+                          input_length,
+                          bias_length,
+                          reinterpret_cast<const CudaT*>(input->template Data<T>()),
+                          has_bias ? reinterpret_cast<const CudaT*>(bias->template Data<T>()) : nullptr,
+                          reinterpret_cast<CudaT*>(output->template MutableData<T>()))) {
+    CUDA_CALL(cudaGetLastError());
+    return Status(common::ONNXRUNTIME, common::FAIL);
+  }
+
+  return Status::OK();
+}
+
+}  //namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/bert/fast_gelu.h

@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+#include "core/providers/cuda/cuda_common.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+using namespace onnxruntime::cuda;
+
+template <typename T>
+class FastGelu final : public CudaKernel {
+ public:
+  FastGelu(const OpKernelInfo& op_kernel_info);
+  Status ComputeInternal(OpKernelContext* ctx) const override;
+};
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.cu

@@ -0,0 +1,102 @@
+/*
+ The implementation of this file is based on gelu plugin in TensorRT demo:
+ https://github.com/NVIDIA/TensorRT/tree/release/5.1/demo/BERT/
+ 
+Copyright 2019 NVIDIA Corporation
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+// Modifications: Add (bias) before Gelu is merged into this op to get better performance.
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/cuda/cuda_common.h"
+#include "core/providers/cuda/cu_inc/common.cuh"
+#include "core/providers/cuda/shared_inc/cuda_call.h"
+#include "fast_gelu_impl.h"
+
+using namespace onnxruntime::cuda;
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+// constants for approximating the normal cdf
+constexpr float A = 0.5;
+
+constexpr float B = 0.7978845608028654;  // sqrt(2.0/M_PI)
+
+constexpr float C = 0.035677408136300125;  // 0.044715 * sqrt(2.0/M_PI)
+
+template <typename T, unsigned TPB>
+__global__ void FastGeluKernel(const T a, const T b, const T c, int input_length, int bias_length, const T* input, const T* bias, T* output) {
+  const int idx = blockIdx.x * TPB + threadIdx.x;
+
+  if (idx < input_length) {
+    const T x = input[idx];
+    const T in = (bias == nullptr) ? x : (x + bias[idx % bias_length]);
+    const T cdf = a + a * _Tanh(in * (c * in * in + b));
+    output[idx] = in * cdf;
+  }
+}
+
+template <unsigned TPB>
+__global__ void FastGeluKernel2(const half2 a, const half2 b, const half2 c, int input_length, int bias_length, const half2* input, const half2* bias, half2* output) {
+// half2 arithmetic functions requires cuda architecture >= 5.3
+#if __CUDA_ARCH__ >= 530
+  const int idx = blockIdx.x * TPB + threadIdx.x;
+
+  if (idx < input_length) {
+    const half2 x = input[idx];
+    const half2 in = (bias == nullptr) ? x : (x + bias[idx % bias_length]);
+    const half2 cdf = a + a * _Tanh(in * (c * in * in + b));
+    output[idx] = in * cdf;
+  }
+#endif
+}
+
+template <>
+bool LaunchFastGeluKernel(cudaStream_t stream, int input_length, int bias_length, const float* input, const float* bias, float* output) {
+  constexpr int blockSize = 256;
+  const int gridSize = (input_length + blockSize - 1) / blockSize;
+  FastGeluKernel<float, blockSize><<<gridSize, blockSize, 0, stream>>>(A, B, C, input_length, bias_length, input, bias, output);
+
+  return CUDA_CALL(cudaPeekAtLastError());
+}
+
+template <>
+bool LaunchFastGeluKernel(cudaStream_t stream, int input_length, int bias_length, const half* input, const half* bias, half* output) {
+  constexpr int blockSize = 256;
+
+  if (0 == (bias_length & 1) && DeviceProp::GetDeviceProps().major >= 7) {
+    const int n = input_length / 2;
+    const int gridSize = (n + blockSize - 1) / blockSize;
+    const half2 A2 = __floats2half2_rn(A, A);
+    const half2 B2 = __floats2half2_rn(B, B);
+    const half2 C2 = __floats2half2_rn(C, C);
+    const half2* input2 = reinterpret_cast<const half2*>(input);
+    const half2* bias2 = reinterpret_cast<const half2*>(bias);
+    half2* output2 = reinterpret_cast<half2*>(output);
+    FastGeluKernel2<blockSize><<<gridSize, blockSize, 0, stream>>>(A2, B2, C2, n, bias_length / 2, input2, bias2, output2);
+  } else {
+    const int gridSize = (input_length + blockSize - 1) / blockSize;
+    FastGeluKernel<half, blockSize><<<gridSize, blockSize, 0, stream>>>(A, B, C, input_length, bias_length, input, bias, output);
+  }
+
+  return CUDA_CALL(cudaPeekAtLastError());
+}
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/bert/fast_gelu_impl.h

@@ -0,0 +1,15 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+template <typename T>
+bool LaunchFastGeluKernel(cudaStream_t stream, int input_length, int bias_length, const T* input, const T* bias, T* output);
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda_contrib_kernels.cc

@@ -9,6 +9,8 @@ using namespace onnxruntime::common;
 namespace onnxruntime {
 namespace contrib {
 namespace cuda {
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, FastGelu);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, FastGelu);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Gelu);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Gelu);
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Gelu);
@@ -48,6 +50,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,
 
 void RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
   static const BuildKernelCreateInfoFn function_table[] = {
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, FastGelu)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, FastGelu)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, Gelu)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, double, Gelu)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, Gelu)>,

onnxruntime/core/graph/contrib_ops/contrib_defs.cc

@@ -314,6 +314,17 @@ void RegisterBertSchemas() {
         updateOutputShape(ctx, 1, mask_index_shape);
       });
 
+  ONNX_CONTRIB_OPERATOR_SCHEMA(FastGelu)
+      .SetDomain(kMSDomain)
+      .SinceVersion(1)
+      .SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL)
+      .SetDoc("Gelu")
+      .Input(0, "X", "input tensor", "T")
+      .Input(1, "bias", "bias tensor", "T", OpSchema::Optional)
+      .Output(0, "Y", "output tensor", "T")
+      .TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float or half tensors.")
+      .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput);
+
   ONNX_CONTRIB_OPERATOR_SCHEMA(SkipLayerNormalization)
       .SetDomain(kMSDomain)
       .SinceVersion(1)

onnxruntime/core/providers/cuda/cu_inc/common.cuh

@@ -7,6 +7,7 @@
 #include <mutex>
 #include <assert.h>
 #include <cuda_runtime.h>
+#include <cuda_fp16.h>
 #include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cuda/shared_inc/cuda_call.h"
 
@@ -154,6 +155,14 @@ __device__ __inline__ double _Tanh(double a) { return tanh(a); }
 template <>
 __device__ __inline__ half _Tanh(half a) { return half(tanhf((float)a)); }
 
+template <>
+__device__ __inline__ half2 _Tanh(half2 a) {
+  float2 tmp = (__half22float2(a));
+  tmp.x = tanhf(tmp.x);
+  tmp.y = tanhf(tmp.y);
+  return __float22half2_rn(tmp);
+}
+
 template <typename T>
 __device__ __inline__ T _Pow(T a, T b);
 

onnxruntime/python/tools/bert/README.md

@@ -1,11 +1,11 @@
 # BERT Model Optimization Tool Overview
 
-This tool converts a BERT ONNX model exported from PyTorch, and generates a optimized model to run faster in NVidia GPU.
+This tool converts a BERT ONNX model exported from PyTorch, and generates an optimized model to run faster in NVidia GPU.
 
 Currently, this script **cannot** process BERT models exported from Tensorflow since the graph has some difference.
 
 ## Export an BERT model from PyTorch
-For example, after using https://github.com/huggingface/transformers to Train a BERT model in PyTorch 1.3, you can use the following function to export ONNX model. 
+For example, after using https://github.com/huggingface/transformers to train a BERT model in PyTorch 1.3, you can use the following function to export ONNX model. 
 
 Please specify do_constant_folding=True. That's required for this tool.
 

onnxruntime/python/tools/bert/bert_model_optimization.py

@@ -350,7 +350,7 @@ class BertOnnxModel(OnnxModel):
 
         # constant node names
         self.normalize_name = "SkipLayerNormalization"
-        self.gelu_name = 'Gelu'
+        self.gelu_name = 'FastGelu'
         self.attention_name = 'Attention'
 
     def get_normalize_nodes(self):
@@ -510,7 +510,7 @@ class BertOnnxModel(OnnxModel):
             if len(matmul_child) != 1 or matmul_child[0].op_type != 'Add':
                 continue
             add_node = matmul_child[0]
-            
+
             children = input_name_to_nodes[add_node.output[0]]
 
             children_types = sorted([child.op_type for child in children])
@@ -525,9 +525,11 @@ class BertOnnxModel(OnnxModel):
             if len(subgraph_nodes) != 5:
                 continue
 
+            nodes_to_remove.append(add_node)
             nodes_to_remove.extend(subgraph_nodes)
+            bias_input = add_node.input[1] if (add_node.input[0] == matmul_node.output[0]) else add_node.input[0]
             gelu_node = onnx.helper.make_node(self.gelu_name,
-                inputs=[add_node.output[0]],
+                inputs=[matmul_node.output[0], bias_input],
                 outputs=[matmul_2.input[0]])
             gelu_node.domain = "com.microsoft"
             nodes_to_add.append(gelu_node)
@@ -882,7 +884,6 @@ def main():
     else:
         bert_model.cast_input_to_int32()
 
-
     if args.float16:
         bert_model.convert_model_float32_to_float16()
 

onnxruntime/test/contrib_ops/fastgelu_op_test.cc

@@ -0,0 +1,172 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "gtest/gtest.h"
+#include "test/common/tensor_op_test_utils.h"
+#include "test/common/cuda_op_test_utils.h"
+#include "test/providers/provider_test_utils.h"
+
+using namespace onnxruntime::test;
+
+namespace onnxruntime {
+namespace test {
+
+const std::vector<float> ComputeGelu(const std::vector<float>& input_data) {
+  std::vector<float> output;
+  output.reserve(input_data.size());
+
+  for (size_t i = 0; i < input_data.size(); i++) {
+    float x = input_data[i];
+    float y = x * (0.5f + 0.5f * tanh(x * (0.035677408136300125f * x * x + 0.7978845608028654f)));
+    output.push_back(y);
+  }
+  return output;
+}
+
+const std::vector<float> AddBias(const std::vector<float>& input_data, const std::vector<float>& bias_data) {
+  size_t bias_length = bias_data.size();
+
+  std::vector<float> output;
+  output.reserve(input_data.size());
+
+  for (size_t i = 0; i < input_data.size(); i++) {
+    output.push_back(input_data[i] + bias_data[i % bias_length]);
+  }
+  return output;
+}
+
+const std::vector<float> GetExpectedResult(const std::vector<float>& input_data, const std::vector<float>& bias_data) {
+  std::vector<float> add_bias_data = AddBias(input_data, bias_data);
+  return ComputeGelu(add_bias_data);
+}
+
+static void RunFastGeluTest(
+    const std::vector<float>& input_data,
+    const std::vector<float>& bias_data,
+    const std::vector<float>& output_data,
+    const std::vector<int64_t>& input_dims,
+    const std::vector<int64_t>& bias_dims,
+    const std::vector<int64_t>& output_dims,
+    bool has_bias = true,
+    bool use_float16 = false) {
+  int min_cuda_architecture = use_float16 ? 530 : 0;
+  if (HasCudaEnvironment(min_cuda_architecture)) {
+    OpTester tester("FastGelu", 1, onnxruntime::kMSDomain);
+
+    if (use_float16) {
+      tester.AddInput<MLFloat16>("X", input_dims, ToFloat16(input_data));
+      if (has_bias) {
+        tester.AddInput<MLFloat16>("bias", bias_dims, ToFloat16(bias_data));
+      }
+      tester.AddOutput<MLFloat16>("Y", output_dims, ToFloat16(output_data));
+    } else {
+      tester.AddInput<float>("X", input_dims, input_data);
+      if (has_bias) {
+        tester.AddInput<float>("bias", bias_dims, bias_data);
+      }
+      tester.AddOutput<float>("Y", output_dims, output_data);
+    }
+
+    std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+    execution_providers.push_back(DefaultCudaExecutionProvider());
+    tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+  }
+}
+
+// This test simulates Gelu in Bert model for float32
+static void RunFastGeluTest(
+    const std::vector<float>& input_data,
+    const std::vector<float>& bias_data,
+    int batch_size,
+    int sequence_length,
+    int hidden_size) {
+  std::vector<float> output_data;
+
+  bool has_bias = (bias_data.size() > 0);
+  if (has_bias) {
+    output_data = GetExpectedResult(input_data, bias_data);
+  } else {
+    output_data = ComputeGelu(input_data);
+  }
+  std::vector<int64_t> input_dims = {batch_size, sequence_length, hidden_size};
+  std::vector<int64_t> bias_dims = {hidden_size};
+  std::vector<int64_t> output_dims = input_dims;
+  RunFastGeluTest(input_data, bias_data, output_data, input_dims, bias_dims, output_dims, has_bias);
+}
+
+TEST(FastGeluTest, FastGeluWithBiasFloat32) {
+  int batch_size = 1;
+  int sequence_length = 2;
+  int hidden_size = 4;
+
+  std::vector<float> input_data = {
+      0.8f, -0.5f, 0.0f, 1.f,
+      0.5f, 0.2f, 0.3f, -0.6f};
+
+  std::vector<float> bias_data = {
+      -0.5f, 0.6f, 1.2f, 2.1f};
+
+  RunFastGeluTest(input_data, bias_data, batch_size, sequence_length, hidden_size);
+}
+
+TEST(FastGeluTest, FastGeluWithoutBiasFloat32) {
+  int batch_size = 1;
+  int sequence_length = 2;
+  int hidden_size = 4;
+
+  std::vector<float> input_data = {
+      0.8f, -0.5f, 0.0f, 1.f,
+      0.5f, 0.2f, 0.3f, -0.6f};
+
+  std::vector<float> bias_data = {};
+
+  RunFastGeluTest(input_data, bias_data, batch_size, sequence_length, hidden_size);
+}
+
+TEST(FastGeluTest, FastGeluWithBiasFloat16) {
+  int batch_size = 1;
+  int sequence_length = 2;
+  int hidden_size = 4;
+
+  std::vector<float> input_data = {
+      0.8f, -0.5f, 0.0f, 1.f,
+      0.5f, 0.2f, 0.3f, -0.6f};
+
+  std::vector<float> bias_data = {
+      -0.5f, 0.6f, 1.2f, 2.1f};
+
+  std::vector<float> output_data = {
+      0.1851806640625f, 0.054046630859375f, 1.0615234375f, 3.095703125f,
+      0, 0.63037109375f, 1.3984375f, 1.3984375f};
+
+  std::vector<int64_t> input_dims = {batch_size, sequence_length, hidden_size};
+  std::vector<int64_t> bias_dims = {hidden_size};
+  std::vector<int64_t> output_dims = input_dims;
+
+  RunFastGeluTest(input_data, bias_data, output_data, input_dims, bias_dims, output_dims, true, true);
+}
+
+TEST(FastGeluTest, FastGeluWithoutBiasFloat16) {
+  int batch_size = 1;
+  int sequence_length = 2;
+  int hidden_size = 4;
+
+  std::vector<float> input_data = {
+      0.8f, -0.5f, 0.0f, 1.f,
+      0.5f, 0.2f, 0.3f, -0.6f};
+
+  std::vector<float> bias_data = {};
+
+  std::vector<float> output_data = {
+      0.63037109375f, -0.154296875f, 0.0f, 0.8408203125f,
+      0.345703125f, 0.11578369140625f, 0.1854248046875f, -0.1646728515625f };
+
+  std::vector<int64_t> input_dims = {batch_size, sequence_length, hidden_size};
+  std::vector<int64_t> bias_dims = {};
+  std::vector<int64_t> output_dims = input_dims;
+
+  RunFastGeluTest(input_data, bias_data, output_data, input_dims, bias_dims, output_dims, false, true);
+}
+
+}  // namespace test
+}  // namespace onnxruntime