From 871e332964808fce1775e6090cf5c7c596934111 Mon Sep 17 00:00:00 2001 From: pangyoki Date: Tue, 30 Aug 2022 15:25:19 +0800 Subject: [PATCH] [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 | 140 ------------------ .../phi/kernels/xpu/layer_norm_grad_kernel.cc | 78 ++++++++++ paddle/phi/kernels/xpu/layer_norm_kernel.cc | 68 +++++++++ 3 files changed, 146 insertions(+), 140 deletions(-) delete mode 100644 paddle/fluid/operators/layer_norm_op_xpu.cc create mode 100644 paddle/phi/kernels/xpu/layer_norm_grad_kernel.cc create mode 100644 paddle/phi/kernels/xpu/layer_norm_kernel.cc diff --git a/paddle/fluid/operators/layer_norm_op_xpu.cc b/paddle/fluid/operators/layer_norm_op_xpu.cc deleted file mode 100644 index ddb580f1a1d..00000000000 --- a/paddle/fluid/operators/layer_norm_op_xpu.cc +++ /dev/null @@ -1,140 +0,0 @@ -/* 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 -class LayerNormXPUKernel : public framework::OpKernel { - using XPUType = typename XPUTypeTrait::Type; - - public: - void Compute(const framework::ExecutionContext& ctx) const override { - const auto begin_norm_axis = ctx.Attr("begin_norm_axis"); - const auto epsilon = ctx.Attr("epsilon"); - const auto* x = ctx.Input("X"); - const auto& x_dims = x->dims(); - auto matrix_dim = phi::flatten_to_2d(x_dims, begin_norm_axis); - int left = static_cast(matrix_dim[0]); - int right = static_cast(matrix_dim[1]); - const auto* scale = ctx.Input("Scale"); - const auto* bias = ctx.Input("Bias"); - auto* y = ctx.Output("Y"); - auto* mean = ctx.Output("Mean"); - auto* variance = ctx.Output("Variance"); - const auto* x_data = x->data(); - const auto* scale_data = - (scale == nullptr ? nullptr : scale->data()); - const auto* bias_data = (bias == nullptr ? nullptr : bias->data()); - auto* y_data = y->mutable_data(ctx.GetPlace()); - auto* mean_data = mean->mutable_data(ctx.GetPlace()); - auto* variance_data = variance->mutable_data(ctx.GetPlace()); - auto& dev_ctx = ctx.template device_context(); - - // 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(x_data), - reinterpret_cast(y_data), - left, - right, - epsilon, - scale_data, - bias_data, - mean_data, - variance_data); - PADDLE_ENFORCE_XDNN_SUCCESS(r, "layer_norm"); - } -}; - -template -class LayerNormGradXPUKernel : public framework::OpKernel { - using XPUType = typename XPUTypeTrait::Type; - - public: - void Compute(const framework::ExecutionContext& ctx) const override { - const auto begin_norm_axis = ctx.Attr("begin_norm_axis"); - const auto epsilon = ctx.Attr("epsilon"); - const auto* x = ctx.Input("X"); - const auto& x_dims = x->dims(); - auto matrix_dim = phi::flatten_to_2d(x_dims, begin_norm_axis); - int left = static_cast(matrix_dim[0]); - int right = static_cast(matrix_dim[1]); - const auto* mean = ctx.Input("Mean"); - const auto* variance = ctx.Input("Variance"); - const auto* scale = ctx.Input("Scale"); - const auto* dy = ctx.Input(framework::GradVarName("Y")); - auto* dx = ctx.Output(framework::GradVarName("X")); - auto* dscale = ctx.Output(framework::GradVarName("Scale")); - auto* dbias = ctx.Output(framework::GradVarName("Bias")); - const auto* x_data = x->data(); - const auto* dy_data = dy->data(); - const auto* mean_data = mean->data(); - const auto* variance_data = variance->data(); - const auto* scale_data = - (scale == nullptr ? nullptr : scale->data()); - auto* dscale_data = - (dscale == nullptr ? nullptr - : dscale->mutable_data(ctx.GetPlace())); - auto* dbias_data = - (dbias == nullptr ? nullptr - : dbias->mutable_data(ctx.GetPlace())); - auto* dx_data = - (dx == nullptr ? nullptr : dx->mutable_data(ctx.GetPlace())); - auto& dev_ctx = ctx.template device_context(); - - // 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(x_data), - reinterpret_cast(dy_data), - reinterpret_cast(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, - ops::LayerNormXPUKernel); -REGISTER_OP_XPU_KERNEL( - layer_norm_grad, - ops::LayerNormGradXPUKernel, - ops::LayerNormGradXPUKernel); - -#endif // PADDLE_WITH_XPU diff --git a/paddle/phi/kernels/xpu/layer_norm_grad_kernel.cc b/paddle/phi/kernels/xpu/layer_norm_grad_kernel.cc new file mode 100644 index 00000000000..d7d8e9058e8 --- /dev/null +++ b/paddle/phi/kernels/xpu/layer_norm_grad_kernel.cc @@ -0,0 +1,78 @@ +// 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 +void LayerNormGradKernel(const Context& ctx, + const DenseTensor& x, + const paddle::optional& scale, + const paddle::optional& 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::Type; + const auto& x_dims = x.dims(); + auto matrix_dim = phi::flatten_to_2d(x_dims, begin_norm_axis); + int left = static_cast(matrix_dim[0]); + int right = static_cast(matrix_dim[1]); + const auto* x_data = x.data(); + const auto* out_grad_data = out_grad.data(); + const auto* mean_data = mean.data(); + const auto* variance_data = variance.data(); + const auto* scale_data = + (scale.get_ptr() == nullptr ? nullptr : scale.get_ptr()->data()); + auto* scale_grad_data = + (scale_grad == nullptr ? nullptr : ctx.template Alloc(scale_grad)); + auto* bias_grad_data = + (bias_grad == nullptr ? nullptr : ctx.template Alloc(bias_grad)); + auto* x_grad_data = + (x_grad == nullptr ? nullptr : ctx.template Alloc(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(x_data), + reinterpret_cast(out_grad_data), + reinterpret_cast(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) {} diff --git a/paddle/phi/kernels/xpu/layer_norm_kernel.cc b/paddle/phi/kernels/xpu/layer_norm_kernel.cc new file mode 100644 index 00000000000..e1895bd0a74 --- /dev/null +++ b/paddle/phi/kernels/xpu/layer_norm_kernel.cc @@ -0,0 +1,68 @@ +// 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 +void LayerNormKernel(const Context& ctx, + const DenseTensor& x, + const paddle::optional& scale, + const paddle::optional& bias, + float epsilon, + int begin_norm_axis, + bool is_test, + DenseTensor* out, + DenseTensor* mean, + DenseTensor* variance) { + using XPUType = typename XPUTypeTrait::Type; + const auto& x_dims = x.dims(); + auto matrix_dim = phi::flatten_to_2d(x_dims, begin_norm_axis); + int left = static_cast(matrix_dim[0]); + int right = static_cast(matrix_dim[1]); + const auto* x_data = x.data(); + const auto* scale_data = + (scale.get_ptr() == nullptr ? nullptr : scale.get_ptr()->data()); + const auto* bias_data = + (bias.get_ptr() == nullptr ? nullptr : bias.get_ptr()->data()); + auto* out_data = ctx.template Alloc(out); + auto* mean_data = ctx.template Alloc(mean); + auto* variance_data = ctx.template Alloc(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(x_data), + reinterpret_cast(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) {} -- GitLab