Merge pull request #12681 from PaddlePaddle/revert-12554-refine_elementwise_add

Revert "Refine elementwise_add op"

Author: tensor-tang

paddle/fluid/operators/elementwise_add_op.cu

@@ -16,60 +16,6 @@ limitations under the License. */
 #include "paddle/fluid/operators/elementwise_add_op.h"
 #include "paddle/fluid/platform/float16.h"
 
-namespace paddle {
-namespace operators {
-
-template <typename T>
-__global__ void ElementwiseAddCUDAKernel(const T *x, const T *y, T *z, int n,
-                                         int post, int size) {
-  int idx_x = threadIdx.x + blockIdx.x * blockDim.x;
-  if (idx_x < size) {
-    int idx_y = idx_x / post - (idx_x / (n * post)) * n;
-    z[idx_x] = x[idx_x] + y[idx_y];
-  }
-}
-
-template <typename T>
-class ElementwiseAddKernel<platform::CUDADeviceContext, T>
-    : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext &ctx) const override {
-    using Tensor = framework::Tensor;
-
-    const auto x = ctx.Input<Tensor>("X");
-    const auto y = ctx.Input<Tensor>("Y");
-    auto z = ctx.Output<Tensor>("Out");
-    auto *z_data = z->mutable_data<T>(ctx.GetPlace());
-
-    auto &device = *(ctx.cuda_device_context().eigen_device());
-    const framework::DDim &x_dim = x->dims();
-    framework::DDim y_dim = y->dims();
-    int size = x->numel();
-    if (x_dim == y_dim) {
-      auto dim = framework::make_ddim({size});
-      auto z_eigen = framework::EigenTensor<T, 1>::From(*z, dim);
-      auto x_eigen = framework::EigenTensor<T, 1>::From(*x, dim);
-      auto y_eigen = framework::EigenTensor<T, 1>::From(*y, dim);
-      z_eigen.device(device) = x_eigen + y_eigen;
-    } else {
-      int axis = ctx.Attr<int>("axis");
-      axis = (axis == -1 ? x_dim.size() - y_dim.size() : axis);
-      y_dim = trim_trailing_singular_dims(y_dim);
-      axis = (y_dim.size() == 0) ? x_dim.size() : axis;
-      int pre, n, post;
-      get_mid_dims(x_dim, y_dim, axis, &pre, &n, &post);
-      int threads = 512;
-      int grids = (size + threads - 1) / threads;
-      auto stream = ctx.cuda_device_context().stream();
-      ElementwiseAddCUDAKernel<T><<<grids, threads, 0, stream>>>(
-          x->data<T>(), y->data<T>(), z_data, n, post, size);
-    }
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
 namespace ops = paddle::operators;
 namespace plat = paddle::platform;
 

paddle/fluid/operators/elementwise_add_op.h

@@ -144,41 +144,16 @@ class ElementwiseAddGradKernel : public framework::OpKernel<T> {
     auto* dout = ctx.Input<Tensor>(framework::GradVarName("Out"));
     auto* dx = ctx.Output<Tensor>(framework::GradVarName("X"));
     auto* dy = ctx.Output<Tensor>(framework::GradVarName("Y"));
+    // skip out, x, y
+    auto* out = dout;
+    auto *x = dout, *y = dout;
 
-    if (dx != nullptr) {
-      // In fact, we can just share memory, but it may cause a bug of memory
-      // optimizer
-      // dx->ShareDataWith(*dout);
-      framework::TensorCopy(*dout, ctx.GetPlace(),
-                            ctx.template device_context<DeviceContext>(), dx);
-    }
-
-    if (dy == nullptr) return;
-
-    const framework::DDim& x_dim = dout->dims();
-    framework::DDim y_dim = dy->dims();
-    if (x_dim == y_dim) {
-      // dy->ShareDataWith(*dout);
-      framework::TensorCopy(*dout, ctx.GetPlace(),
-                            ctx.template device_context<DeviceContext>(), dy);
+    if (platform::is_cpu_place(ctx.GetPlace()) && dx != nullptr &&
+        dy != nullptr && (dx->dims() == dy->dims())) {
+      elementwise_add_grad<DeviceContext, T>(ctx, x, y, out, dout, dx, dy);
     } else {
-      dy->mutable_data<T>(ctx.GetPlace());
-      // Perform reduction to dout to calculate dy
-      int axis = ctx.Attr<int>("axis");
-      axis = (axis == -1 ? x_dim.size() - y_dim.size() : axis);
-      y_dim = trim_trailing_singular_dims(y_dim);
-      axis = (y_dim.size() == 0) ? x_dim.size() : axis;
-
-      auto& device =
-          *(ctx.template device_context<DeviceContext>().eigen_device());
-      int pre, n, post;
-      get_mid_dims(x_dim, y_dim, axis, &pre, &n, &post);
-      auto eigen_dout = framework::EigenTensor<T, 3>::From(
-          *dout, framework::make_ddim({pre, n, post}));
-      auto eigen_dy =
-          framework::EigenTensor<T, 1>::From(*dy, framework::make_ddim({n}));
-      eigen_dy.device(device) = eigen_dout.sum(
-          framework::EigenDim<2>::From(framework::make_ddim({0, 2})));
+      default_elementwise_add_grad<DeviceContext, T>(ctx, x, y, out, dout, dx,
+                                                     dy);
     }
   }
 };