[PHI] move layer_norm/layer_norm_grad xpu kernel to phi (#45524)

* move layer_norm xpu kernel to phi, test=kunlun

* fix, test=kunlun

paddle/fluid/operators/layer_norm_op_xpu.cc

-/* Copyright (c) 2022 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. */
-
-#ifdef PADDLE_WITH_XPU
-
-#include "paddle/fluid/framework/op_registry.h"
-#include "paddle/fluid/platform/device/device_wrapper.h"
-
-namespace paddle {
-namespace operators {
-
-using Tensor = framework::Tensor;
-using DDim = framework::DDim;
-
-template <typename DeviceContext, typename T>
-class LayerNormXPUKernel : public framework::OpKernel<T> {
-  using XPUType = typename XPUTypeTrait<T>::Type;
-
- public:
-  void Compute(const framework::ExecutionContext& ctx) const override {
-    const auto begin_norm_axis = ctx.Attr<int>("begin_norm_axis");
-    const auto epsilon = ctx.Attr<float>("epsilon");
-    const auto* x = ctx.Input<Tensor>("X");
-    const auto& x_dims = x->dims();
-    auto matrix_dim = phi::flatten_to_2d(x_dims, begin_norm_axis);
-    int left = static_cast<int>(matrix_dim[0]);
-    int right = static_cast<int>(matrix_dim[1]);
-    const auto* scale = ctx.Input<Tensor>("Scale");
-    const auto* bias = ctx.Input<Tensor>("Bias");
-    auto* y = ctx.Output<Tensor>("Y");
-    auto* mean = ctx.Output<Tensor>("Mean");
-    auto* variance = ctx.Output<Tensor>("Variance");
-    const auto* x_data = x->data<T>();
-    const auto* scale_data =
-        (scale == nullptr ? nullptr : scale->data<float>());
-    const auto* bias_data = (bias == nullptr ? nullptr : bias->data<float>());
-    auto* y_data = y->mutable_data<T>(ctx.GetPlace());
-    auto* mean_data = mean->mutable_data<float>(ctx.GetPlace());
-    auto* variance_data = variance->mutable_data<float>(ctx.GetPlace());
-    auto& dev_ctx = ctx.template device_context<DeviceContext>();
-
-    // int layer_norm(Context* ctx, const T* x, T* y, int m, int n, float eps,
-    // const float* scale, const float* bias, float* mean, float* var);
-    int r = xpu::layer_norm(dev_ctx.x_context(),
-                            reinterpret_cast<const XPUType*>(x_data),
-                            reinterpret_cast<XPUType*>(y_data),
-                            left,
-                            right,
-                            epsilon,
-                            scale_data,
-                            bias_data,
-                            mean_data,
-                            variance_data);
-    PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm");
-  }
-};
-
-template <typename DeviceContext, typename T>
-class LayerNormGradXPUKernel : public framework::OpKernel<T> {
-  using XPUType = typename XPUTypeTrait<T>::Type;
-
- public:
-  void Compute(const framework::ExecutionContext& ctx) const override {
-    const auto begin_norm_axis = ctx.Attr<int>("begin_norm_axis");
-    const auto epsilon = ctx.Attr<float>("epsilon");
-    const auto* x = ctx.Input<Tensor>("X");
-    const auto& x_dims = x->dims();
-    auto matrix_dim = phi::flatten_to_2d(x_dims, begin_norm_axis);
-    int left = static_cast<int>(matrix_dim[0]);
-    int right = static_cast<int>(matrix_dim[1]);
-    const auto* mean = ctx.Input<Tensor>("Mean");
-    const auto* variance = ctx.Input<Tensor>("Variance");
-    const auto* scale = ctx.Input<Tensor>("Scale");
-    const auto* dy = ctx.Input<Tensor>(framework::GradVarName("Y"));
-    auto* dx = ctx.Output<Tensor>(framework::GradVarName("X"));
-    auto* dscale = ctx.Output<Tensor>(framework::GradVarName("Scale"));
-    auto* dbias = ctx.Output<Tensor>(framework::GradVarName("Bias"));
-    const auto* x_data = x->data<T>();
-    const auto* dy_data = dy->data<T>();
-    const auto* mean_data = mean->data<float>();
-    const auto* variance_data = variance->data<float>();
-    const auto* scale_data =
-        (scale == nullptr ? nullptr : scale->data<float>());
-    auto* dscale_data =
-        (dscale == nullptr ? nullptr
-                           : dscale->mutable_data<float>(ctx.GetPlace()));
-    auto* dbias_data =
-        (dbias == nullptr ? nullptr
-                          : dbias->mutable_data<float>(ctx.GetPlace()));
-    auto* dx_data =
-        (dx == nullptr ? nullptr : dx->mutable_data<T>(ctx.GetPlace()));
-    auto& dev_ctx = ctx.template device_context<DeviceContext>();
-
-    // int layer_norm_grad(Context* ctx, const T* x, const T* dy, T* dx, int m,
-    // int n, float eps, const float* scale, const float* mean, const float*
-    // var, float* dscale, float* dbias);
-    int r = xpu::layer_norm_grad(dev_ctx.x_context(),
-                                 reinterpret_cast<const XPUType*>(x_data),
-                                 reinterpret_cast<const XPUType*>(dy_data),
-                                 reinterpret_cast<XPUType*>(dx_data),
-                                 left,
-                                 right,
-                                 epsilon,
-                                 scale_data,
-                                 mean_data,
-                                 variance_data,
-                                 dscale_data,
-                                 dbias_data);
-    PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm_grad");
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-
-REGISTER_OP_XPU_KERNEL(
-    layer_norm,
-    ops::LayerNormXPUKernel<paddle::platform::XPUDeviceContext, float>,
-    ops::LayerNormXPUKernel<paddle::platform::XPUDeviceContext,
-                            paddle::platform::float16>);
-REGISTER_OP_XPU_KERNEL(
-    layer_norm_grad,
-    ops::LayerNormGradXPUKernel<paddle::platform::XPUDeviceContext, float>,
-    ops::LayerNormGradXPUKernel<paddle::platform::XPUDeviceContext,
-                                paddle::platform::float16>);
-
-#endif  // PADDLE_WITH_XPU

paddle/phi/kernels/xpu/layer_norm_grad_kernel.cc

+// Copyright (c) 2022 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 "paddle/phi/kernels/layer_norm_grad_kernel.h"
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void LayerNormGradKernel(const Context& ctx,
+                         const DenseTensor& x,
+                         const paddle::optional<DenseTensor>& scale,
+                         const paddle::optional<DenseTensor>& bias,
+                         const DenseTensor& mean,
+                         const DenseTensor& variance,
+                         const DenseTensor& out_grad,
+                         float epsilon,
+                         int begin_norm_axis,
+                         bool is_test,
+                         DenseTensor* x_grad,
+                         DenseTensor* scale_grad,
+                         DenseTensor* bias_grad) {
+  using XPUType = typename XPUTypeTrait<T>::Type;
+  const auto& x_dims = x.dims();
+  auto matrix_dim = phi::flatten_to_2d(x_dims, begin_norm_axis);
+  int left = static_cast<int>(matrix_dim[0]);
+  int right = static_cast<int>(matrix_dim[1]);
+  const auto* x_data = x.data<T>();
+  const auto* out_grad_data = out_grad.data<T>();
+  const auto* mean_data = mean.data<float>();
+  const auto* variance_data = variance.data<float>();
+  const auto* scale_data =
+      (scale.get_ptr() == nullptr ? nullptr : scale.get_ptr()->data<float>());
+  auto* scale_grad_data =
+      (scale_grad == nullptr ? nullptr : ctx.template Alloc<float>(scale_grad));
+  auto* bias_grad_data =
+      (bias_grad == nullptr ? nullptr : ctx.template Alloc<float>(bias_grad));
+  auto* x_grad_data =
+      (x_grad == nullptr ? nullptr : ctx.template Alloc<T>(x_grad));
+
+  // int layer_norm_grad(Context* ctx, const T* x, const T* dy, T* dx, int m,
+  // int n, float eps, const float* scale, const float* mean, const float*
+  // var, float* dscale, float* dbias);
+  int r = xpu::layer_norm_grad(ctx.x_context(),
+                               reinterpret_cast<const XPUType*>(x_data),
+                               reinterpret_cast<const XPUType*>(out_grad_data),
+                               reinterpret_cast<XPUType*>(x_grad_data),
+                               left,
+                               right,
+                               epsilon,
+                               scale_data,
+                               mean_data,
+                               variance_data,
+                               scale_grad_data,
+                               bias_grad_data);
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm_grad");
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(layer_norm_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::LayerNormGradKernel,
+                   float,
+                   phi::dtype::float16) {}

paddle/phi/kernels/xpu/layer_norm_kernel.cc

+// Copyright (c) 2022 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 "paddle/phi/kernels/layer_norm_kernel.h"
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void LayerNormKernel(const Context& ctx,
+                     const DenseTensor& x,
+                     const paddle::optional<DenseTensor>& scale,
+                     const paddle::optional<DenseTensor>& bias,
+                     float epsilon,
+                     int begin_norm_axis,
+                     bool is_test,
+                     DenseTensor* out,
+                     DenseTensor* mean,
+                     DenseTensor* variance) {
+  using XPUType = typename XPUTypeTrait<T>::Type;
+  const auto& x_dims = x.dims();
+  auto matrix_dim = phi::flatten_to_2d(x_dims, begin_norm_axis);
+  int left = static_cast<int>(matrix_dim[0]);
+  int right = static_cast<int>(matrix_dim[1]);
+  const auto* x_data = x.data<T>();
+  const auto* scale_data =
+      (scale.get_ptr() == nullptr ? nullptr : scale.get_ptr()->data<float>());
+  const auto* bias_data =
+      (bias.get_ptr() == nullptr ? nullptr : bias.get_ptr()->data<float>());
+  auto* out_data = ctx.template Alloc<T>(out);
+  auto* mean_data = ctx.template Alloc<float>(mean);
+  auto* variance_data = ctx.template Alloc<float>(variance);
+
+  // int layer_norm(Context* ctx, const T* x, T* y, int m, int n, float eps,
+  // const float* scale, const float* bias, float* mean, float* var);
+  int r = xpu::layer_norm(ctx.x_context(),
+                          reinterpret_cast<const XPUType*>(x_data),
+                          reinterpret_cast<XPUType*>(out_data),
+                          left,
+                          right,
+                          epsilon,
+                          scale_data,
+                          bias_data,
+                          mean_data,
+                          variance_data);
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm");
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(layer_norm,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::LayerNormKernel,
+                   float,
+                   phi::dtype::float16) {}