Merge pull request #12554 from sneaxiy/refine_elementwise_add

Refine elementwise_add op

paddle/fluid/operators/elementwise_add_op.cu

@@ -16,6 +16,60 @@ 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,16 +144,41 @@ 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 (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);
+    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);
     } else {
-      default_elementwise_add_grad<DeviceContext, T>(ctx, x, y, out, dout, dx,
-                                                     dy);
+      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})));
     }
   }
 };