code refine

paddle/operators/layer_norm_op.h

@@ -97,15 +97,15 @@ class LayerNormKernel : public framework::OpKernel<T> {
     auto &dev_ctx = ctx.template device_context<DeviceContext>();
     math::RowwiseMean<DeviceContext, T> row_mean;
 
-    // functor-> get mean
+    // get mean
     row_mean(dev_ctx, x, mean);
 
-    // functor-> get variance
+    // get variance
     ElementwiseComputeEx<SubAndSquareFunctor<T>, DeviceContext, T>(
         ctx, &x, mean, /*axis*/ 0, SubAndSquareFunctor<T>(), &out);
     row_mean(dev_ctx, out, var);
 
-    // functor-> get norm_out
+    // get x_norm
     ElementwiseComputeEx<SubFunctor<T>, DeviceContext, T>(
         ctx, &x, mean, /*axis*/ 0, SubFunctor<T>(), &out);
     ElementwiseComputeEx<DivAndSqrtFunctor<T>, DeviceContext, T>(
@@ -129,9 +129,11 @@ class LayerNormGradKernel : public framework::OpKernel<T> {
   void Compute(const framework::ExecutionContext &ctx) const override {
     const float epsilon = ctx.Attr<float>("epsilon");
     auto x = *ctx.Input<Tensor>("X");
-    auto mean = *ctx.Input<Tensor>("Mean");
-    auto var = *ctx.Input<Tensor>("Variance");
-    auto scale = *ctx.Input<Tensor>("Scale");
+    auto *y = ctx.Input<Tensor>("Y");
+    auto *mean = ctx.Input<Tensor>("Mean");
+    auto *var = ctx.Input<Tensor>("Variance");
+    auto *scale = ctx.Input<Tensor>("Scale");
+    auto *bias = ctx.Input<Tensor>("Bias");
     auto d_y = *ctx.Input<Tensor>(framework::GradVarName("Y"));
     const auto begin_norm_axis = ctx.Attr<int>("begin_norm_axis");
 
@@ -155,14 +157,19 @@ class LayerNormGradKernel : public framework::OpKernel<T> {
     if (d_scale || d_x) {
       x.Resize(matrix_shape);
       temp.mutable_data<T>(matrix_shape, ctx.GetPlace());
-      temp_norm.mutable_data<T>(matrix_shape, ctx.GetPlace());
 
-      // get x_norm
-      ElementwiseComputeEx<SubFunctor<T>, DeviceContext, T>(
-          ctx, &x, &mean, /*axis*/ 0, SubFunctor<T>(), &temp_norm);
-      ElementwiseComputeEx<DivAndSqrtFunctor<T>, DeviceContext, T>(
-          ctx, &temp_norm, &var, /*axis*/ 0,
-          DivAndSqrtFunctor<T>(static_cast<T>(epsilon)), &temp_norm);
+      if (!(bias && scale)) {
+        temp_norm.ShareDataWith(*y);
+        temp_norm.Resize(matrix_shape);
+      } else {
+        temp_norm.mutable_data<T>(matrix_shape, ctx.GetPlace());
+        // get x_norm
+        ElementwiseComputeEx<SubFunctor<T>, DeviceContext, T>(
+            ctx, &x, mean, /*axis*/ 0, SubFunctor<T>(), &temp_norm);
+        ElementwiseComputeEx<DivAndSqrtFunctor<T>, DeviceContext, T>(
+            ctx, &temp_norm, var, /*axis*/ 0,
+            DivAndSqrtFunctor<T>(static_cast<T>(epsilon)), &temp_norm);
+      }
     }
 
     if (d_bias) {
@@ -188,7 +195,7 @@ class LayerNormGradKernel : public framework::OpKernel<T> {
       if (d_scale) {
         // dy_dx
         ElementwiseComputeEx<MulFunctor<T>, DeviceContext, T>(
-            ctx, &d_y, &scale, /*axis*/ 1, MulFunctor<T>(), &temp);
+            ctx, &d_y, scale, /*axis*/ 1, MulFunctor<T>(), &temp);
         framework::Copy(temp, ctx.GetPlace(), ctx.device_context(), d_x);
 
         // dy_dmean_dx
@@ -199,7 +206,6 @@ class LayerNormGradKernel : public framework::OpKernel<T> {
         // dy_var_dx
         ElementwiseComputeEx<MulFunctor<T>, DeviceContext, T>(
             ctx, &temp, &temp_norm, /*axis*/ 0, MulFunctor<T>(), &temp);
-
       } else {
         // dy_dx
         framework::Copy(d_y, ctx.GetPlace(), ctx.device_context(), d_x);
@@ -216,12 +222,12 @@ class LayerNormGradKernel : public framework::OpKernel<T> {
       // dy_var_dx
       row_mean(dev_ctx, temp, &temp_vec);
       ElementwiseComputeEx<MulFunctor<T>, DeviceContext, T>(
-          ctx, &temp_norm, &temp_vec, /*axis*/ 0, MulFunctor<T>(), &temp_norm);
+          ctx, &temp_norm, &temp_vec, /*axis*/ 0, MulFunctor<T>(), &temp);
       ElementwiseComputeEx<SubFunctor<T>, DeviceContext, T>(
-          ctx, d_x, &temp_norm, /*axis*/ 0, SubFunctor<T>(), d_x);
+          ctx, d_x, &temp, /*axis*/ 0, SubFunctor<T>(), d_x);
 
       ElementwiseComputeEx<DivAndSqrtFunctor<T>, DeviceContext, T>(
-          ctx, d_x, &var, /*axis*/ 0,
+          ctx, d_x, var, /*axis*/ 0,
           DivAndSqrtFunctor<T>(static_cast<T>(epsilon)), d_x);
       d_x->Resize(dx_dim);
     }