chore: import upstream snapshot with attribution

test/cpp/fluid/fused/CMakeLists.txt

@@ -0,0 +1,31 @@
+if(WITH_GPU OR WITH_ROCM)
+  # fusion_group
+  if(NOT APPLE AND NOT WIN32)
+    paddle_test(test_fusion_group_op SRCS fusion_group_op_test.cc)
+  endif()
+  if(NOT WITH_ROCM)
+    nv_test(
+      test_fused_residual_dropout_bias
+      SRCS fused_residual_dropout_bias_test.cu
+      DEPS tensor op_registry dropout_op generated_op phi common)
+    nv_test(
+      test_fused_dropout_act_bias
+      SRCS fused_dropout_act_bias_test.cu
+      DEPS tensor op_registry dropout_op generated_op phi common)
+    nv_test(
+      test_fused_layernorm_residual_dropout_bias
+      SRCS fused_layernorm_residual_dropout_bias_test.cu
+      DEPS tensor
+           op_registry
+           dropout_op
+           generated_op
+           phi
+           common
+           ${CINN_DEPS})
+    nv_test(
+      test_cudnn_norm_conv
+      SRCS cudnn_norm_conv_test.cc
+      DEPS generated_op tensor op_registry phi common)
+    paddle_test(test_cudnn_bn_add_relu SRCS cudnn_bn_add_relu_test.cc)
+  endif()
+endif()

test/cpp/fluid/fused/cudnn_bn_add_relu_test.cc

@@ -0,0 +1,947 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+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. */
+
+#include <random>
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "paddle/common/flags.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+#include "paddle/fluid/framework/program_desc.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/phi/common/float16.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/fusion/gpu/cudnn_bn_stats_finalize.cu.h"
+#include "paddle/phi/kernels/fusion/gpu/cudnn_scale_bias_add_relu.cu.h"
+
+COMMON_DECLARE_bool(cudnn_batchnorm_spatial_persistent);
+
+namespace framework = paddle::framework;
+namespace platform = paddle::platform;
+
+template <typename T>
+void InitRandomTensor(const std::vector<int64_t> &dims,
+                      phi::DenseTensor *cpu_out) {
+  T *cpu_out_ptr =
+      cpu_out->mutable_data<T>(common::make_ddim(dims), phi::CPUPlace());
+  std::default_random_engine random(0);
+  std::uniform_real_distribution<float> dis(-1.0, 1.0);
+  for (int i = 0; i < cpu_out->numel(); ++i) {
+    cpu_out_ptr[i] = static_cast<T>(dis(random));
+  }
+}
+
+template <typename T>
+void InitConstantTensor(const std::vector<int64_t> &dims,
+                        T value,
+                        phi::DenseTensor *cpu_out) {
+  T *cpu_out_ptr =
+      cpu_out->mutable_data<T>(common::make_ddim(dims), phi::CPUPlace());
+  for (int i = 0; i < cpu_out->numel(); ++i) {
+    cpu_out_ptr[i] = value;
+  }
+}
+
+template <typename T>
+void CheckOutput(std::string name,
+                 const phi::DenseTensor &cpu_res,
+                 const phi::DenseTensor &cpu_base,
+                 float diff,
+                 bool is_relative_atol = false) {
+  if (cpu_res.dims().size() == cpu_base.dims().size()) {
+    EXPECT_EQ(cpu_res.dims(), cpu_base.dims());
+  } else {
+    EXPECT_EQ(cpu_res.numel(), cpu_base.numel());
+  }
+
+  const T *cpu_res_ptr = cpu_res.data<T>();
+  const T *cpu_base_ptr = cpu_base.data<T>();
+  float max_diff = 0;
+  int index = 0;
+  for (int i = 0; i < cpu_res.numel(); ++i) {
+    float cur_diff;
+    if (is_relative_atol) {
+      cur_diff = static_cast<float>(
+          std::abs((cpu_res_ptr[i] - cpu_base_ptr[i]) / cpu_base_ptr[i]));
+      EXPECT_LT(static_cast<float>(std::abs((cpu_res_ptr[i] - cpu_base_ptr[i]) /
+                                            cpu_base_ptr[i])),
+                diff);
+    } else {
+      cur_diff = static_cast<float>(std::abs(cpu_res_ptr[i] - cpu_base_ptr[i]));
+      EXPECT_LT(static_cast<float>(std::abs(cpu_res_ptr[i] - cpu_base_ptr[i])),
+                diff);
+    }
+    if (cur_diff > max_diff) {
+      max_diff = cur_diff;
+      index = i;
+    }
+  }
+  std::string error_type = is_relative_atol ? "relative" : "absolute";
+  LOG(INFO) << "[" << name << "] The dims is [" << cpu_res.dims()
+            << "], maximum " << error_type << " error is " << max_diff << ": "
+            << cpu_res_ptr[index] << " vs " << cpu_base_ptr[index];
+}
+
+template <typename T>
+void ComputeSumAndSquareSum(const phi::DenseTensor &cpu_x,
+                            phi::DenseTensor *cpu_sum,
+                            phi::DenseTensor *cpu_sum_of_square) {
+  // x is in NHWC format.
+  const auto &dims = cpu_x.dims();
+  int64_t c = dims[3];
+
+  const T *cpu_x_ptr = cpu_x.data<T>();
+  float *cpu_sum_ptr =
+      cpu_sum->mutable_data<float>({1, 1, 1, c}, phi::CPUPlace());
+  float *cpu_sum_square_ptr =
+      cpu_sum_of_square->mutable_data<float>({1, 1, 1, c}, phi::CPUPlace());
+
+  for (int j = 0; j < c; ++j) {
+    float tmp_sum = 0.0f;
+    float tmp_sum_of_squares = 0.0f;
+    for (int i = 0; i < cpu_x.numel() / c; ++i) {
+      float tmp_x = static_cast<float>(cpu_x_ptr[i * c + j]);
+      tmp_sum += tmp_x;
+      tmp_sum_of_squares += tmp_x * tmp_x;
+    }
+    cpu_sum_ptr[j] = tmp_sum;
+    cpu_sum_square_ptr[j] = tmp_sum_of_squares;
+  }
+}
+
+template <typename T>
+void ComputeInplaceAdd(const phi::DenseTensor &cpu_x, phi::DenseTensor *cpu_y) {
+  EXPECT_EQ(cpu_x.dims(), cpu_y->dims());
+
+  const T *cpu_x_ptr = cpu_x.data<T>();
+  T *cpu_y_ptr = cpu_y->data<T>();
+  for (int64_t i = 0; i < cpu_x.numel(); ++i) {
+    cpu_y_ptr[i] += cpu_x_ptr[i];
+  }
+}
+
+template <typename T>
+void ComputeInplaceRelu(phi::DenseTensor *cpu_x) {
+  T *cpu_x_ptr = cpu_x->data<T>();
+  for (int64_t i = 0; i < cpu_x->numel(); ++i) {
+    cpu_x_ptr[i] =
+        cpu_x_ptr[i] > static_cast<T>(0) ? cpu_x_ptr[i] : static_cast<T>(0);
+  }
+}
+
+void ComputeBatchNormForward(const phi::GPUContext &ctx,
+                             const phi::DenseTensor &cpu_x,
+                             const phi::DenseTensor &cpu_scale,
+                             const phi::DenseTensor &cpu_bias,
+                             phi::DenseTensor *cpu_mean,
+                             phi::DenseTensor *cpu_var,
+                             phi::DenseTensor *cpu_saved_mean,
+                             phi::DenseTensor *cpu_saved_var,
+                             phi::DenseTensor *cpu_y,
+                             phi::DenseTensor *saved_reserve_space) {
+  framework::Scope scope;
+  auto *x = scope.Var("X")->GetMutable<phi::DenseTensor>();
+  auto *scale = scope.Var("Scale")->GetMutable<phi::DenseTensor>();
+  auto *bias = scope.Var("Bias")->GetMutable<phi::DenseTensor>();
+  auto *mean = scope.Var("Mean")->GetMutable<phi::DenseTensor>();
+  auto *var = scope.Var("Variance")->GetMutable<phi::DenseTensor>();
+  auto *y = scope.Var("Y")->GetMutable<phi::DenseTensor>();
+  auto *saved_mean = scope.Var("SavedMean")->GetMutable<phi::DenseTensor>();
+  auto *saved_var = scope.Var("SavedVariance")->GetMutable<phi::DenseTensor>();
+  auto *reserve_space =
+      scope.Var("ReserveSpace")->GetMutable<phi::DenseTensor>();
+
+  auto place = ctx.GetPlace();
+  paddle::framework::TensorCopySync(cpu_x, place, x);
+  paddle::framework::TensorCopySync(cpu_scale, place, scale);
+  paddle::framework::TensorCopySync(cpu_bias, place, bias);
+  paddle::framework::TensorCopySync(*cpu_mean, place, mean);
+  paddle::framework::TensorCopySync(*cpu_var, place, var);
+
+  int64_t channels = x->dims()[3];
+  scale->Resize({channels});
+  bias->Resize({channels});
+  mean->Resize({channels});
+  var->Resize({channels});
+
+  framework::AttributeMap attrs;
+  std::string data_layout = "NHWC";
+  attrs.insert({"data_layout", data_layout});
+
+  auto op =
+      framework::OpRegistry::CreateOp("batch_norm",
+                                      {{"X", {"X"}},
+                                       {"Scale", {"Scale"}},
+                                       {"Bias", {"Bias"}},
+                                       {"Mean", {"Mean"}},
+                                       {"Variance", {"Variance"}}},
+                                      {{"Y", {"Y"}},
+                                       {"MeanOut", {"Mean"}},
+                                       {"VarianceOut", {"Variance"}},
+                                       {"SavedMean", {"SavedMean"}},
+                                       {"SavedVariance", {"SavedVariance"}},
+                                       {"ReserveSpace", {"ReserveSpace"}}},
+                                      attrs);
+  op->Run(scope, ctx.GetPlace());
+
+  paddle::framework::TensorCopySync(*y, phi::CPUPlace(), cpu_y);
+  paddle::framework::TensorCopySync(*mean, phi::CPUPlace(), cpu_mean);
+  paddle::framework::TensorCopySync(*var, phi::CPUPlace(), cpu_var);
+  paddle::framework::TensorCopySync(
+      *saved_mean, phi::CPUPlace(), cpu_saved_mean);
+  paddle::framework::TensorCopySync(*saved_var, phi::CPUPlace(), cpu_saved_var);
+  // reserved_space will stay on GPU and used in grad op.
+  saved_reserve_space->ShareDataWith(*reserve_space);
+}
+
+void ComputeFusedBNAddReluForward(const phi::GPUContext &ctx,
+                                  const phi::DenseTensor &cpu_x,
+                                  const phi::DenseTensor &cpu_z,
+                                  const phi::DenseTensor &cpu_scale,
+                                  const phi::DenseTensor &cpu_bias,
+                                  phi::DenseTensor *cpu_mean,
+                                  phi::DenseTensor *cpu_var,
+                                  phi::DenseTensor *cpu_saved_mean,
+                                  phi::DenseTensor *cpu_saved_var,
+                                  phi::DenseTensor *cpu_y,
+                                  phi::DenseTensor *saved_reserve_space) {
+  framework::Scope scope;
+  auto *x = scope.Var("X")->GetMutable<phi::DenseTensor>();
+  auto *z = scope.Var("Z")->GetMutable<phi::DenseTensor>();
+  auto *scale = scope.Var("Scale")->GetMutable<phi::DenseTensor>();
+  auto *bias = scope.Var("Bias")->GetMutable<phi::DenseTensor>();
+  auto *mean = scope.Var("Mean")->GetMutable<phi::DenseTensor>();
+  auto *var = scope.Var("Variance")->GetMutable<phi::DenseTensor>();
+  auto *y = scope.Var("Y")->GetMutable<phi::DenseTensor>();
+  auto *saved_mean = scope.Var("SavedMean")->GetMutable<phi::DenseTensor>();
+  auto *saved_var = scope.Var("SavedVariance")->GetMutable<phi::DenseTensor>();
+  auto *reserve_space =
+      scope.Var("ReserveSpace")->GetMutable<phi::DenseTensor>();
+
+  auto place = ctx.GetPlace();
+  paddle::framework::TensorCopySync(cpu_x, place, x);
+  paddle::framework::TensorCopySync(cpu_z, place, z);
+  paddle::framework::TensorCopySync(cpu_scale, place, scale);
+  paddle::framework::TensorCopySync(cpu_bias, place, bias);
+  paddle::framework::TensorCopySync(*cpu_mean, place, mean);
+  paddle::framework::TensorCopySync(*cpu_var, place, var);
+
+  int64_t channels = x->dims()[3];
+  scale->Resize({channels});
+  bias->Resize({channels});
+  mean->Resize({channels});
+  var->Resize({channels});
+
+  framework::AttributeMap attrs;
+
+  auto op =
+      framework::OpRegistry::CreateOp("fused_bn_add_activation",
+                                      {{"X", {"X"}},
+                                       {"Z", {"Z"}},
+                                       {"Scale", {"Scale"}},
+                                       {"Bias", {"Bias"}},
+                                       {"Mean", {"Mean"}},
+                                       {"Variance", {"Variance"}}},
+                                      {{"Y", {"Y"}},
+                                       {"MeanOut", {"Mean"}},
+                                       {"VarianceOut", {"Variance"}},
+                                       {"SavedMean", {"SavedMean"}},
+                                       {"SavedVariance", {"SavedVariance"}},
+                                       {"ReserveSpace", {"ReserveSpace"}}},
+                                      attrs);
+  op->Run(scope, ctx.GetPlace());
+
+  paddle::framework::TensorCopySync(*y, phi::CPUPlace(), cpu_y);
+  paddle::framework::TensorCopySync(*mean, phi::CPUPlace(), cpu_mean);
+  paddle::framework::TensorCopySync(*var, phi::CPUPlace(), cpu_var);
+  paddle::framework::TensorCopySync(
+      *saved_mean, phi::CPUPlace(), cpu_saved_mean);
+  paddle::framework::TensorCopySync(*saved_var, phi::CPUPlace(), cpu_saved_var);
+  // reserved_space will stay on GPU and used in grad op.
+  saved_reserve_space->ShareDataWith(*reserve_space);
+}
+
+void ComputeFusedBNAddReluBackward(const phi::GPUContext &ctx,
+                                   const phi::DenseTensor &cpu_dy,
+                                   const phi::DenseTensor &cpu_x,
+                                   const phi::DenseTensor &cpu_scale,
+                                   const phi::DenseTensor &cpu_bias,
+                                   const phi::DenseTensor &cpu_saved_mean,
+                                   const phi::DenseTensor &cpu_saved_var,
+                                   const phi::DenseTensor &cpu_y,
+                                   const phi::DenseTensor &saved_reserve_space,
+                                   phi::DenseTensor *cpu_dx,
+                                   phi::DenseTensor *cpu_dz,
+                                   phi::DenseTensor *cpu_dscale,
+                                   phi::DenseTensor *cpu_dbias) {
+  framework::Scope scope;
+  auto *x = scope.Var("X")->GetMutable<phi::DenseTensor>();
+  auto *y = scope.Var("Y")->GetMutable<phi::DenseTensor>();
+  auto *dy = scope.Var("Y@GRAD")->GetMutable<phi::DenseTensor>();
+  auto *scale = scope.Var("Scale")->GetMutable<phi::DenseTensor>();
+  auto *bias = scope.Var("Bias")->GetMutable<phi::DenseTensor>();
+  auto *saved_mean = scope.Var("SavedMean")->GetMutable<phi::DenseTensor>();
+  auto *saved_var = scope.Var("SavedVariance")->GetMutable<phi::DenseTensor>();
+  auto *reserve_space =
+      scope.Var("ReserveSpace")->GetMutable<phi::DenseTensor>();
+  auto *dx = scope.Var("X@GRAD")->GetMutable<phi::DenseTensor>();
+  auto *dz = scope.Var("Z@GRAD")->GetMutable<phi::DenseTensor>();
+  auto *dscale = scope.Var("Scale@GRAD")->GetMutable<phi::DenseTensor>();
+  auto *dbias = scope.Var("Bias@GRAD")->GetMutable<phi::DenseTensor>();
+
+  auto place = ctx.GetPlace();
+  paddle::framework::TensorCopySync(cpu_x, place, x);
+  paddle::framework::TensorCopySync(cpu_y, place, y);
+  paddle::framework::TensorCopySync(cpu_dy, place, dy);
+  paddle::framework::TensorCopySync(cpu_scale, place, scale);
+  paddle::framework::TensorCopySync(cpu_bias, place, bias);
+  paddle::framework::TensorCopySync(cpu_saved_mean, place, saved_mean);
+  paddle::framework::TensorCopySync(cpu_saved_var, place, saved_var);
+  reserve_space->ShareDataWith(saved_reserve_space);
+
+  int64_t channels = x->dims()[3];
+  scale->Resize({channels});
+  bias->Resize({channels});
+  saved_mean->Resize({channels});
+  saved_var->Resize({channels});
+
+  framework::AttributeMap attrs;
+  float momentum = 0.9;
+  float epsilon = 1e-5;
+  std::string act_type = "relu";
+  attrs.insert({"momentum", momentum});
+  attrs.insert({"epsilon", epsilon});
+  attrs.insert({"act_type", act_type});
+
+  auto op =
+      framework::OpRegistry::CreateOp("fused_bn_add_activation_grad",
+                                      {{"X", {"X"}},
+                                       {"Y", {"Y"}},
+                                       {"Y@GRAD", {"Y@GRAD"}},
+                                       {"Scale", {"Scale"}},
+                                       {"Bias", {"Bias"}},
+                                       {"SavedMean", {"SavedMean"}},
+                                       {"SavedVariance", {"SavedVariance"}},
+                                       {"ReserveSpace", {"ReserveSpace"}}},
+                                      {{"X@GRAD", {"X@GRAD"}},
+                                       {"Z@GRAD", {"Z@GRAD"}},
+                                       {"Scale@GRAD", {"Scale@GRAD"}},
+                                       {"Bias@GRAD", {"Bias@GRAD"}}},
+                                      attrs);
+  op->Run(scope, ctx.GetPlace());
+
+  paddle::framework::TensorCopySync(*dx, phi::CPUPlace(), cpu_dx);
+  paddle::framework::TensorCopySync(*dz, phi::CPUPlace(), cpu_dz);
+  paddle::framework::TensorCopySync(*dscale, phi::CPUPlace(), cpu_dscale);
+  paddle::framework::TensorCopySync(*dbias, phi::CPUPlace(), cpu_dbias);
+}
+
+template <typename T>
+class CudnnBNAddReluTester {
+ public:
+  CudnnBNAddReluTester(int batch_size,
+                       int height,
+                       int width,
+                       int channels,
+                       std::string act_type,
+                       bool fuse_add,
+                       bool has_shortcut) {
+    batch_size_ = batch_size;
+    height_ = height;
+    width_ = width;
+    channels_ = channels;
+    ele_count_ = batch_size_ * height_ * width_;
+    act_type_ = act_type;
+    fuse_add_ = fuse_add;
+    has_shortcut_ = has_shortcut;
+    SetUp();
+  }
+
+  ~CudnnBNAddReluTester() = default;
+
+  void CheckForward(float diff, bool is_relative_atol = false) {
+    LOG(INFO) << "[CheckForward, diff=" << diff
+              << ", is_relative_atol=" << is_relative_atol
+              << "] act_type=" << act_type_ << ", fuse_add=" << fuse_add_
+              << ", has_shortcut=" << has_shortcut_;
+    phi::GPUContext *ctx = static_cast<phi::GPUContext *>(
+        phi::DeviceContextPool::Instance().Get(phi::GPUPlace(0)));
+
+    auto select = [&](phi::DenseTensor *in) {
+      return has_shortcut_ ? in : nullptr;
+    };
+
+    phi::DenseTensor cpu_mean_base_x;
+    phi::DenseTensor cpu_var_base_x;
+    phi::DenseTensor cpu_mean_base_z;
+    phi::DenseTensor cpu_var_base_z;
+    if (!has_shortcut_ && fuse_add_ && (act_type_ == "relu")) {
+      BaselineForwardFusedBNAddRelu(*ctx,
+                                    &cpu_mean_base_x,
+                                    &cpu_var_base_x,
+                                    &cpu_saved_mean_base_x_,
+                                    &cpu_saved_var_base_x_,
+                                    &cpu_y_base_,
+                                    &saved_reserve_space_x_);
+    } else {
+      BaselineForward(*ctx,
+                      &cpu_mean_base_x,
+                      &cpu_var_base_x,
+                      &cpu_saved_mean_base_x_,
+                      &cpu_saved_var_base_x_,
+                      &cpu_y_base_,
+                      &saved_reserve_space_x_,
+                      select(&cpu_mean_base_z),
+                      select(&cpu_var_base_z),
+                      select(&cpu_saved_mean_base_z_),
+                      select(&cpu_saved_var_base_z_),
+                      select(&saved_reserve_space_z_));
+    }
+
+    phi::DenseTensor cpu_mean_x;
+    phi::DenseTensor cpu_var_x;
+    phi::DenseTensor cpu_y;
+    phi::DenseTensor cpu_mean_z;
+    phi::DenseTensor cpu_var_z;
+    FusedForward(*ctx,
+                 &cpu_mean_x,
+                 &cpu_var_x,
+                 &cpu_saved_mean_x_,
+                 &cpu_saved_var_x_,
+                 &cpu_y,
+                 &cpu_bitmask_,
+                 select(&cpu_mean_z),
+                 select(&cpu_var_z),
+                 select(&cpu_saved_mean_z_),
+                 select(&cpu_saved_var_z_));
+
+    CheckOutput<float>(
+        "Mean", cpu_mean_x, cpu_mean_base_x, diff, is_relative_atol);
+    CheckOutput<float>(
+        "Variance", cpu_var_x, cpu_var_base_x, diff, is_relative_atol);
+    CheckOutput<float>("SavedMean",
+                       cpu_saved_mean_x_,
+                       cpu_saved_mean_base_x_,
+                       diff,
+                       is_relative_atol);
+    CheckOutput<float>("SavedVariance",
+                       cpu_saved_var_x_,
+                       cpu_saved_var_base_x_,
+                       diff,
+                       is_relative_atol);
+    if (has_shortcut_) {
+      CheckOutput<float>(
+          "MeanZ", cpu_mean_z, cpu_mean_base_z, diff, is_relative_atol);
+      CheckOutput<float>(
+          "VarianceZ", cpu_var_z, cpu_var_base_z, diff, is_relative_atol);
+      CheckOutput<float>("SavedMeanZ",
+                         cpu_saved_mean_z_,
+                         cpu_saved_mean_base_z_,
+                         diff,
+                         is_relative_atol);
+      CheckOutput<float>("SavedVarianceZ",
+                         cpu_saved_var_z_,
+                         cpu_saved_var_base_z_,
+                         diff,
+                         is_relative_atol);
+    }
+    CheckOutput<T>("Y", cpu_y, cpu_y_base_, diff, is_relative_atol);
+  }
+
+  void CheckBackward(float diff, bool is_relative_atol = false) {
+    phi::GPUContext *ctx = static_cast<phi::GPUContext *>(
+        phi::DeviceContextPool::Instance().Get(phi::GPUPlace(0)));
+
+    phi::DenseTensor cpu_dx_base;
+    phi::DenseTensor cpu_dz_base;
+    phi::DenseTensor cpu_dscale_base;
+    phi::DenseTensor cpu_dbias_base;
+    BaselineBackwardFusedBNAddRelu(
+        *ctx, &cpu_dx_base, &cpu_dz_base, &cpu_dscale_base, &cpu_dbias_base);
+
+    phi::DenseTensor cpu_dx;
+    phi::DenseTensor cpu_dz;
+    phi::DenseTensor cpu_dscale;
+    phi::DenseTensor cpu_dbias;
+    FusedBackward(*ctx, &cpu_dx, &cpu_dz, &cpu_dscale, &cpu_dbias);
+
+    CheckOutput<T>("DX", cpu_dx, cpu_dx_base, diff, is_relative_atol);
+    CheckOutput<T>("DZ", cpu_dz, cpu_dz_base, diff, is_relative_atol);
+    CheckOutput<float>(
+        "DScale", cpu_dscale, cpu_dscale_base, diff, is_relative_atol);
+    CheckOutput<float>(
+        "DBias", cpu_dbias, cpu_dbias_base, diff, is_relative_atol);
+  }
+
+ private:
+  void SetUp() {
+    InitRandomTensor<T>({batch_size_, height_, width_, channels_}, &cpu_x_);
+    InitRandomTensor<float>({channels_}, &cpu_bn_scale_x_);
+    InitRandomTensor<float>({channels_}, &cpu_bn_bias_x_);
+
+    if (has_shortcut_) {
+      InitRandomTensor<T>({batch_size_, height_, width_, channels_}, &cpu_z_);
+      InitRandomTensor<float>({channels_}, &cpu_bn_scale_z_);
+      InitRandomTensor<float>({channels_}, &cpu_bn_bias_z_);
+    } else {
+      if (fuse_add_) {
+        InitRandomTensor<T>({batch_size_, height_, width_, channels_}, &cpu_z_);
+      }
+    }
+
+    InitRandomTensor<T>({batch_size_, height_, width_, channels_}, &cpu_dy_);
+  }
+
+  void InitMeanVar(phi::DenseTensor *cpu_mean,
+                   phi::DenseTensor *cpu_var,
+                   phi::DenseTensor *cpu_saved_mean,
+                   phi::DenseTensor *cpu_saved_var) {
+    InitConstantTensor<float>({channels_}, static_cast<float>(0.0f), cpu_mean);
+    InitConstantTensor<float>({channels_}, static_cast<float>(1.0f), cpu_var);
+    InitConstantTensor<float>(
+        {channels_}, static_cast<float>(0.0f), cpu_saved_mean);
+    InitConstantTensor<float>(
+        {channels_}, static_cast<float>(0.0f), cpu_saved_var);
+  }
+
+  void BaselineForward(const phi::GPUContext &ctx,
+                       phi::DenseTensor *cpu_mean_x,
+                       phi::DenseTensor *cpu_var_x,
+                       phi::DenseTensor *cpu_saved_mean_x,
+                       phi::DenseTensor *cpu_saved_var_x,
+                       phi::DenseTensor *cpu_y,
+                       phi::DenseTensor *saved_reserve_space_x,
+                       phi::DenseTensor *cpu_mean_z = nullptr,
+                       phi::DenseTensor *cpu_var_z = nullptr,
+                       phi::DenseTensor *cpu_saved_mean_z = nullptr,
+                       phi::DenseTensor *cpu_saved_var_z = nullptr,
+                       phi::DenseTensor *saved_reserve_space_z = nullptr) {
+    InitMeanVar(cpu_mean_x, cpu_var_x, cpu_saved_mean_x, cpu_saved_var_x);
+    ComputeBatchNormForward(ctx,
+                            cpu_x_,
+                            cpu_bn_scale_x_,
+                            cpu_bn_bias_x_,
+                            cpu_mean_x,
+                            cpu_var_x,
+                            cpu_saved_mean_x,
+                            cpu_saved_var_x,
+                            cpu_y,
+                            saved_reserve_space_x);
+    if (has_shortcut_) {
+      phi::DenseTensor cpu_z_out;
+      InitMeanVar(cpu_mean_z, cpu_var_z, cpu_saved_mean_z, cpu_saved_var_z);
+      ComputeBatchNormForward(ctx,
+                              cpu_z_,
+                              cpu_bn_scale_z_,
+                              cpu_bn_bias_z_,
+                              cpu_mean_z,
+                              cpu_var_z,
+                              cpu_saved_mean_z,
+                              cpu_saved_var_z,
+                              &cpu_z_out,
+                              saved_reserve_space_z);
+      ComputeInplaceAdd<T>(cpu_z_out, cpu_y);
+    } else {
+      if (fuse_add_) {
+        ComputeInplaceAdd<T>(cpu_z_, cpu_y);
+      }
+    }
+    if (act_type_ == "relu") {
+      ComputeInplaceRelu<T>(cpu_y);
+    }
+  }
+
+  void BaselineForwardFusedBNAddRelu(const phi::GPUContext &ctx,
+                                     phi::DenseTensor *cpu_mean,
+                                     phi::DenseTensor *cpu_var,
+                                     phi::DenseTensor *cpu_saved_mean,
+                                     phi::DenseTensor *cpu_saved_var,
+                                     phi::DenseTensor *cpu_y,
+                                     phi::DenseTensor *saved_reserve_space) {
+    InitMeanVar(cpu_mean, cpu_var, cpu_saved_mean, cpu_saved_var);
+    ComputeFusedBNAddReluForward(ctx,
+                                 cpu_x_,
+                                 cpu_z_,
+                                 cpu_bn_scale_x_,
+                                 cpu_bn_bias_x_,
+                                 cpu_mean,
+                                 cpu_var,
+                                 cpu_saved_mean,
+                                 cpu_saved_var,
+                                 cpu_y,
+                                 saved_reserve_space);
+  }
+
+  void BaselineBackwardFusedBNAddRelu(const phi::GPUContext &ctx,
+                                      phi::DenseTensor *cpu_dx,
+                                      phi::DenseTensor *cpu_dz,
+                                      phi::DenseTensor *cpu_dscale,
+                                      phi::DenseTensor *cpu_dbias) {
+    ComputeFusedBNAddReluBackward(ctx,
+                                  cpu_dy_,
+                                  cpu_x_,
+                                  cpu_bn_scale_x_,
+                                  cpu_bn_bias_x_,
+                                  cpu_saved_mean_base_x_,
+                                  cpu_saved_var_base_x_,
+                                  cpu_y_base_,
+                                  saved_reserve_space_x_,
+                                  cpu_dx,
+                                  cpu_dz,
+                                  cpu_dscale,
+                                  cpu_dbias);
+  }
+
+  void ComputeFusedBNStatsFinalize(const phi::GPUContext &ctx,
+                                   const phi::DenseTensor &cpu_x,
+                                   const phi::DenseTensor &cpu_bn_scale,
+                                   const phi::DenseTensor &cpu_bn_bias,
+                                   phi::DenseTensor *sum,
+                                   phi::DenseTensor *sum_of_square,
+                                   phi::DenseTensor *bn_scale,
+                                   phi::DenseTensor *bn_bias,
+                                   phi::DenseTensor *mean,
+                                   phi::DenseTensor *var,
+                                   phi::DenseTensor *saved_mean,
+                                   phi::DenseTensor *saved_var,
+                                   phi::DenseTensor *equiv_scale,
+                                   phi::DenseTensor *equiv_bias) {
+    phi::DenseTensor cpu_sum;
+    phi::DenseTensor cpu_sum_of_square;
+    ComputeSumAndSquareSum<T>(cpu_x, &cpu_sum, &cpu_sum_of_square);
+
+    auto place = ctx.GetPlace();
+    paddle::framework::TensorCopySync(cpu_sum, place, sum);
+    paddle::framework::TensorCopySync(cpu_sum_of_square, place, sum_of_square);
+    paddle::framework::TensorCopySync(cpu_bn_scale, place, bn_scale);
+    paddle::framework::TensorCopySync(cpu_bn_bias, place, bn_bias);
+
+    bn_scale->Resize({1, 1, 1, channels_});
+    bn_bias->Resize({1, 1, 1, channels_});
+
+    // input
+    mean->Resize({1, 1, 1, channels_});
+    var->Resize({1, 1, 1, channels_});
+
+    // output
+    equiv_scale->Resize({1, 1, 1, channels_});
+    equiv_bias->Resize({1, 1, 1, channels_});
+    saved_mean->Resize({1, 1, 1, channels_});
+    saved_var->Resize({1, 1, 1, channels_});
+
+    auto param_shape = common::vectorize<int>(bn_scale->dims());
+    phi::fusion::CudnnBNStatsFinalize<T> bn_op(ctx, param_shape);
+    bn_op.Forward(ctx,
+                  *sum,
+                  *sum_of_square,
+                  *bn_scale,
+                  *bn_bias,
+                  saved_mean,
+                  saved_var,
+                  mean,
+                  var,
+                  equiv_scale,
+                  equiv_bias,
+                  eps_,
+                  momentum_,
+                  ele_count_,
+                  true);
+  }
+
+  // Get forward results of CudnnBNStatsFinalize + CudnnScaleBiasAddRelu
+  void FusedForward(const phi::GPUContext &ctx,
+                    phi::DenseTensor *cpu_mean_x,
+                    phi::DenseTensor *cpu_var_x,
+                    phi::DenseTensor *cpu_saved_mean_x,
+                    phi::DenseTensor *cpu_saved_var_x,
+                    phi::DenseTensor *cpu_y,
+                    phi::DenseTensor *cpu_bitmask,
+                    phi::DenseTensor *cpu_mean_z = nullptr,
+                    phi::DenseTensor *cpu_var_z = nullptr,
+                    phi::DenseTensor *cpu_saved_mean_z = nullptr,
+                    phi::DenseTensor *cpu_saved_var_z = nullptr) {
+    phi::DenseTensor x;
+    phi::DenseTensor sum_x;
+    phi::DenseTensor sum_of_square_x;
+    phi::DenseTensor bn_scale_x;
+    phi::DenseTensor bn_bias_x;
+
+    phi::DenseTensor z;
+    phi::DenseTensor sum_z;
+    phi::DenseTensor sum_of_square_z;
+    phi::DenseTensor bn_scale_z;
+    phi::DenseTensor bn_bias_z;
+
+    auto place = ctx.GetPlace();
+    paddle::framework::TensorCopySync(cpu_x_, place, &x);
+    if (fuse_add_ || has_shortcut_) {
+      paddle::framework::TensorCopySync(cpu_z_, place, &z);
+    }
+
+    phi::DenseTensor mean_x;
+    phi::DenseTensor var_x;
+    phi::DenseTensor saved_mean_x;
+    phi::DenseTensor saved_var_x;
+    phi::DenseTensor equiv_scale_x;
+    phi::DenseTensor equiv_bias_x;
+
+    phi::DenseTensor mean_z;
+    phi::DenseTensor var_z;
+    phi::DenseTensor saved_mean_z;
+    phi::DenseTensor saved_var_z;
+    phi::DenseTensor equiv_scale_z;
+    phi::DenseTensor equiv_bias_z;
+
+    phi::DenseTensor y;
+    phi::DenseTensor bitmask;
+
+    InitMeanVar(cpu_mean_x, cpu_var_x, cpu_saved_mean_x, cpu_saved_var_x);
+    paddle::framework::TensorCopySync(*cpu_mean_x, place, &mean_x);
+    paddle::framework::TensorCopySync(*cpu_var_x, place, &var_x);
+    if (has_shortcut_) {
+      InitMeanVar(cpu_mean_z, cpu_var_z, cpu_saved_mean_z, cpu_saved_var_z);
+      paddle::framework::TensorCopySync(*cpu_mean_z, place, &mean_z);
+      paddle::framework::TensorCopySync(*cpu_var_z, place, &var_z);
+    }
+
+    // 1. BN Stats Finalize
+    ComputeFusedBNStatsFinalize(ctx,
+                                cpu_x_,
+                                cpu_bn_scale_x_,
+                                cpu_bn_bias_x_,
+                                &sum_x,
+                                &sum_of_square_x,
+                                &bn_scale_x,
+                                &bn_bias_x,
+                                &mean_x,
+                                &var_x,
+                                &saved_mean_x,
+                                &saved_var_x,
+                                &equiv_scale_x,
+                                &equiv_bias_x);
+    if (has_shortcut_) {
+      ComputeFusedBNStatsFinalize(ctx,
+                                  cpu_z_,
+                                  cpu_bn_scale_z_,
+                                  cpu_bn_bias_z_,
+                                  &sum_z,
+                                  &sum_of_square_z,
+                                  &bn_scale_z,
+                                  &bn_bias_z,
+                                  &mean_z,
+                                  &var_z,
+                                  &saved_mean_z,
+                                  &saved_var_z,
+                                  &equiv_scale_z,
+                                  &equiv_bias_z);
+    }
+
+    y.Resize(common::make_ddim({batch_size_, height_, width_, channels_}));
+
+    int c = channels_;
+    int64_t nhw = ele_count_;
+    int32_t c_int32_elems = ((c + 63) & ~63) / 32;
+    int32_t nhw_int32_elems = (static_cast<int32_t>(nhw) + 31) & ~31;
+    bitmask.Resize(common::make_ddim({nhw_int32_elems, c_int32_elems, 1}));
+
+    auto data_shape = common::vectorize<int>(x.dims());
+    auto param_shape = common::vectorize<int>(bn_scale_x.dims());
+    auto bitmask_shape = common::vectorize<int>(bitmask.dims());
+
+    // 2. Scale Bias + Relu
+    phi::fusion::CudnnScaleBiasAddRelu<T> sbar_op(ctx,
+                                                  act_type_,
+                                                  fuse_add_,
+                                                  has_shortcut_,
+                                                  data_shape,
+                                                  param_shape,
+                                                  bitmask_shape);
+    sbar_op.Forward(ctx,
+                    x,
+                    equiv_scale_x,
+                    equiv_bias_x,
+                    &z,
+                    &equiv_scale_z,
+                    &equiv_bias_z,
+                    &y,
+                    &bitmask);
+
+    paddle::framework::TensorCopySync(mean_x, phi::CPUPlace(), cpu_mean_x);
+    paddle::framework::TensorCopySync(var_x, phi::CPUPlace(), cpu_var_x);
+    paddle::framework::TensorCopySync(
+        saved_mean_x, phi::CPUPlace(), cpu_saved_mean_x);
+    paddle::framework::TensorCopySync(
+        saved_var_x, phi::CPUPlace(), cpu_saved_var_x);
+    if (has_shortcut_) {
+      paddle::framework::TensorCopySync(mean_z, phi::CPUPlace(), cpu_mean_z);
+      paddle::framework::TensorCopySync(var_z, phi::CPUPlace(), cpu_var_z);
+      paddle::framework::TensorCopySync(
+          saved_mean_z, phi::CPUPlace(), cpu_saved_mean_z);
+      paddle::framework::TensorCopySync(
+          saved_var_z, phi::CPUPlace(), cpu_saved_var_z);
+    }
+    paddle::framework::TensorCopySync(y, phi::CPUPlace(), cpu_y);
+    paddle::framework::TensorCopySync(bitmask, phi::CPUPlace(), cpu_bitmask);
+  }
+
+  // Get backward results of CudnnBNStatsFinalize + CudnnScaleBiasAddRelu
+  void FusedBackward(const phi::GPUContext &ctx,
+                     phi::DenseTensor *cpu_dx,
+                     phi::DenseTensor *cpu_dz,
+                     phi::DenseTensor *cpu_dscale,
+                     phi::DenseTensor *cpu_dbias) {
+    phi::DenseTensor dy;
+    phi::DenseTensor x;
+    phi::DenseTensor bn_scale;
+    phi::DenseTensor bn_bias;
+    phi::DenseTensor saved_mean;
+    phi::DenseTensor saved_var;
+    phi::DenseTensor bitmask;
+    phi::DenseTensor dx;
+    phi::DenseTensor dz;
+    phi::DenseTensor dscale;
+    phi::DenseTensor dbias;
+
+    auto place = ctx.GetPlace();
+    paddle::framework::TensorCopySync(cpu_dy_, place, &dy);
+    paddle::framework::TensorCopySync(cpu_x_, place, &x);
+    paddle::framework::TensorCopySync(cpu_bn_scale_x_, place, &bn_scale);
+    paddle::framework::TensorCopySync(cpu_bn_bias_x_, place, &bn_bias);
+    paddle::framework::TensorCopySync(cpu_saved_mean_x_, place, &saved_mean);
+    paddle::framework::TensorCopySync(cpu_saved_var_x_, place, &saved_var);
+    paddle::framework::TensorCopySync(cpu_bitmask_, place, &bitmask);
+
+    bn_scale.Resize({1, 1, 1, channels_});
+    bn_bias.Resize({1, 1, 1, channels_});
+    saved_mean.Resize({1, 1, 1, channels_});
+    saved_var.Resize({1, 1, 1, channels_});
+
+    dx.Resize(common::make_ddim({batch_size_, height_, width_, channels_}));
+    dz.Resize(common::make_ddim({batch_size_, height_, width_, channels_}));
+    dscale.Resize(common::make_ddim({1, 1, 1, channels_}));
+    dbias.Resize(common::make_ddim({1, 1, 1, channels_}));
+
+    auto data_shape = common::vectorize<int>(x.dims());
+    auto param_shape = common::vectorize<int>(bn_scale.dims());
+    auto bitmask_shape = common::vectorize<int>(bitmask.dims());
+
+    std::string act_type = "relu";
+    phi::fusion::CudnnScaleBiasAddRelu<T> sbar_op(
+        ctx, act_type, true, false, data_shape, param_shape, bitmask_shape);
+    sbar_op.Backward(ctx,
+                     dy,
+                     x,
+                     bn_scale,
+                     bn_bias,
+                     saved_mean,
+                     saved_var,
+                     &bitmask,
+                     &dx,
+                     &dz,
+                     &dscale,
+                     &dbias,
+                     eps_);
+
+    paddle::framework::TensorCopySync(dx, phi::CPUPlace(), cpu_dx);
+    paddle::framework::TensorCopySync(dz, phi::CPUPlace(), cpu_dz);
+    paddle::framework::TensorCopySync(dscale, phi::CPUPlace(), cpu_dscale);
+    paddle::framework::TensorCopySync(dbias, phi::CPUPlace(), cpu_dbias);
+  }
+
+ private:
+  int batch_size_;
+  int height_;
+  int width_;
+  int channels_;
+  int ele_count_;
+
+  std::string act_type_;
+  bool fuse_add_;
+  bool has_shortcut_;
+
+  // Forward input
+  phi::DenseTensor cpu_x_;
+  phi::DenseTensor cpu_bn_scale_x_;
+  phi::DenseTensor cpu_bn_bias_x_;
+  phi::DenseTensor cpu_z_;
+  phi::DenseTensor cpu_bn_scale_z_;
+  phi::DenseTensor cpu_bn_bias_z_;
+
+  // Backward input
+  phi::DenseTensor cpu_dy_;
+  phi::DenseTensor cpu_bitmask_;
+  phi::DenseTensor cpu_saved_mean_x_;
+  phi::DenseTensor cpu_saved_var_x_;
+  phi::DenseTensor cpu_saved_mean_z_;
+  phi::DenseTensor cpu_saved_var_z_;
+  phi::DenseTensor cpu_saved_mean_base_x_;
+  phi::DenseTensor cpu_saved_var_base_x_;
+  phi::DenseTensor saved_reserve_space_x_;
+  phi::DenseTensor cpu_saved_mean_base_z_;
+  phi::DenseTensor cpu_saved_var_base_z_;
+  phi::DenseTensor saved_reserve_space_z_;
+  phi::DenseTensor cpu_y_base_;
+
+  double eps_ = 1e-5;
+  float momentum_ = 0.9;
+};
+
+TEST(CudnnBNAddReluFp16, BNAdd) {
+  int batch_size = 4;
+  int height = 8;
+  int width = 8;
+  int channels = 64;
+  std::string act_type = "";
+  bool has_shortcut = false;
+  FLAGS_cudnn_batchnorm_spatial_persistent = true;
+  for (auto fuse_add : {false, true}) {
+    CudnnBNAddReluTester<phi::dtype::float16> test(
+        batch_size, height, width, channels, act_type, fuse_add, has_shortcut);
+    test.CheckForward(2e-3);
+  }
+}
+
+TEST(CudnnBNAddReluFp16, BNAddRelu) {
+  int batch_size = 4;
+  int height = 8;
+  int width = 8;
+  int channels = 64;
+  std::string act_type = "relu";
+  bool has_shortcut = false;
+  FLAGS_cudnn_batchnorm_spatial_persistent = true;
+  for (auto fuse_add : {false, true}) {
+    CudnnBNAddReluTester<phi::dtype::float16> test(
+        batch_size, height, width, channels, act_type, fuse_add, has_shortcut);
+    test.CheckForward(2e-3);
+    if (fuse_add) {
+      test.CheckBackward(2e-4);
+    }
+  }
+}
+
+TEST(CudnnBNAddReluFp16, HasShortcut) {
+  int batch_size = 4;
+  int height = 8;
+  int width = 8;
+  int channels = 64;
+  std::string act_type = "";
+  bool fuse_add = false;
+  bool has_shortcut = true;
+  FLAGS_cudnn_batchnorm_spatial_persistent = true;
+  CudnnBNAddReluTester<phi::dtype::float16> test(
+      batch_size, height, width, channels, act_type, fuse_add, has_shortcut);
+  test.CheckForward(5e-3);
+}

test/cpp/fluid/fused/cudnn_norm_conv_test.cc

@@ -0,0 +1,538 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+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. */
+
+#include <random>
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+#include "paddle/fluid/framework/program_desc.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/phi/common/float16.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/fusion/gpu/cudnn_norm_conv.cu.h"
+
+namespace framework = paddle::framework;
+namespace platform = paddle::platform;
+
+USE_OP_ITSELF(conv2d);
+USE_OP_ITSELF(conv2d_grad);
+PD_DECLARE_KERNEL(conv2d, GPUDNN, ALL_LAYOUT);
+PD_DECLARE_KERNEL(conv2d_grad, GPUDNN, ALL_LAYOUT);
+
+template <typename T>
+void InitRandomTensor(const std::vector<int64_t> &dims,
+                      phi::DenseTensor *cpu_out) {
+  T *cpu_out_ptr =
+      cpu_out->mutable_data<T>(common::make_ddim(dims), phi::CPUPlace());
+
+  std::default_random_engine random(0);
+  std::uniform_real_distribution<float> dis(0.0, 1.0);
+  for (int i = 0; i < cpu_out->numel(); ++i) {
+    cpu_out_ptr[i] = static_cast<T>(dis(random));
+  }
+}
+
+template <typename T>
+void TransposeNchwToNhwc(const phi::DenseTensor &cpu_in,
+                         phi::DenseTensor *cpu_out) {
+  const auto &in_dims = cpu_in.dims();
+  EXPECT_EQ(cpu_in.dims().size(), 4);
+
+  const T *cpu_in_ptr = cpu_in.data<T>();
+  T *cpu_out_ptr = cpu_out->mutable_data<T>(
+      {in_dims[0], in_dims[2], in_dims[3], in_dims[1]}, phi::CPUPlace());
+
+  int64_t n = in_dims[0];
+  int64_t c = in_dims[1];
+  int64_t hw = in_dims[2] * in_dims[3];
+  for (int i = 0; i < n; ++i) {
+    for (int j = 0; j < hw; ++j) {
+      for (int k = 0; k < c; ++k) {
+        int dst_idx = i * hw * c + j * c + k;   // NOLINT
+        int src_idx = i * c * hw + k * hw + j;  // NOLINT
+        cpu_out_ptr[dst_idx] = cpu_in_ptr[src_idx];
+      }
+    }
+  }
+}
+
+template <typename T>
+void CheckOutput(const phi::DenseTensor &cpu_res,
+                 const phi::DenseTensor &cpu_base,
+                 float diff,
+                 bool is_relative_atol = false) {
+  EXPECT_EQ(cpu_res.dims(), cpu_base.dims());
+
+  const T *cpu_res_ptr = cpu_res.data<T>();
+  const T *cpu_base_ptr = cpu_base.data<T>();
+  for (int i = 0; i < cpu_res.numel(); ++i) {
+    if (is_relative_atol) {
+      EXPECT_LT(static_cast<float>(std::abs((cpu_res_ptr[i] - cpu_base_ptr[i]) /
+                                            cpu_base_ptr[i])),
+                diff);
+    } else {
+      EXPECT_LT(static_cast<float>(std::abs(cpu_res_ptr[i] - cpu_base_ptr[i])),
+                diff);
+    }
+  }
+}
+
+// Use Paddle conv2d op results as baseline
+void ComputeConv2DForward(const phi::GPUContext &ctx,
+                          const phi::DenseTensor &cpu_input,
+                          const phi::DenseTensor &cpu_filter,
+                          phi::DenseTensor *cpu_output,
+                          int stride,
+                          int padding) {
+  framework::Scope scope;
+  auto *input = scope.Var("Input")->GetMutable<phi::DenseTensor>();
+  auto *filter = scope.Var("Filter")->GetMutable<phi::DenseTensor>();
+  auto *output = scope.Var("Output")->GetMutable<phi::DenseTensor>();
+
+  auto place = ctx.GetPlace();
+  paddle::framework::TensorCopySync(cpu_input, place, input);
+  paddle::framework::TensorCopySync(cpu_filter, place, filter);
+
+  framework::AttributeMap attrs;
+  bool use_cudnn = true;
+  std::string data_format = "NHWC";
+  std::vector<int> strides = {stride, stride};
+  std::vector<int> paddings = {padding, padding};
+  attrs.insert({"strides", strides});
+  attrs.insert({"paddings", paddings});
+  attrs.insert({"use_cudnn", use_cudnn});
+  attrs.insert({"data_format", data_format});
+
+  auto op = framework::OpRegistry::CreateOp(
+      "conv2d",
+      {{"Input", {"Input"}}, {"Filter", {"Filter"}}},
+      {{"Output", {"Output"}}},
+      attrs);
+  op->Run(scope, ctx.GetPlace());
+
+  paddle::framework::TensorCopySync(*output, phi::CPUPlace(), cpu_output);
+}
+
+// Use Paddle conv2d_grad op results as baseline
+void ComputeConv2DBackward(const phi::GPUContext &ctx,
+                           const phi::DenseTensor &cpu_input,
+                           const phi::DenseTensor &cpu_filter,
+                           const phi::DenseTensor &cpu_output_grad,
+                           phi::DenseTensor *cpu_input_grad,
+                           phi::DenseTensor *cpu_filter_grad,
+                           int stride,
+                           int padding,
+                           int dilation) {
+  framework::Scope scope;
+  auto *input = scope.Var("Input")->GetMutable<phi::DenseTensor>();
+  auto *filter = scope.Var("Filter")->GetMutable<phi::DenseTensor>();
+  auto *output_grad = scope.Var("Output@GRAD")->GetMutable<phi::DenseTensor>();
+  auto *input_grad = scope.Var("Input@GRAD")->GetMutable<phi::DenseTensor>();
+  auto *filter_grad = scope.Var("Filter@GRAD")->GetMutable<phi::DenseTensor>();
+
+  auto place = ctx.GetPlace();
+  paddle::framework::TensorCopySync(cpu_input, place, input);
+  paddle::framework::TensorCopySync(cpu_filter, place, filter);
+  paddle::framework::TensorCopySync(cpu_output_grad, place, output_grad);
+
+  framework::AttributeMap attrs;
+  bool use_cudnn = true;
+  std::string data_format = "NHWC";
+  std::string padding_algorithm = "EXPLICIT";
+  std::vector<int> strides = {stride, stride};
+  std::vector<int> paddings = {padding, padding};
+  std::vector<int> dilations = {dilation, dilation};
+  int groups = 1;
+  bool exhaustive_search = false;
+  bool use_addto = false;
+  attrs.insert({"use_cudnn", use_cudnn});
+  attrs.insert({"data_format", data_format});
+  attrs.insert({"padding_algorithm", padding_algorithm});
+  attrs.insert({"strides", strides});
+  attrs.insert({"paddings", paddings});
+  attrs.insert({"dilations", dilations});
+  attrs.insert({"groups", groups});
+  attrs.insert({"exhaustive_search", exhaustive_search});
+  attrs.insert({"use_addto", use_addto});
+  attrs.insert({"workspace_size_MB", 512});
+
+  auto op = framework::OpRegistry::CreateOp(
+      "conv2d_grad",
+      {{"Input", {"Input"}},
+       {"Filter", {"Filter"}},
+       {"Output@GRAD", {"Output@GRAD"}}},
+      {{"Input@GRAD", {"Input@GRAD"}}, {"Filter@GRAD", {"Filter@GRAD"}}},
+      attrs);
+  op->Run(scope, ctx.GetPlace());
+
+  paddle::framework::TensorCopySync(
+      *input_grad, phi::CPUPlace(), cpu_input_grad);
+  paddle::framework::TensorCopySync(
+      *filter_grad, phi::CPUPlace(), cpu_filter_grad);
+}
+
+template <typename T>
+void ComputeSumAndSquareSum(const phi::DenseTensor &cpu_out,
+                            phi::DenseTensor *cpu_sum,
+                            phi::DenseTensor *cpu_sum_of_square) {
+  const auto &dims = cpu_out.dims();
+  int64_t c = dims[3];
+
+  const T *cpu_out_ptr = cpu_out.data<T>();
+  float *cpu_sum_ptr =
+      cpu_sum->mutable_data<float>({1, 1, 1, c}, phi::CPUPlace());
+  float *cpu_sum_square_ptr =
+      cpu_sum_of_square->mutable_data<float>({1, 1, 1, c}, phi::CPUPlace());
+
+  for (int j = 0; j < c; ++j) {
+    float tmp_sum = 0.0f;
+    float tmp_sum_of_squares = 0.0f;
+    for (int i = 0; i < cpu_out.numel() / c; ++i) {
+      float tmp_out = static_cast<float>(cpu_out_ptr[i * c + j]);
+      tmp_sum += tmp_out;
+      tmp_sum_of_squares += tmp_out * tmp_out;
+    }
+    cpu_sum_ptr[j] = tmp_sum;
+    cpu_sum_square_ptr[j] = tmp_sum_of_squares;
+  }
+}
+
+template <typename T>
+class CudnnNormConvolutionTester {
+ public:
+  CudnnNormConvolutionTester(int batch_size,
+                             int height,
+                             int width,
+                             int input_channels,
+                             int output_channels,
+                             int kernel_size,
+                             int stride) {
+    batch_size_ = batch_size;
+    height_ = height;
+    width_ = width;
+    input_channels_ = input_channels;
+    output_channels_ = output_channels;
+    kernel_size_ = kernel_size;
+    stride_ = stride;
+    padding_ = (kernel_size_ - 1) / 2;
+    out_height_ = (height_ + 2 * padding_ - kernel_size_) / stride_ + 1;
+    out_width_ = (width_ + 2 * padding_ - kernel_size_) / stride_ + 1;
+    SetUp();
+  }
+
+  ~CudnnNormConvolutionTester() = default;
+
+  void CheckForward(float diff, bool is_relative_atol = false) {
+    phi::GPUContext *ctx = static_cast<phi::GPUContext *>(
+        phi::DeviceContextPool::Instance().Get(phi::GPUPlace(0)));
+
+    phi::DenseTensor cpu_output_base;
+    phi::DenseTensor cpu_sum_base;
+    phi::DenseTensor cpu_sum_of_square_base;
+    BaselineForward(
+        *ctx, &cpu_output_base, &cpu_sum_base, &cpu_sum_of_square_base);
+
+    phi::DenseTensor cpu_output;
+    phi::DenseTensor cpu_sum;
+    phi::DenseTensor cpu_sum_of_square;
+    FusedForward(*ctx, &cpu_output, &cpu_sum, &cpu_sum_of_square);
+
+    // Check forward correctness between baseline and results of normconv.
+    CheckOutput<T>(cpu_output, cpu_output_base, diff, is_relative_atol);
+    CheckOutput<float>(cpu_sum, cpu_sum_base, diff, is_relative_atol);
+    CheckOutput<float>(
+        cpu_sum_of_square, cpu_sum_of_square_base, diff, is_relative_atol);
+  }
+
+  void CheckBackward(float diff, bool is_relative_atol = false) {
+    phi::GPUContext *ctx = static_cast<phi::GPUContext *>(
+        phi::DeviceContextPool::Instance().Get(phi::GPUPlace(0)));
+
+    phi::DenseTensor cpu_input_grad_base;
+    phi::DenseTensor cpu_filter_nchw_grad_base;
+    phi::DenseTensor cpu_filter_nhwc_grad_base;
+    BaselineBackward(*ctx, &cpu_input_grad_base, &cpu_filter_nchw_grad_base);
+    TransposeNchwToNhwc<T>(cpu_filter_nchw_grad_base,
+                           &cpu_filter_nhwc_grad_base);
+
+    phi::DenseTensor cpu_input_grad;
+    phi::DenseTensor cpu_filter_nhwc_grad;
+    FusedBackward(*ctx, &cpu_input_grad, &cpu_filter_nhwc_grad);
+
+    // Check backward correctness between baseline and results of normconv.
+    CheckOutput<T>(cpu_input_grad, cpu_input_grad_base, diff, is_relative_atol);
+    CheckOutput<T>(cpu_filter_nhwc_grad,
+                   cpu_filter_nhwc_grad_base,
+                   diff,
+                   is_relative_atol);
+  }
+
+ private:
+  void SetUp() {
+    InitRandomTensor<T>({batch_size_, height_, width_, input_channels_},
+                        &cpu_input_);
+    InitRandomTensor<T>(
+        {output_channels_, input_channels_, kernel_size_, kernel_size_},
+        &cpu_filter_nchw_);
+    // transpoes for filter, NCHW -> NHWC
+    TransposeNchwToNhwc<T>(cpu_filter_nchw_, &cpu_filter_nhwc_);
+    InitRandomTensor<T>(
+        {batch_size_, out_height_, out_width_, output_channels_},
+        &cpu_output_grad_);
+  }
+
+  void BaselineForward(const phi::GPUContext &ctx,
+                       phi::DenseTensor *cpu_output_base,
+                       phi::DenseTensor *cpu_sum_base,
+                       phi::DenseTensor *cpu_sum_of_square_base) {
+    ComputeConv2DForward(
+        ctx, cpu_input_, cpu_filter_nchw_, cpu_output_base, stride_, padding_);
+    ComputeSumAndSquareSum<T>(
+        *cpu_output_base, cpu_sum_base, cpu_sum_of_square_base);
+  }
+
+  void BaselineBackward(const phi::GPUContext &ctx,
+                        phi::DenseTensor *cpu_input_grad_base,
+                        phi::DenseTensor *cpu_filter_grad_base) {
+    ComputeConv2DBackward(ctx,
+                          cpu_input_,
+                          cpu_filter_nchw_,
+                          cpu_output_grad_,
+                          cpu_input_grad_base,
+                          cpu_filter_grad_base,
+                          stride_,
+                          padding_,
+                          dilation_);
+  }
+
+  // get forward results of cudnn_norm_conv
+  void FusedForward(const phi::GPUContext &ctx,
+                    phi::DenseTensor *cpu_output,
+                    phi::DenseTensor *cpu_sum,
+                    phi::DenseTensor *cpu_sum_of_square) {
+    phi::DenseTensor input;
+    phi::DenseTensor filter_nhwc;
+    phi::DenseTensor output;
+    phi::DenseTensor sum;
+    phi::DenseTensor sum_of_square;
+
+    auto place = ctx.GetPlace();
+    paddle::framework::TensorCopySync(cpu_input_, place, &input);
+    paddle::framework::TensorCopySync(cpu_filter_nhwc_, place, &filter_nhwc);
+
+    output.Resize(common::make_ddim(
+        {batch_size_, out_height_, out_width_, output_channels_}));
+    sum.Resize(common::make_ddim({1, 1, 1, output_channels_}));
+    sum_of_square.Resize(common::make_ddim({1, 1, 1, output_channels_}));
+
+    auto input_shape = common::vectorize<int>(input.dims());
+    auto filter_shape = common::vectorize<int>(filter_nhwc.dims());
+    auto output_shape = common::vectorize<int>(output.dims());
+    phi::fusion::CudnnNormConvolution<T> conv_op(ctx,
+                                                 input_shape,
+                                                 filter_shape,
+                                                 output_shape,
+                                                 padding_,
+                                                 stride_,
+                                                 dilation_,
+                                                 group_);
+    conv_op.Forward(ctx, input, filter_nhwc, &output, &sum, &sum_of_square);
+
+    paddle::framework::TensorCopySync(output, phi::CPUPlace(), cpu_output);
+    paddle::framework::TensorCopySync(sum, phi::CPUPlace(), cpu_sum);
+    paddle::framework::TensorCopySync(
+        sum_of_square, phi::CPUPlace(), cpu_sum_of_square);
+  }
+
+  void FusedBackward(const phi::GPUContext &ctx,
+                     phi::DenseTensor *cpu_input_grad,
+                     phi::DenseTensor *cpu_filter_grad) {
+    phi::DenseTensor input;
+    phi::DenseTensor filter_nhwc;
+    phi::DenseTensor output_grad;
+    phi::DenseTensor input_grad;
+    phi::DenseTensor filter_grad;
+
+    auto place = ctx.GetPlace();
+    paddle::framework::TensorCopySync(cpu_input_, place, &input);
+    paddle::framework::TensorCopySync(cpu_filter_nhwc_, place, &filter_nhwc);
+    paddle::framework::TensorCopySync(cpu_output_grad_, place, &output_grad);
+
+    input_grad.Resize(input.dims());
+    filter_grad.Resize(filter_nhwc.dims());
+
+    auto input_shape = common::vectorize<int>(input.dims());
+    auto filter_shape = common::vectorize<int>(filter_nhwc.dims());
+    auto output_shape = common::vectorize<int>(output_grad.dims());
+    phi::fusion::CudnnNormConvolutionGrad<T> conv_grad_op(ctx,
+                                                          input_shape,
+                                                          filter_shape,
+                                                          output_shape,
+                                                          padding_,
+                                                          stride_,
+                                                          dilation_,
+                                                          group_);
+    conv_grad_op.Backward(
+        ctx, input, filter_nhwc, output_grad, &input_grad, &filter_grad);
+
+    paddle::framework::TensorCopySync(
+        input_grad, phi::CPUPlace(), cpu_input_grad);
+    paddle::framework::TensorCopySync(
+        filter_grad, phi::CPUPlace(), cpu_filter_grad);
+  }
+
+ private:
+  int batch_size_;
+  int height_;
+  int width_;
+  int out_height_;
+  int out_width_;
+  int input_channels_;
+  int output_channels_;
+  int kernel_size_;
+  int stride_;
+  int padding_;
+  const int dilation_ = 1;
+  const int group_ = 1;
+
+  // Forward input
+  phi::DenseTensor cpu_input_;
+  phi::DenseTensor cpu_filter_nchw_;
+  phi::DenseTensor cpu_filter_nhwc_;
+
+  // Backward input
+  phi::DenseTensor cpu_output_grad_;
+};
+
+// test for fp16, kernel = 1, output_channels = input_channels
+TEST(CudnnNormConvFp16, K1S1) {
+  int batch_size = 4;
+  int height = 56;
+  int width = 56;
+  int input_channels = 32;
+  int output_channels = 32;
+  int kernel_size = 1;
+  int stride = 1;
+  CudnnNormConvolutionTester<phi::dtype::float16> test(batch_size,
+                                                       height,
+                                                       width,
+                                                       input_channels,
+                                                       output_channels,
+                                                       kernel_size,
+                                                       stride);
+  phi::GPUContext *ctx = static_cast<phi::GPUContext *>(
+      phi::DeviceContextPool::Instance().Get(phi::GPUPlace(0)));
+
+  if (ctx->GetComputeCapability() < 70 || ctx->GetComputeCapability() >= 90) {
+    ASSERT_THROW(test.CheckForward(1e-3, true),
+                 paddle::platform::EnforceNotMet);
+    ASSERT_THROW(test.CheckBackward(1e-3, true),
+                 paddle::platform::EnforceNotMet);
+  } else {
+    ASSERT_NO_THROW(test.CheckForward(1e-3, true));
+    ASSERT_NO_THROW(test.CheckBackward(1e-3, true));
+  }
+}
+
+// test for fp16, kernel = 3, output_channels = input_channels
+TEST(CudnnNormConvFp16, K3S1) {
+  int batch_size = 4;
+  int height = 56;
+  int width = 56;
+  int input_channels = 32;
+  int output_channels = 32;
+  int kernel_size = 3;
+  int stride = 1;
+  CudnnNormConvolutionTester<phi::dtype::float16> test(batch_size,
+                                                       height,
+                                                       width,
+                                                       input_channels,
+                                                       output_channels,
+                                                       kernel_size,
+                                                       stride);
+  phi::GPUContext *ctx = static_cast<phi::GPUContext *>(
+      phi::DeviceContextPool::Instance().Get(phi::GPUPlace(0)));
+
+  if (ctx->GetComputeCapability() < 70 || ctx->GetComputeCapability() >= 90) {
+    ASSERT_THROW(test.CheckForward(1e-3, true),
+                 paddle::platform::EnforceNotMet);
+    ASSERT_THROW(test.CheckBackward(1e-3, true),
+                 paddle::platform::EnforceNotMet);
+  } else {
+    ASSERT_NO_THROW(test.CheckForward(1e-3, true));
+    ASSERT_NO_THROW(test.CheckBackward(1e-3, true));
+  }
+}
+
+// test for fp16, kernel = 1, output_channels = input_channels * 4
+TEST(CudnnNormConvFp16, K1S1O4) {
+  int batch_size = 4;
+  int height = 56;
+  int width = 56;
+  int input_channels = 32;
+  int output_channels = 128;
+  int kernel_size = 1;
+  int stride = 1;
+  CudnnNormConvolutionTester<phi::dtype::float16> test(batch_size,
+                                                       height,
+                                                       width,
+                                                       input_channels,
+                                                       output_channels,
+                                                       kernel_size,
+                                                       stride);
+  phi::GPUContext *ctx = static_cast<phi::GPUContext *>(
+      phi::DeviceContextPool::Instance().Get(phi::GPUPlace(0)));
+
+  if (ctx->GetComputeCapability() < 70 || ctx->GetComputeCapability() >= 90) {
+    ASSERT_THROW(test.CheckForward(1e-3, true),
+                 paddle::platform::EnforceNotMet);
+    ASSERT_THROW(test.CheckBackward(1e-3, true),
+                 paddle::platform::EnforceNotMet);
+  } else {
+    ASSERT_NO_THROW(test.CheckForward(1e-3, true));
+    ASSERT_NO_THROW(test.CheckBackward(1e-3, true));
+  }
+}
+
+// test for fp16, kernel = 1, stride = 2, output_channels = input_channels * 4
+TEST(CudnnNormConvFp16, K1S2O4) {
+  int batch_size = 4;
+  int height = 8;
+  int width = 8;
+  int input_channels = 32;
+  int output_channels = 128;
+  int kernel_size = 1;
+  int stride = 2;
+  CudnnNormConvolutionTester<phi::dtype::float16> test(batch_size,
+                                                       height,
+                                                       width,
+                                                       input_channels,
+                                                       output_channels,
+                                                       kernel_size,
+                                                       stride);
+  phi::GPUContext *ctx = static_cast<phi::GPUContext *>(
+      phi::DeviceContextPool::Instance().Get(phi::GPUPlace(0)));
+
+  if (ctx->GetComputeCapability() <= 70 || ctx->GetComputeCapability() >= 90) {
+    ASSERT_THROW(test.CheckForward(1e-3, true),
+                 paddle::platform::EnforceNotMet);
+    ASSERT_THROW(test.CheckBackward(1e-3), paddle::platform::EnforceNotMet);
+  } else {
+    ASSERT_NO_THROW(test.CheckForward(1e-3, true));
+    ASSERT_NO_THROW(test.CheckBackward(1e-3));
+  }
+}

test/cpp/fluid/fused/fused_dropout_act_bias_test.cu

@@ -0,0 +1,403 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+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. */
+
+#include <time.h>
+
+#include <random>
+#include <vector>
+
+#include "paddle/phi/common/amp_type_traits.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/functors.h"
+#include "paddle/phi/kernels/fusion/gpu/fused_dropout_act_bias.h"
+#include "test/cpp/fluid/fused/fused_dropout_test.h"
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_DECLARE_KERNEL(dropout, GPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(dropout_grad, GPU, ALL_LAYOUT);
+#endif
+
+namespace framework = paddle::framework;
+namespace platform = paddle::platform;
+
+/**
+ * @brief the unittest of fused_dropout_act_bias
+ * 1. random input data
+ * 2. add bias, call activation, call paddle dropout, and get the base result
+ * 3. call FusedDropoutActBias function get fused result
+ * 4. compare the base result and fused result
+ */
+
+template <typename T, typename Functor, typename GradFunctor>
+struct TestFusedDropoutActBias {
+  uint32_t rows;
+  uint32_t cols;
+  uint64_t seed;
+  float dropout_prob;
+  bool is_upscale_in_train;
+  bool is_test;  // default false,  Set to true for inference only
+  bool has_bias = true;
+  phi::DenseTensor src, bias, out, mask;
+  phi::DenseTensor dsrc, dbias;
+
+  std::vector<T> src_vec, bias_vec, out_vec, mask_vec;
+  std::vector<T> correct_out, correct_dsrc, correct_dbias;
+  std::vector<uint8_t> correct_mask;
+
+  phi::GPUPlace place;
+  phi::GPUContext *ctx;
+
+  TestFusedDropoutActBias() {
+    rows = 32;
+    cols = 32;
+    seed = 0;
+    dropout_prob = 0.0;
+    is_upscale_in_train = false;
+    is_test = false;
+    has_bias = true;
+    phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
+    auto devicectx = pool.Get(place);
+    ctx = reinterpret_cast<phi::GPUContext *>(devicectx);
+  }
+
+  TestFusedDropoutActBias(int rows_,
+                          int cols_,
+                          uint64_t seed_ = 0,
+                          float dropout_prob_ = 0.0,
+                          bool is_upscale_in_train_ = false,
+                          bool is_test_ = false) {
+    rows = rows_;
+    cols = cols_;
+    seed = seed_;
+    dropout_prob = dropout_prob_;
+    is_upscale_in_train = is_upscale_in_train_;
+    is_test = is_test_;
+    has_bias = true;
+    phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
+    auto devicectx = pool.Get(place);
+    ctx = reinterpret_cast<phi::GPUContext *>(devicectx);
+  }
+
+  ~TestFusedDropoutActBias() = default;
+
+  void SetUp() {
+    const int n = rows * cols;
+    correct_out.resize(n);
+    correct_mask.resize(n);
+    correct_dsrc.resize(n);
+    correct_dbias.resize(cols);
+
+    src_vec.resize(n);
+    bias_vec.resize(cols);
+    std::default_random_engine random(time(NULL));
+    std::uniform_real_distribution<float> dis(0.0, 1.0);
+
+    for (int i = 0; i < rows; i++) {
+      for (int j = 0; j < cols; j++) {
+        src_vec[i * cols + j] = static_cast<T>(dis(random));
+        if (i == 0) bias_vec[j] = dis(random);
+      }
+    }
+
+    framework::TensorFromVector<T>(src_vec, *ctx, &src);
+    src.Resize({rows, cols});
+    if (has_bias) {
+      framework::TensorFromVector<T>(bias_vec, *ctx, &bias);
+      bias.Resize({cols});
+    }
+
+    {
+      out.mutable_data<T>({rows, cols}, place);
+      mask.mutable_data<uint8_t>({rows, cols}, place);
+      dsrc.mutable_data<T>({rows, cols}, place);
+
+      if (has_bias) {
+        dbias.mutable_data<T>({cols}, place);
+      }
+    }
+  }
+
+  void BaseForward() {
+    std::vector<T> out1(rows * cols);
+    Functor act;
+    if (has_bias) {
+      // add bias and call activation
+      for (int i = 0; i < rows; i++) {
+        for (int j = 0; j < cols; j++) {
+          const T tmp = src_vec[i * cols + j] + bias_vec[j];
+          out1[i * cols + j] = act(tmp);
+        }
+      }
+      // call dropout
+      Dropout<T>(out1,
+                 src.dims(),
+                 &correct_out,
+                 &correct_mask,
+                 *ctx,
+                 seed,
+                 dropout_prob,
+                 is_upscale_in_train,
+                 is_test);
+    } else {
+      for (int i = 0; i < rows; i++) {
+        for (int j = 0; j < cols; j++) {
+          const T tmp = src_vec[i * cols + j];
+          out1[i * cols + j] = act(tmp);
+        }
+      }
+
+      Dropout<T>(out1,
+                 src.dims(),
+                 &correct_out,
+                 &correct_mask,
+                 *ctx,
+                 seed,
+                 dropout_prob,
+                 is_upscale_in_train,
+                 is_test);
+    }
+    ctx->Wait();
+  }
+
+  void BaseBackward() {
+    std::vector<T> _out(rows * cols);
+    // call dropout_grad
+    DropoutGrad<T>(&_out,
+                   src.dims(),
+                   correct_out,
+                   correct_mask,
+                   *ctx,
+                   dropout_prob,
+                   is_upscale_in_train);
+
+    // calculate dbias
+    memset(&correct_dbias[0], 0, cols * sizeof(T));
+    GradFunctor act_grad;
+    for (int i = 0; i < rows; i++) {
+      for (int j = 0; j < cols; j++) {
+        T args[2];
+        args[0] = _out[i * cols + j];
+        if (has_bias) {
+          args[1] = src_vec[i * cols + j] + bias_vec[j];
+        } else {
+          args[1] = src_vec[i * cols + j];
+        }
+        T val = args[0] * act_grad.UseOut(args[1]);
+        correct_dsrc[i * cols + j] = val;
+      }
+    }
+
+    if (has_bias) {
+      // reduce_sum: keep the same calculate order as the GPU
+      ReduceSum<T>(correct_dsrc, &correct_dbias, rows, cols);
+    }
+  }
+
+  void FusedForward() {
+    const int VecSize = MAX_CACHE_BYTES / sizeof(T);
+    auto config =
+        phi::fusion::Get1DBlocksAnd2DGrids(*ctx,
+                                           static_cast<uint64_t>(rows),
+                                           static_cast<uint64_t>(cols),
+                                           VecSize);
+    const int increment = ((cols - 1) / (config.thread_per_block.x *
+                                         config.block_per_grid.x * VecSize) +
+                           1) *
+                          VecSize;
+
+    T *bias_ptr = nullptr;
+    if (has_bias) {
+      bias_ptr = bias.data<T>();
+    }
+    Functor act;
+    phi::fusion::LaunchDropoutActBias<T, uint8_t, Functor>(act,
+                                                           seed,
+                                                           rows,
+                                                           cols,
+                                                           increment,
+                                                           dropout_prob,
+                                                           is_upscale_in_train,
+                                                           is_test,
+                                                           src.data<T>(),
+                                                           bias_ptr,
+                                                           out.data<T>(),
+                                                           mask.data<uint8_t>(),
+                                                           *ctx);
+    ctx->Wait();
+  }
+
+  void FusedBackward() {
+    if (is_test) return;
+
+    T *bias_ptr = nullptr;
+    T *dbias_ptr = nullptr;
+    if (has_bias) {
+      dbias_ptr = dbias.data<T>();
+      bias_ptr = bias.data<T>();
+    }
+    GradFunctor act_grad;
+    phi::fusion::LaunchDropoutActBiasGrad<T, uint8_t, GradFunctor>(
+        act_grad,
+        out.data<T>(),
+        mask.data<uint8_t>(),
+        src.data<T>(),
+        bias_ptr,
+        dropout_prob,
+        is_upscale_in_train,
+        rows,
+        cols,
+        dsrc.data<T>(),
+        dbias_ptr,
+        *ctx);
+  }
+
+  void Run() {
+    SetUp();
+    BaseForward();
+    FusedForward();
+    BaseBackward();
+    FusedBackward();
+  }
+
+  void CheckOut(const T diff) {
+    const int n = rows * cols;
+    std::vector<T> _out(n);
+    std::vector<uint8_t> _mask(n);
+    framework::TensorToVector(out, *ctx, &_out);
+    if (!is_test) {
+      framework::TensorToVector<uint8_t>(mask, *ctx, &_mask);
+    }
+    ctx->Wait();
+
+    for (int i = 0; i < n; i++) {
+      EXPECT_LT(std::abs(_out[i] - correct_out[i]), diff);
+      if (!is_test) EXPECT_EQ(_mask[i], correct_mask[i]);
+    }
+  }
+
+  void CheckGrad(const T diff) {
+    if (is_test) return;
+
+    const int n = rows * cols;
+
+    std::vector<T> _dsrc(n);
+    framework::TensorToVector(dsrc, *ctx, &_dsrc);
+
+    for (int i = 0; i < n; i++) {
+      EXPECT_LT(std::abs(_dsrc[i] - correct_dsrc[i]), diff);
+    }
+
+    if (has_bias) {
+      std::vector<T> _dbias(cols);
+      framework::TensorToVector(dbias, *ctx, &_dbias);
+      ctx->Wait();
+      for (int i = 0; i < cols; i++) {
+        EXPECT_LT(std::abs(_dbias[i] - correct_dbias[i]), diff);
+      }
+    }
+  }
+};
+
+// test the shape , bias, activation
+template <typename T, typename Functor, typename GradFunctor>
+static void BaseTest(const bool is_fp16 = false) {
+  const int rows = 16;
+  std::vector<int> cols_list = {16, 17};
+  bool has_bias[2] = {true, false};
+  T default_diff = !is_fp16 ? static_cast<T>(1e-5) : static_cast<T>(1e-1);
+  for (auto cols : {16, 17}) {
+    for (auto has_bias : {true, false}) {
+      TestFusedDropoutActBias<T, Functor, GradFunctor> test(rows, cols);
+      test.has_bias = has_bias;
+      test.Run();
+      test.CheckOut(default_diff);
+      test.CheckGrad(default_diff);
+    }
+  }
+}
+
+TEST(FusedDropout, GPUFusedDorpoutActBias) {
+  BaseTest<float,
+           phi::funcs::ReluFunctor<float>,
+           phi::funcs::ReluGradFunctor<float>>();
+  BaseTest<float,
+           phi::fusion::LayerNormParamTypeGeluFunctor<float>,
+           phi::fusion::GeluGradFunctor<float>>();
+}
+TEST(FusedDropout, GPUFusedDropoutActBiasDouble) {
+  BaseTest<double,
+           phi::funcs::ReluFunctor<double>,
+           phi::funcs::ReluGradFunctor<double>>();
+  BaseTest<double,
+           phi::fusion::LayerNormParamTypeGeluFunctor<double>,
+           phi::fusion::GeluGradFunctor<double>>();
+}
+
+// test fp16, For inference, check_grad is not required. ref: test_dropout_op.py
+TEST(FusedDropout, GPUFusedDropoutActBiasFp16) {
+  using fp16 = phi::dtype::float16;
+  BaseTest<fp16,
+           phi::funcs::ReluFunctor<fp16>,
+           phi::funcs::ReluGradFunctor<fp16>>(true);
+}
+
+TEST(FusedDropout, GPUFusedDropoutActBiasIsUpscaleInTrain) {
+  const int rows = 16;
+  const int cols = 16;
+  for (auto is_upscale_in_train : {true, false}) {
+    TestFusedDropoutActBias<float,
+                            phi::funcs::ReluFunctor<float>,
+                            phi::funcs::ReluGradFunctor<float>>
+        test(rows, cols, 0, 1.0, is_upscale_in_train, false);
+    test.Run();
+    test.CheckOut(static_cast<float>(1e-5));
+    test.CheckGrad(static_cast<float>(1e-3));
+  }
+}
+
+TEST(FusedDropout, GPUFusedDropoutActBiasIsTest) {
+  const int rows = 16;
+  const int cols = 16;
+  TestFusedDropoutActBias<float,
+                          phi::funcs::ReluFunctor<float>,
+                          phi::funcs::ReluGradFunctor<float>>
+      test(rows, cols, 0, 0.35, true, true);
+  test.Run();
+  test.CheckOut(static_cast<float>(1e-5));
+  test.CheckGrad(static_cast<float>(1e-3));
+}
+
+TEST(FusedDropout, GPUFusedDropoutActBiasSeed) {
+  const int rows = 16;
+  const int cols = 16;
+  TestFusedDropoutActBias<float,
+                          phi::funcs::ReluFunctor<float>,
+                          phi::funcs::ReluGradFunctor<float>>
+      test(rows, cols, 125, 0.0, false, false);
+  test.Run();
+  test.CheckOut(static_cast<float>(1e-5));
+  test.CheckGrad(static_cast<float>(1e-3));
+}
+
+TEST(FusedDropout, GPUFusedDropoutActBiasLargeShape) {
+  const int rows = 256;
+  const int cols = 4096;
+  TestFusedDropoutActBias<float,
+                          phi::funcs::ReluFunctor<float>,
+                          phi::funcs::ReluGradFunctor<float>>
+      test(rows, cols);
+  test.Run();
+  test.CheckOut(static_cast<float>(1e-5));
+  test.CheckGrad(static_cast<float>(1e-3));
+}

test/cpp/fluid/fused/fused_dropout_test.h

@@ -0,0 +1,223 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+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. */
+
+#pragma once
+
+#include <random>
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+#include "paddle/fluid/framework/program_desc.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/memory/memory.h"
+#include "paddle/phi/kernels/funcs/layer_norm_impl.cu.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/layer_norm_kernel.h"
+#include "paddle/utils/string/printf.h"
+
+namespace framework = paddle::framework;
+namespace platform = paddle::platform;
+namespace memory = paddle::memory;
+
+USE_OP_ITSELF(dropout);
+
+template <typename T>
+using CudnnDataType = phi::backends::gpu::CudnnDataType<T>;
+template <typename T>
+using LayerNormParamType = typename CudnnDataType<T>::BatchNormParamType;
+
+/**
+ * @brief call paddle dropout op
+ */
+template <typename T>
+void Dropout(const std::vector<T> &x,
+             const phi::DDim &x_dim,
+             std::vector<T> *out,
+             std::vector<uint8_t> *mask,
+             const phi::GPUContext &ctx,
+             uint64_t seed,
+             float dropout_prob,
+             bool is_upscale_in_train,
+             bool is_test) {
+  framework::Scope scope;
+  auto var_x = scope.Var("X");
+  auto tensor_x = var_x->GetMutable<phi::DenseTensor>();
+  framework::TensorFromVector(x, ctx, tensor_x);
+  tensor_x->Resize(x_dim);
+
+  auto var_out = scope.Var("Out");
+  auto tensor_out = var_out->GetMutable<phi::DenseTensor>();
+
+  auto var_mask = scope.Var("Mask");
+  auto tensor_mask = var_mask->GetMutable<phi::DenseTensor>();
+
+  framework::AttributeMap attrs;
+  attrs.insert({"fix_seed", 1});
+  attrs.insert({"seed", static_cast<int>(seed)});
+  attrs.insert({"dropout_prob", dropout_prob});
+  if (is_upscale_in_train) {
+    attrs.insert({"dropout_implementation", std::string("upscale_in_train")});
+  }
+
+  if (is_test) {
+    attrs.insert({"is_test", true});
+  }
+
+  auto op = framework::OpRegistry::CreateOp(
+      "dropout", {{"X", {"X"}}}, {{"Out", {"Out"}}, {"Mask", {"Mask"}}}, attrs);
+  op->Run(scope, ctx.GetPlace());
+
+  framework::TensorToVector<T>(*tensor_out, ctx, out);
+  if (!is_test) {
+    framework::TensorToVector<uint8_t>(*tensor_mask, ctx, mask);
+  }
+  ctx.Wait();
+}
+
+/**
+ * @brief call paddle dropout_grad op
+ */
+template <typename T>
+void DropoutGrad(std::vector<T> *dx,
+                 const phi::DDim &x_dim,
+                 const std::vector<T> &dout,
+                 const std::vector<uint8_t> &mask,
+                 const phi::GPUContext &ctx,
+                 float dropout_prob,
+                 bool is_upscale_in_train) {
+  framework::Scope scope;
+  const size_t n = x_dim[0] * x_dim[1];
+  auto var_out = scope.Var("DOut");
+  auto tensor_out = var_out->GetMutable<phi::DenseTensor>();
+  framework::TensorFromVector(dout, ctx, tensor_out);
+  tensor_out->Resize(x_dim);
+
+  auto var_mask = scope.Var("Mask");
+  auto tensor_mask = var_mask->GetMutable<phi::DenseTensor>();
+  framework::TensorFromVector(mask, ctx, tensor_mask);
+  tensor_mask->Resize(x_dim);
+
+  auto var_dx = scope.Var("DX");
+  auto tensor_dx = var_dx->GetMutable<phi::DenseTensor>();
+
+  framework::AttributeMap attrs;
+  attrs.insert({"dropout_prob", dropout_prob});
+  attrs.insert({"is_test", false});
+  if (is_upscale_in_train) {
+    attrs.insert({"dropout_implementation", std::string("upscale_in_train")});
+  } else {
+    attrs.insert({"dropout_implementation", std::string("downgrade_in_infer")});
+  }
+
+  auto op = framework::OpRegistry::CreateOp(
+      "dropout_grad",
+      {{"Out@GRAD", {"DOut"}}, {"Mask", {"Mask"}}},
+      {{"X@GRAD", {"DX"}}},
+      attrs);
+  op->Run(scope, ctx.GetPlace());
+
+  framework::TensorToVector(*tensor_dx, ctx, dx);
+  ctx.Wait();
+}
+
+/**
+ * @brief call paddle layer_norm op
+ */
+template <typename T>
+void LayerNorm(const std::vector<LayerNormParamType<T>> &scale,
+               const std::vector<LayerNormParamType<T>> &bias,
+               const std::vector<T> &x,
+               std::vector<LayerNormParamType<T>> *means,
+               std::vector<LayerNormParamType<T>> *vars,
+               std::vector<T> *y,
+               const double epsilon,
+               const int rows,
+               const int cols,
+               const phi::GPUContext &ctx) {
+  framework::Scope scope;
+  auto place = ctx.GetPlace();
+  paddle::optional<phi::DenseTensor> scale_opt;
+  if (scale.size() > 0) {
+    auto var_scale = scope.Var("Scale");
+    auto tensor_scale = var_scale->GetMutable<phi::DenseTensor>();
+    framework::TensorFromVector(scale, ctx, tensor_scale);
+    tensor_scale->Resize({cols});
+    scale_opt = *tensor_scale;
+  }
+
+  paddle::optional<phi::DenseTensor> bias_opt;
+  if (bias.size() > 0) {
+    auto var_bias = scope.Var("Bias");
+    auto tensor_bias = var_bias->GetMutable<phi::DenseTensor>();
+    framework::TensorFromVector(bias, ctx, tensor_bias);
+    tensor_bias->Resize({cols});
+
+    bias_opt = *tensor_bias;
+  }
+
+  auto var_x = scope.Var("X");
+  auto tensor_x = var_x->GetMutable<phi::DenseTensor>();
+  framework::TensorFromVector(x, ctx, tensor_x);
+  tensor_x->Resize({rows, cols});
+
+  auto var_y = scope.Var("Y");
+  auto tensor_y = var_y->GetMutable<phi::DenseTensor>();
+  tensor_y->Resize({rows, cols});
+
+  auto var_mean = scope.Var("Mean");
+  auto tensor_mean = var_mean->GetMutable<phi::DenseTensor>();
+  tensor_mean->Resize({rows});
+
+  auto var_variance = scope.Var("Variance");
+  auto tensor_variance = var_variance->GetMutable<phi::DenseTensor>();
+  tensor_variance->Resize({rows});
+  ctx.Wait();
+  phi::LayerNormKernel<T>(static_cast<const phi::GPUContext &>(ctx),
+                          *tensor_x,
+                          scale_opt,
+                          bias_opt,
+                          1e-5,
+                          1,
+                          tensor_y,
+                          tensor_mean,
+                          tensor_variance);
+  framework::TensorToVector(*tensor_y, ctx, y);
+  framework::TensorToVector(*tensor_mean, ctx, means);
+  framework::TensorToVector(*tensor_variance, ctx, vars);
+  ctx.Wait();
+}
+
+template <typename T>
+inline void ReduceSum(const std::vector<T> &dout,
+                      std::vector<T> *dbias,
+                      const int rows,
+                      const int cols) {
+  for (int j = 0; j < cols; j++) {
+    std::vector<T> tmp_dbias(rows);
+    for (int i = 0; i < rows; i++) {
+      tmp_dbias[i] = dout[i * cols + j];
+    }
+    int tmp_rows = rows / 2;
+    while (tmp_rows) {
+      for (int i = 0; i < tmp_rows; i++) {
+        tmp_dbias[i] += tmp_dbias[i + tmp_rows];
+      }
+      tmp_rows /= 2;
+    }
+    (*dbias)[j] = tmp_dbias[0];
+  }
+}

test/cpp/fluid/fused/fused_layernorm_residual_dropout_bias_test.cu

@@ -0,0 +1,385 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+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. */
+
+#include <time.h>
+
+#include <random>
+#include <vector>
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/fusion/gpu/fused_dropout_helper.h"
+#include "test/cpp/fluid/fused/fused_dropout_test.h"
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_DECLARE_KERNEL(dropout, GPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(dropout_grad, GPU, ALL_LAYOUT);
+#endif
+
+/**
+ * @brief The unit test of fused_layernorm_residual_dropout_bias
+ */
+
+template <typename T>
+struct TestFusedLayernormResidualDropoutBias {
+  uint32_t rows;
+  uint32_t cols;
+  uint64_t seed;
+  float dropout_prob, epsilon;
+  bool is_upscale_in_train;
+  bool is_test;  // default false,  Set to true for inference only
+  bool has_bias = true;
+  bool has_scale = true;
+  bool has_layernorm_bias = true;
+  phi::DenseTensor src, residual, bias, out, mask, scale, layernorm_bias,
+      layernorm_out, means, vars;
+  phi::DenseTensor dsrc, dbias;
+
+  std::vector<T> src_vec, residual_vec, bias_vec;
+  std::vector<phi::funcs::LayerNormParamType<T>> means_vec, vars_vec, scale_vec,
+      layernorm_bias_vec;
+  std::vector<T> correct_out, correct_dsrc, correct_dbias,
+      correct_layernorm_out;
+  std::vector<phi::funcs::LayerNormParamType<T>> correct_means, correct_vars;
+  std::vector<uint8_t> correct_mask;
+
+  phi::GPUPlace place;
+  phi::GPUContext *ctx;
+
+  TestFusedLayernormResidualDropoutBias() {
+    rows = 32;
+    cols = 32;
+    seed = 0;
+    dropout_prob = 0.0;
+    is_upscale_in_train = false;
+    is_test = false;
+    has_bias = true;
+    has_scale = true;
+    has_layernorm_bias = true;
+    epsilon = 0.00001f;
+    phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
+    auto devicectx = pool.Get(place);
+    ctx = reinterpret_cast<phi::GPUContext *>(devicectx);
+  }
+
+  TestFusedLayernormResidualDropoutBias(int _rows,
+                                        int _cols,
+                                        uint64_t _seed = 0,
+                                        float _dropout_prob = 0.0,
+                                        float _epsilon = 0.00001f,
+                                        bool _is_upscale_in_train = false,
+                                        bool _is_test = false,
+                                        bool _has_bias = true) {
+    rows = _rows;
+    cols = _cols;
+    seed = _seed;
+    dropout_prob = _dropout_prob;
+    epsilon = _epsilon;
+    is_upscale_in_train = _is_upscale_in_train;
+    is_test = _is_test;
+    has_bias = _has_bias;
+    has_scale = true;
+    has_layernorm_bias = true;
+    phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
+    auto devicectx = pool.Get(place);
+    ctx = reinterpret_cast<phi::GPUContext *>(devicectx);
+  }
+
+  ~TestFusedLayernormResidualDropoutBias() = default;
+
+  void SetUp() {
+    using U = phi::funcs::LayerNormParamType<T>;
+    const int n = rows * cols;
+    correct_out.resize(n);
+    correct_mask.resize(n);
+    correct_dsrc.resize(n);
+    correct_dbias.resize(cols);
+    correct_means.resize(rows);
+    correct_vars.resize(rows);
+    correct_layernorm_out.resize(n);
+
+    src_vec.resize(n);
+    residual_vec.resize(n);
+    if (has_bias) {
+      bias_vec.resize(cols);
+    }
+    if (has_scale) {
+      scale_vec.resize(cols);
+    }
+    if (has_layernorm_bias) {
+      layernorm_bias_vec.resize(cols);
+    }
+    std::default_random_engine random(time(NULL));
+    std::uniform_real_distribution<float> dis(0.0, 1.0);
+
+    for (int i = 0; i < rows; i++) {
+      for (int j = 0; j < cols; j++) {
+        src_vec[i * cols + j] = static_cast<T>(dis(random));
+        residual_vec[i * cols + j] = static_cast<T>(dis(random));
+        if (i == 0) {
+          if (has_bias) {
+            bias_vec[j] = static_cast<T>(dis(random));
+          }
+          if (has_scale) {
+            scale_vec[j] = static_cast<U>(dis(random));
+          }
+          if (has_layernorm_bias) {
+            layernorm_bias_vec[j] = static_cast<U>(dis(random));
+          }
+        }
+      }
+    }
+
+    framework::TensorFromVector<T>(src_vec, *ctx, &src);
+    src.Resize({rows, cols});
+    framework::TensorFromVector<T>(residual_vec, *ctx, &residual);
+    residual.Resize({rows, cols});
+    if (has_bias) {
+      framework::TensorFromVector<T>(bias_vec, *ctx, &bias);
+      bias.Resize({cols});
+    }
+    if (has_scale) {
+      framework::TensorFromVector<U>(scale_vec, *ctx, &scale);
+      scale.Resize({cols});
+    }
+    if (has_layernorm_bias) {
+      framework::TensorFromVector<U>(layernorm_bias_vec, *ctx, &layernorm_bias);
+      layernorm_bias.Resize({cols});
+    }
+
+    {
+      out.Resize({rows, cols});
+      out.mutable_data<T>(place);
+      mask.Resize({rows, cols});
+      mask.mutable_data<uint8_t>(place);
+      means.Resize({rows});
+      means.mutable_data<U>(place);
+      vars.Resize({rows});
+      vars.mutable_data<U>(place);
+      layernorm_out.Resize({rows, cols});
+      layernorm_out.mutable_data<T>(place);
+      dsrc.Resize({rows, cols});
+      dsrc.mutable_data<T>(place);
+
+      if (has_bias) {
+        dbias.Resize({cols});
+        dbias.mutable_data<T>(place);
+      }
+    }
+  }
+
+  void BaseForward() {
+    using U = phi::funcs::LayerNormParamType<T>;
+    std::vector<T> out1(rows * cols), out2(rows * cols);
+    if (has_bias) {
+      // add bias
+      for (int i = 0; i < rows; i++) {
+        for (int j = 0; j < cols; j++) {
+          out1[i * cols + j] = src_vec[i * cols + j] + bias_vec[j];
+        }
+      }
+      // call dropout
+      Dropout<T>(out1,
+                 src.dims(),
+                 &out2,
+                 &correct_mask,
+                 *ctx,
+                 seed,
+                 dropout_prob,
+                 is_upscale_in_train,
+                 is_test);
+    } else {
+      Dropout<T>(src_vec,
+                 src.dims(),
+                 &out2,
+                 &correct_mask,
+                 *ctx,
+                 seed,
+                 dropout_prob,
+                 is_upscale_in_train,
+                 is_test);
+    }
+    // add residual
+    for (int i = 0; i < rows; i++) {
+      for (int j = 0; j < cols; j++) {
+        correct_out[i * cols + j] =
+            residual_vec[i * cols + j] + out2[i * cols + j];
+      }
+    }
+    LayerNorm<T>(scale_vec,
+                 layernorm_bias_vec,
+                 correct_out,
+                 &correct_means,
+                 &correct_vars,
+                 &correct_layernorm_out,
+                 epsilon,
+                 rows,
+                 cols,
+                 *ctx);
+    ctx->Wait();
+  }
+
+  void FusedForward() {
+    using U = phi::funcs::LayerNormParamType<T>;
+    int VecSize = MAX_CACHE_BYTES / sizeof(T);
+    if (cols % 4 != 0) {
+      VecSize = 1;
+    }
+    int threads = phi::funcs::GetDesiredBlockDim(cols / VecSize);
+    const int increment = ((cols - 1) / (threads * VecSize) + 1) * VecSize;
+
+    T *bias_ptr = nullptr;
+    U *scale_ptr = nullptr;
+    U *layernorm_bias_ptr = nullptr;
+    if (has_bias) {
+      bias_ptr = bias.data<T>();
+    }
+    if (has_scale) {
+      scale_ptr = scale.data<U>();
+    }
+    if (has_layernorm_bias) {
+      layernorm_bias_ptr = layernorm_bias.data<U>();
+    }
+
+    phi::fusion::LaunchLayernormResidualDropoutBias<T, uint8_t, U, false>(
+        rows,
+        cols,
+        increment,
+        seed,
+        dropout_prob,
+        epsilon,
+        is_upscale_in_train,
+        is_test,
+        src.data<T>(),
+        residual.data<T>(),
+        bias_ptr,
+        scale_ptr,
+        layernorm_bias_ptr,
+        mask.data<uint8_t>(),
+        out.data<T>(),
+        layernorm_out.data<T>(),
+        means.data<U>(),
+        vars.data<U>(),
+        *ctx);
+    ctx->Wait();
+  }
+
+  void Run() {
+    SetUp();
+    BaseForward();
+    FusedForward();
+  }
+
+  void CheckOut(const T diff) {
+    using U = phi::funcs::LayerNormParamType<T>;
+    const int n = rows * cols;
+    std::vector<T> _out(n), _layernorm_out(n);
+    std::vector<U> _means(rows), _vars(cols);
+    std::vector<uint8_t> _mask(n);
+    framework::TensorToVector(out, *ctx, &_out);
+    framework::TensorToVector(layernorm_out, *ctx, &_layernorm_out);
+    framework::TensorToVector(means, *ctx, &_means);
+    framework::TensorToVector(vars, *ctx, &_vars);
+    if (!is_test && dropout_prob != 0.0f) {
+      framework::TensorToVector(mask, *ctx, &_mask);
+    }
+    ctx->Wait();
+
+    for (int i = 0; i < n; i++) {
+      EXPECT_LT(std::abs(_out[i] - correct_out[i]), diff);
+      EXPECT_LT(std::abs(_layernorm_out[i] - correct_layernorm_out[i]), diff);
+      if (!is_test && dropout_prob != 0.0f) {
+        EXPECT_EQ(_mask[i], correct_mask[i]);
+      }
+    }
+    for (int i = 0; i < rows; i++) {
+      EXPECT_LT(std::abs(_means[i] - correct_means[i]), static_cast<U>(diff));
+      EXPECT_LT(std::abs(_vars[i] - correct_vars[i]), static_cast<U>(diff));
+    }
+  }
+};
+
+template <typename T>
+static void BaseTest(const bool is_fp16 = false) {
+  const int rows = 16;
+  T default_diff = !is_fp16 ? static_cast<T>(1e-4) : static_cast<T>(1e-2);
+  for (auto cols : {16, 17}) {
+    for (auto has_bias : {true, false}) {
+      for (auto has_scale : {true, false}) {
+        for (auto has_layernorm_bias : {true, false}) {
+          TestFusedLayernormResidualDropoutBias<T> test(rows, cols);
+          test.has_bias = has_bias;
+          test.has_scale = has_scale;
+          test.has_layernorm_bias = has_layernorm_bias;
+          test.Run();
+          test.CheckOut(default_diff);
+        }
+      }
+    }
+  }
+}
+TEST(FusedDropout, GPUFusedLayernormResidualDropoutBias) { BaseTest<float>(); }
+
+TEST(FusedDropout, GPUFusedLayernormResidualDropoutBiasDouble) {
+  BaseTest<double>();
+}
+
+TEST(FusedDropout, GPUFusedLayernormResidualDropoutBiasFp16) {
+  BaseTest<phi::dtype::float16>(true);
+}
+
+TEST(FusedDropout, GPUFusedLayernormResidualDropoutBiasIsUpscaleInTrain) {
+  const int rows = 16;
+  const int cols = 16;
+  for (auto is_upscale_in_train : {true, false}) {
+    TestFusedLayernormResidualDropoutBias<float> test(
+        rows, cols, 0, 1.0, 0.00001f, is_upscale_in_train, false);
+    test.Run();
+    test.CheckOut(static_cast<float>(1e-4));
+  }
+}
+
+TEST(FusedDropout, GPUFusedLayernormResidualDropoutBiasIsTest) {
+  const int rows = 16;
+  const int cols = 16;
+  TestFusedLayernormResidualDropoutBias<float> test(
+      rows, cols, 0, 0.35, 0.00001f, true, true);
+  test.Run();
+  test.CheckOut(static_cast<float>(1e-4));
+}
+
+TEST(FusedDropout, GPUFusedLayernormResidualDropoutBiasSeed) {
+  const int rows = 16;
+  const int cols = 16;
+  TestFusedLayernormResidualDropoutBias<float> test(
+      rows, cols, 125, 0.0, 0.00001f, false, false);
+  test.Run();
+  test.CheckOut(static_cast<float>(1e-4));
+}
+
+TEST(FusedDropout, GPUFusedLayernormResidualDropoutLargeShape) {
+  const int rows = 512;
+  const int cols = 512;
+  TestFusedLayernormResidualDropoutBias<float> test(rows, cols);
+  test.Run();
+  test.CheckOut(static_cast<float>(1e-4));
+}
+
+TEST(FusedDropout, GPUFusedLayernormResidualDropoutFp16MLperf) {
+  const int rows = 512;
+  const int cols = 1024;
+  TestFusedLayernormResidualDropoutBias<phi::dtype::float16> test(
+      rows, cols, 0, 0, 0.00001f, false, false, false);
+  test.Run();
+  test.CheckOut(static_cast<phi::dtype::float16>(1e-2));
+}

test/cpp/fluid/fused/fused_residual_dropout_bias_test.cu

@@ -0,0 +1,396 @@
+/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+
+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. */
+
+#include <time.h>
+
+#include <random>
+#include <vector>
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/fusion/gpu/fused_residual_dropout_bias.h"
+#include "test/cpp/fluid/fused/fused_dropout_test.h"
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_DECLARE_KERNEL(dropout, GPU, ALL_LAYOUT);
+PD_DECLARE_KERNEL(dropout_grad, GPU, ALL_LAYOUT);
+#endif
+
+namespace framework = paddle::framework;
+namespace platform = paddle::platform;
+
+bool CheckEqual(float value, float ref) { return std::abs(value - ref) < 1e-5; }
+
+/**
+ * @brief the unittest of FusedResidualDropoutBias
+ * 1. random input data
+ * 2. add bias, call paddle dropout op, add residual, and get the base result
+ * 3. call FusedResidualDropoutBias function get fused result
+ * 4. compare the base result and fused result
+ */
+
+template <typename T>
+struct FusedResidualDropoutBiasTester {
+  uint32_t rows;
+  uint32_t cols;
+  uint64_t seed;
+  float dropout_prob;
+  bool is_upscale_in_train;
+  bool is_test;  // default false,  Set to true for inference only
+  bool has_bias = true;
+  bool add_residual = true;
+
+  phi::DenseTensor src, residual, bias, out, mask;
+  phi::DenseTensor dsrc, dbias;
+
+  std::vector<T> src_vec, residual_vec, bias_vec;
+  std::vector<T> correct_out, correct_dsrc, correct_dbias;
+  std::vector<uint8_t> correct_mask;
+
+  phi::GPUPlace place;
+  phi::GPUContext *ctx;
+
+  FusedResidualDropoutBiasTester() {
+    rows = 32;
+    cols = 32;
+    seed = 0;
+    dropout_prob = 0.0;
+    is_upscale_in_train = false;
+    is_test = false;
+    phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
+    auto device_ctx = pool.Get(place);
+    ctx = reinterpret_cast<phi::GPUContext *>(device_ctx);
+  }
+
+  FusedResidualDropoutBiasTester(int rows,
+                                 int cols,
+                                 uint64_t seed = 0,
+                                 float dropout_prob = 0.0,
+                                 bool is_upscale_in_train = false,
+                                 bool is_test = false)
+      : rows(rows),
+        cols(cols),
+        seed(seed),
+        dropout_prob(dropout_prob),
+        is_upscale_in_train(is_upscale_in_train),
+        is_test(is_test) {
+    phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
+    auto device_ctx = pool.Get(place);
+    ctx = reinterpret_cast<phi::GPUContext *>(device_ctx);
+  }
+
+  void SetUp() {
+    const int n = rows * cols;
+    correct_out.resize(n);
+    correct_mask.resize(n);
+    correct_dsrc.resize(n);
+    correct_dbias.resize(cols);
+
+    src_vec.resize(n);
+    if (add_residual) {
+      residual_vec.resize(n);
+    }
+    bias_vec.resize(cols);
+    std::default_random_engine random(time(NULL));
+    std::uniform_real_distribution<float> dis(0.0, 1.0);
+
+    for (int i = 0; i < rows; i++) {
+      for (int j = 0; j < cols; j++) {
+        src_vec[i * cols + j] = static_cast<T>(dis(random));
+        if (add_residual) {
+          residual_vec[i * cols + j] = static_cast<T>(dis(random));
+        }
+        if (i == 0) {
+          bias_vec[j] = dis(random);
+        }
+      }
+    }
+
+    framework::TensorFromVector<T>(src_vec, *ctx, &src);
+    src.Resize({rows, cols});
+    if (add_residual) {
+      framework::TensorFromVector<T>(residual_vec, *ctx, &residual);
+      residual.Resize({rows, cols});
+    }
+    if (has_bias) {
+      framework::TensorFromVector<T>(bias_vec, *ctx, &bias);
+      bias.Resize({cols});
+    }
+
+    out.mutable_data<T>({rows, cols}, place);
+    mask.mutable_data<uint8_t>({rows, cols}, place);
+    dsrc.mutable_data<T>({rows, cols}, place);
+
+    if (has_bias) {
+      dbias.mutable_data<T>({cols}, place);
+    }
+  }
+
+  void BaseForward() {
+    if (has_bias) {
+      // add bias
+      std::vector<T> bias_out(rows * cols);
+      for (int i = 0; i < rows; i++) {
+        for (int j = 0; j < cols; j++) {
+          bias_out[i * cols + j] = src_vec[i * cols + j] + bias_vec[j];
+        }
+      }
+      // call dropout
+      Dropout<T>(bias_out,
+                 src.dims(),
+                 &correct_out,
+                 &correct_mask,
+                 *ctx,
+                 seed,
+                 dropout_prob,
+                 is_upscale_in_train,
+                 is_test);
+    } else {
+      Dropout<T>(src_vec,
+                 src.dims(),
+                 &correct_out,
+                 &correct_mask,
+                 *ctx,
+                 seed,
+                 dropout_prob,
+                 is_upscale_in_train,
+                 is_test);
+    }
+    ctx->Wait();
+    PADDLE_ENFORCE_GPU_SUCCESS(platform::GpuGetLastError());
+    if (add_residual) {
+      // add residual
+      for (int i = 0; i < rows; i++) {
+        for (int j = 0; j < cols; j++) {
+          int idx = i * cols + j;
+          correct_out[idx] = residual_vec[idx] + correct_out[idx];
+        }
+      }
+    }
+  }
+
+  void BaseBackward() {
+    DropoutGrad<T>(&correct_dsrc,
+                   src.dims(),
+                   correct_out,
+                   correct_mask,
+                   *ctx,
+                   dropout_prob,
+                   is_upscale_in_train);
+    // calc dbias
+    memset(&correct_dbias[0], 0, cols * sizeof(T));
+    if (has_bias) {
+      ReduceSum<T>(correct_out, &correct_dbias, rows, cols);
+    }
+  }
+
+  void FusedForward() {
+    const int VecSize = MAX_CACHE_BYTES / sizeof(T);
+    auto config =
+        phi::fusion::Get1DBlocksAnd2DGrids(*ctx,
+                                           static_cast<uint64_t>(rows),
+                                           static_cast<uint64_t>(cols),
+                                           VecSize);
+
+    const int increment = ((cols - 1) / (config.thread_per_block.x *
+                                         config.block_per_grid.x * VecSize) +
+                           1) *
+                          VecSize;
+
+    T *bias_ptr = has_bias ? bias.data<T>() : nullptr;
+    T *residual_ptr = add_residual ? residual.data<T>() : nullptr;
+    phi::fusion::LaunchResidualDropoutBias<T, uint8_t>(rows,
+                                                       cols,
+                                                       increment,
+                                                       seed,
+                                                       dropout_prob,
+                                                       is_test,
+                                                       is_upscale_in_train,
+                                                       src.data<T>(),
+                                                       residual_ptr,
+                                                       bias_ptr,
+                                                       mask.data<uint8_t>(),
+                                                       out.data<T>(),
+                                                       *ctx);
+    ctx->Wait();
+    PADDLE_ENFORCE_GPU_SUCCESS(platform::GpuGetLastError());
+  }
+
+  void FusedBackward() {
+    if (is_test) {
+      return;
+    }
+
+    T *bias_ptr = has_bias ? dbias.data<T>() : nullptr;
+    phi::fusion::LaunchResidualDropoutBiasGrad<T, uint8_t>(out.data<T>(),
+                                                           mask.data<uint8_t>(),
+                                                           dropout_prob,
+                                                           is_upscale_in_train,
+                                                           rows,
+                                                           cols,
+                                                           dsrc.data<T>(),
+                                                           bias_ptr,
+                                                           *ctx);
+  }
+
+  void Run() {
+    SetUp();
+    BaseForward();
+    FusedForward();
+    BaseBackward();
+    FusedBackward();
+  }
+
+  void CheckOut(const T diff) {
+    const int n = rows * cols;
+    std::vector<T> fused_out(n);
+    std::vector<uint8_t> fused_mask(n);
+    framework::TensorToVector(out, *ctx, &fused_out);
+    if (!is_test && dropout_prob != 0.0f) {
+      framework::TensorToVector<uint8_t>(mask, *ctx, &fused_mask);
+    }
+    ctx->Wait();
+
+    for (int i = 0; i < n; i++) {
+      EXPECT_LT(std::abs(fused_out[i] - correct_out[i]), diff);
+      if (!is_test && dropout_prob != 0.0f) {
+        EXPECT_EQ(fused_mask[i], correct_mask[i]);
+      }
+    }
+  }
+
+  void CheckGrad(const T diff) {
+    if (is_test) {
+      return;
+    }
+
+    const int n = rows * cols;
+
+    std::vector<T> _dsrc(n);
+    framework::TensorToVector(dsrc, *ctx, &_dsrc);
+
+    for (int i = 0; i < n; i++) {
+      EXPECT_LT(std::abs(_dsrc[i] - correct_dsrc[i]), diff);
+    }
+
+    if (has_bias) {
+      std::vector<T> _dbias(cols);
+      framework::TensorToVector(dbias, *ctx, &_dbias);
+      ctx->Wait();
+      for (int i = 0; i < cols; i++) {
+        EXPECT_LT(std::abs(_dbias[i] - correct_dbias[i]), diff);
+      }
+    }
+  }
+};
+
+// test the shape and bias
+template <typename T>
+static void BaseTest() {
+  const int rows = 16;
+  T max_diff = static_cast<T>(0);
+  if (std::is_same<T, phi::dtype::float16>::value) {
+    max_diff = static_cast<T>(1e-1);
+  } else {
+    max_diff = static_cast<T>(1e-5);
+  }
+  for (auto cols : {16, 17}) {
+    for (auto has_bias : {true, false}) {
+      FusedResidualDropoutBiasTester<T> test(rows, cols);
+      test.has_bias = has_bias;
+      test.Run();
+      test.CheckOut(max_diff);
+      test.CheckGrad(max_diff);
+    }
+  }
+}
+
+TEST(FusedDropout, GPUFusedResidualDropoutBias) { BaseTest<float>(); }
+
+TEST(FusedDropout, GPUFusedResidualDropoutBiasDouble) { BaseTest<double>(); }
+
+TEST(FusedDropout, GPUFusedResidualDropoutBiasFp16) {
+  BaseTest<phi::dtype::float16>();
+}
+
+TEST(FusedDropout, GPUFusedResidualDropoutBiasIsUpscaleInTrain) {
+  const int rows = 16;
+  const int cols = 16;
+  for (auto is_upscale_in_train : {true, false}) {
+    FusedResidualDropoutBiasTester<float> test(
+        rows, cols, 0, 1.0, is_upscale_in_train, false);
+    test.Run();
+    test.CheckOut(static_cast<float>(1e-5));
+    test.CheckGrad(static_cast<float>(1e-5));
+  }
+}
+
+TEST(FusedDropout, GPUFusedResidualDropoutBiasIsTest) {
+  const int rows = 16;
+  const int cols = 16;
+  FusedResidualDropoutBiasTester<float> test(rows, cols, 0, 0.35, true, true);
+  test.Run();
+  test.CheckOut(static_cast<float>(1e-5));
+  test.CheckGrad(static_cast<float>(1e-5));
+}
+
+TEST(FusedDropout, GPUFusedResidualDropoutBiasSeed) {
+  const int rows = 16;
+  const int cols = 16;
+  FusedResidualDropoutBiasTester<float> test(
+      rows, cols, 125, 0.0, false, false);
+  test.Run();
+  test.CheckOut(static_cast<float>(1e-5));
+  test.CheckGrad(static_cast<float>(1e-5));
+}
+
+TEST(FusedDropout, NoResidual) {
+  const int rows = 16;
+  const int cols = 16;
+  for (float p : {0.0f, 0.5f, 1.0f}) {
+    FusedResidualDropoutBiasTester<float> test(rows, cols, 0, p, false, false);
+    test.add_residual = false;
+    test.Run();
+    // For a non 0 or 1 dropout_prob, just test whether it can run successly.
+    if (CheckEqual(p, 0.0f) || CheckEqual(p, 1.0f)) {
+      test.CheckOut(static_cast<float>(1e-5));
+      test.CheckGrad(static_cast<float>(1e-5));
+    }
+  }
+}
+
+TEST(FusedDropout, GPUFusedResidualDropoutBiasLargeShape) {
+  const int rows = 256;
+  const int cols = 4096;
+  FusedResidualDropoutBiasTester<float> test(rows, cols);
+  test.Run();
+  test.CheckOut(static_cast<float>(1e-5));
+  test.CheckGrad(static_cast<float>(1e-3));
+}
+
+TEST(FusedDropout, GPUFusedResidualDropoutBiasLargeShapeFp16) {
+  // Used to test that `cudaErrorLaunchOutOfResources` will not occur
+  int rows = 1;
+  int cols = 12288;
+  if (std::getenv("_rows") != nullptr) {
+    rows = atoi(std::getenv("_rows"));
+  }
+  if (std::getenv("_cols") != nullptr) {
+    cols = atoi(std::getenv("_cols"));
+  }
+  FusedResidualDropoutBiasTester<phi::dtype::float16> test(
+      rows, cols, 0, 0.0, true, true);
+  test.Run();
+  test.CheckOut(static_cast<phi::dtype::float16>(1e-1));
+  test.CheckGrad(static_cast<phi::dtype::float16>(1e-1));
+}

test/cpp/fluid/fused/fusion_group_op_test.cc

@@ -0,0 +1,237 @@
+/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+
+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. */
+
+#include "gtest/gtest.h"
+#include "paddle/fluid/framework/op_desc.h"
+#include "paddle/fluid/framework/op_proto_maker.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/program_desc.h"
+#include "paddle/fluid/platform/init.h"
+#include "paddle/phi/backends/device_code.h"
+
+namespace paddle {
+namespace operators {
+
+using CPUKernelFunc = std::function<void(size_t n, std::vector<void*> args)>;
+
+template <typename T>
+phi::DenseTensor* CreateTensor(framework::Scope* scope,
+                               const phi::Place& place,
+                               const std::string& name,
+                               const std::vector<int64_t>& shape) {
+  auto* var = scope->Var(name);
+  auto* tensor = var->GetMutable<phi::DenseTensor>();
+  if (!shape.empty()) {
+    tensor->mutable_data<T>(common::make_ddim(shape), place);
+  }
+  return tensor;
+}
+
+template <typename T>
+void SetupRandomCPUTensor(phi::DenseTensor* tensor,
+                          const std::vector<int64_t>& shape) {
+  static unsigned int seed = 100;
+  std::mt19937 rng(seed++);
+  std::uniform_real_distribution<double> uniform_dist(0, 1);
+
+  T* ptr = tensor->mutable_data<T>(common::make_ddim(shape), phi::CPUPlace());
+  for (int64_t i = 0; i < tensor->numel(); ++i) {
+    ptr[i] = static_cast<T>(uniform_dist(rng)) - static_cast<T>(0.5);
+  }
+}
+
+framework::OpDesc* CreateFusionGroupOp(
+    framework::ProgramDesc* program,
+    const std::vector<std::string>& input_names,
+    const std::vector<std::vector<int64_t>>& input_shapes,
+    const std::vector<std::string>& output_names,
+    int type,
+    std::string func_name) {
+  EXPECT_EQ(input_names.size(), input_shapes.size());
+
+  std::vector<int> input_dtypes(input_names.size(),
+                                framework::proto::VarType::FP32);
+  std::vector<int> output_dtypes(output_names.size(),
+                                 framework::proto::VarType::FP32);
+
+  for (size_t i = 0; i < input_names.size(); ++i) {
+    auto* var = program->MutableBlock(0)->Var(input_names[i]);
+    var->SetType(framework::proto::VarType::DENSE_TENSOR);
+    var->SetDataType(framework::proto::VarType::FP32);
+    var->SetShape(input_shapes[i]);
+  }
+  for (const auto& output_name : output_names) {
+    auto* var = program->MutableBlock(0)->Var(output_name);
+    var->SetType(framework::proto::VarType::DENSE_TENSOR);
+    var->SetDataType(framework::proto::VarType::FP32);
+  }
+
+  auto* op = program->MutableBlock(0)->AppendOp();
+  op->SetType("fusion_group");
+  op->SetInput("Inputs", input_names);
+  op->SetOutput("Outs", output_names);
+  op->SetAttr("inputs_dtype", input_dtypes);
+  op->SetAttr("outs_dtype", output_dtypes);
+  op->SetAttr("type", type);
+  op->SetAttr("func_name", func_name);
+  op->SetAttr(framework::OpProtoAndCheckerMaker::OpRoleAttrName(),
+              static_cast<int>(framework::OpRole::kForward));
+  return op;
+}
+
+void PrepareDeviceCode(phi::Place place,
+                       std::string func_name,
+                       std::string cuda_kernel_str) {
+  phi::DeviceCodePool& pool = phi::DeviceCodePool::Init({place});
+
+  std::unique_ptr<phi::DeviceCode> code(
+      new phi::GPUDeviceCode(place, func_name, cuda_kernel_str));
+  code->Compile();
+  pool.Set(std::move(code));
+}
+
+void CheckOutputs(framework::Scope* scope,
+                  const std::vector<std::string>& output_names,
+                  std::vector<phi::DenseTensor>* cpu_tensors,
+                  size_t num_inputs,
+                  CPUKernelFunc cpu_kernel_func) {
+  std::vector<phi::DenseTensor> cpu_outputs;
+  cpu_outputs.resize(output_names.size());
+  for (size_t j = 0; j < output_names.size(); ++j) {
+    auto* var = scope->Var(output_names[j]);
+    const auto& dev_tensor = var->Get<phi::DenseTensor>();
+    paddle::framework::TensorCopySync(
+        dev_tensor, phi::CPUPlace(), &(cpu_outputs[j]));
+
+    cpu_tensors->at(num_inputs + j)
+        .mutable_data<float>(dev_tensor.dims(), phi::CPUPlace());
+  }
+
+  size_t n = cpu_tensors->at(0).numel();
+  std::vector<void*> args;
+  for (auto& cpu_tensor : *cpu_tensors) {
+    args.push_back(cpu_tensor.data<float>());
+  }
+  cpu_kernel_func(n, args);
+
+  for (size_t j = 0; j < output_names.size(); ++j) {
+    auto* dev_ptr = cpu_outputs[j].data<float>();
+    auto* cpu_ptr = cpu_tensors->at(num_inputs + j).data<float>();
+    int64_t length = cpu_outputs[j].numel();
+    LOG(INFO) << "Check the " << j << "th output...";
+    for (int64_t i = 0; i < length; ++i) {
+      EXPECT_NEAR(dev_ptr[i], cpu_ptr[i], 1.E-05);
+    }
+  }
+}
+
+void TestMain(const std::vector<std::string>& input_names,
+              const std::vector<std::vector<int64_t>>& input_shapes,
+              const std::vector<std::string>& output_names,
+              int type,
+              std::string func_name,
+              std::string cuda_kernel_str,
+              CPUKernelFunc cpu_kernel_func) {
+  // Compile the device code
+  paddle::framework::InitDevices({0});
+  phi::GPUPlace place = phi::GPUPlace(0);
+  PrepareDeviceCode(place, func_name, cuda_kernel_str);
+
+  // Create a ProgramDesc that has a fusion_group_op.
+  framework::ProgramDesc program;
+  framework::OpDesc* op_desc = CreateFusionGroupOp(
+      &program, input_names, input_shapes, output_names, type, func_name);
+  auto fusion_group_op = framework::OpRegistry::CreateOp(*op_desc);
+
+  framework::Scope scope;
+
+  // Prepare input tensors.
+  std::vector<phi::DenseTensor> cpu_tensors;
+  cpu_tensors.resize(input_names.size() + output_names.size());
+  for (size_t i = 0; i < input_names.size(); ++i) {
+    SetupRandomCPUTensor<float>(&(cpu_tensors[i]), input_shapes[i]);
+    phi::DenseTensor* dev_tensor =
+        CreateTensor<float>(&scope, place, input_names[i], input_shapes[i]);
+    paddle::framework::TensorCopySync(cpu_tensors[i], place, dev_tensor);
+  }
+  // Create output tensors.
+  std::vector<int64_t> empty_shape;
+  for (const auto& output_name : output_names) {
+    CreateTensor<float>(&scope, place, output_name, empty_shape);
+  }
+
+  fusion_group_op->Run(scope, place);
+
+  auto* dev_ctx = phi::DeviceContextPool::Instance().Get(place);
+  dev_ctx->Wait();
+
+  // Check the output.
+  CheckOutputs(
+      &scope, output_names, &cpu_tensors, input_names.size(), cpu_kernel_func);
+}
+
+TEST(FusionGroupOp, elementwise) {
+  if (!phi::dynload::HasNVRTC() || !phi::dynload::HasCUDADriver()) {
+    return;
+  }
+
+  // z = relu(x + y)
+  std::vector<std::string> input_names = {"x", "y"};
+  std::vector<std::string> output_names = {"z"};
+  std::vector<std::vector<int64_t>> input_shapes = {{256, 256}, {256, 256}};
+  constexpr auto kernel = R"(
+static inline __device__ float relu(float x) {
+  return x * (x > 0);
+}
+
+extern "C" __global__
+void elementwise_cuda_kernel_0(size_t n, float *x, float* y, float* z) {
+  for (size_t tid = blockIdx.x * blockDim.x + threadIdx.x; tid < n;
+       tid += blockDim.x * gridDim.x) {
+    float tmp_0 = x[tid];
+    float tmp_1 = y[tid];
+    float tmp_2 = tmp_0 + tmp_1;
+    float tmp_3 = relu(tmp_2);
+    z[tid] = tmp_3;
+  }
+})";
+
+  auto elementwise_cpu_kernel_0 = [](size_t n,
+                                     std::vector<void*> args) -> void {
+    float* x = static_cast<float*>(args[0]);
+    float* y = static_cast<float*>(args[1]);
+    float* z = static_cast<float*>(args[2]);
+    for (size_t i = 0; i < n; ++i) {
+      float tmp_0 = x[i];
+      float tmp_1 = y[i];
+      float tmp_2 = tmp_0 + tmp_1;
+      float tmp_3 = tmp_2 > 0 ? tmp_2 : 0;
+      z[i] = tmp_3;
+    }
+  };
+
+  TestMain(input_names,
+           input_shapes,
+           output_names,
+           0,
+           "elementwise_cuda_kernel_0",
+           kernel,
+           elementwise_cpu_kernel_0);
+}
+
+}  // namespace operators
+}  // namespace paddle
+
+USE_OP_ITSELF(fusion_group);
+PD_DECLARE_KERNEL(fusion_group, GPU, ALL_LAYOUT);