/* Copyright (c) 2021 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/fluid/operators/layer_norm_op.h"
#include "paddle/fluid/operators/npu_op_runner.h"

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
using DDim = framework::DDim;

using DataLayout = framework::DataLayout;

template <typename T>
class NormDataType;

template <>
class NormDataType<platform::float16> {
 public:
  // The scaling param type is float for HALF and FLOAT tensors
  using ScalingParamType = const float;
  using BatchNormParamType = float;
};

template <>
class NormDataType<float> {
 public:
  using ScalingParamType = const float;
  using BatchNormParamType = float;
};

template <typename T>
using NormDataType = NormDataType<T>;
template <typename T>
using LayerNormParamType = typename NormDataType<T>::BatchNormParamType;

template <typename T>
class LayerNormNPUKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
    using U = LayerNormParamType<T>;
    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* 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_dims = x->dims();
    std::vector<int> axes;
    auto matrix_dim = framework::flatten_to_2d(x_dims, begin_norm_axis);
    int right = static_cast<int>(matrix_dim[1]);

    // The shape of scale and bias should be equal to x.shape[begin_norm_axis:],
    // required by Ascend.
    for (auto i = begin_norm_axis; i < x_dims.size(); ++i) {
      axes.push_back(x_dims[i]);
    }

    auto place = ctx.GetPlace();
    auto stream =
        ctx.template device_context<paddle::platform::NPUDeviceContext>()
            .stream();

    Tensor default_scale(x->type());
    if (!scale) {
      default_scale.mutable_data<T>(framework::make_ddim(axes), place);
      Tensor value(x->type());
      value.mutable_data<T>({1}, place);
      TensorFromVector(std::vector<T>{static_cast<T>(1.0)},
                       ctx.device_context(), &value);
      auto runner =
          NpuOpRunner("FillD", {value}, {default_scale}, {{"dims", axes}});
      runner.Run(stream);
      scale = &default_scale;
    } else {
      const_cast<Tensor*>(scale)->Resize(framework::make_ddim(axes));
    }

    Tensor default_bias(x->type());
    if (!bias) {
      default_bias.mutable_data<T>(framework::make_ddim(axes), place);
      Tensor value(x->type());
      value.mutable_data<T>({1}, place);
      TensorFromVector(std::vector<T>{static_cast<T>(0)}, ctx.device_context(),
                       &value);
      auto runner =
          NpuOpRunner("FillD", {value}, {default_bias}, {{"dims", axes}});
      runner.Run(stream);
      bias = &default_bias;
    } else {
      const_cast<Tensor*>(bias)->Resize(framework::make_ddim(axes));
    }

    // cast scale from LayerNormParamType to T if needed
    Tensor cast_scale(x->type());
    if (x->type() == framework::proto::VarType::FP16 &&
        scale->type() == framework::proto::VarType::FP32) {
      cast_scale.Resize(scale->dims());
      cast_scale.mutable_data<T>(ctx.GetPlace());
      auto dst_dtype = ConvertToNpuDtype(x->type());
      auto runner_cast_scale =
          NpuOpRunner("Cast", {*scale}, {cast_scale},
                      {{"dst_type", static_cast<int>(dst_dtype)}});
      runner_cast_scale.Run(stream);
    } else {
      cast_scale.ShareDataWith(*scale);
    }

    // cast bias from LayerNormParamType to T if needed
    Tensor cast_bias(x->type());
    if (x->type() == framework::proto::VarType::FP16 &&
        bias->type() == framework::proto::VarType::FP32) {
      cast_bias.Resize(bias->dims());
      cast_bias.mutable_data<T>(ctx.GetPlace());
      auto dst_dtype = ConvertToNpuDtype(x->type());
      auto runner_cast_bias =
          NpuOpRunner("Cast", {*bias}, {cast_bias},
                      {{"dst_type", static_cast<int>(dst_dtype)}});
      runner_cast_bias.Run(stream);
    } else {
      cast_bias.ShareDataWith(*bias);
    }

    y->mutable_data<T>(ctx.GetPlace());

    // mean should be of  U type
    Tensor* tmp_mean = mean;
    Tensor cast_mean(x->type());
    if (x->type() == framework::proto::VarType::FP16 &&
        (scale->type() == framework::proto::VarType::FP32 ||
         bias->type() == framework::proto::VarType::FP32)) {
      cast_mean.Resize(mean->dims());
      cast_mean.mutable_data<T>(ctx.GetPlace());
      tmp_mean = &cast_mean;
      mean->mutable_data<U>(ctx.GetPlace());
    } else {
      mean->mutable_data<T>(ctx.GetPlace());
    }

    // same for variance
    Tensor* tmp_variance = variance;
    Tensor cast_variance(x->type());
    if (x->type() == framework::proto::VarType::FP16 &&
        (scale->type() == framework::proto::VarType::FP32 ||
         bias->type() == framework::proto::VarType::FP32)) {
      cast_variance.Resize(variance->dims());
      cast_variance.mutable_data<T>(ctx.GetPlace());
      tmp_variance = &cast_variance;
      variance->mutable_data<U>(ctx.GetPlace());
    } else {
      variance->mutable_data<T>(ctx.GetPlace());
    }

    auto runner = NpuOpRunner("LayerNorm", {*x, cast_scale, cast_bias},
                              {*y, *tmp_mean, *tmp_variance},
                              {{"begin_norm_axis", begin_norm_axis},
                               {"begin_params_axis", begin_norm_axis},
                               {"epsilon", epsilon}});
    runner.Run(stream);

    // cast back from FP16 to FP32
    if (x->type() == framework::proto::VarType::FP16 &&
        mean->type() == framework::proto::VarType::FP32) {
      auto dst_dtype = ConvertToNpuDtype(mean->type());
      auto runner_cast_mean =
          NpuOpRunner("Cast", {*tmp_mean}, {*mean},
                      {{"dst_type", static_cast<int>(dst_dtype)}});
      runner_cast_mean.Run(stream);
    }
    // same for variance
    if (x->type() == framework::proto::VarType::FP16 &&
        variance->type() == framework::proto::VarType::FP32) {
      auto dst_dtype = ConvertToNpuDtype(variance->type());
      auto runner_cast_variance =
          NpuOpRunner("Cast", {*tmp_variance}, {*variance},
                      {{"dst_type", static_cast<int>(dst_dtype)}});
      runner_cast_variance.Run(stream);
    }

    // revert shape of scale and bias
    // TODO(zhiqiu): better implementation, use tmp tensor to avoid write input
    // tensor.
    const_cast<Tensor*>(scale)->Resize(framework::make_ddim({right}));
    const_cast<Tensor*>(bias)->Resize(framework::make_ddim({right}));
  }
};

template <typename T>
class LayerNormGradNPUKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
    using U = LayerNormParamType<T>;
    const auto begin_norm_axis = ctx.Attr<int>("begin_norm_axis");
    const auto* x = ctx.Input<Tensor>("X");
    const auto& x_dims = x->dims();
    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"));

    auto matrix_dim = framework::flatten_to_2d(x_dims, begin_norm_axis);
    int right = static_cast<int>(matrix_dim[1]);

    std::vector<int> axes;
    for (auto i = begin_norm_axis; i < x_dims.size(); ++i) {
      axes.push_back(x_dims[i]);
    }

    auto place = ctx.GetPlace();
    auto stream =
        ctx.template device_context<paddle::platform::NPUDeviceContext>()
            .stream();

    // No need to compute any gradient, jusr return
    if (!dx && !dscale && !dbias) {
      return;
    }

    // The rank of mean should be equal to x, required by Ascend.
    std::vector<int> new_shape;
    for (auto i = 0; i < begin_norm_axis; ++i) {
      new_shape.push_back(x_dims[i]);
    }
    for (auto i = begin_norm_axis; i < x_dims.size(); ++i) {
      new_shape.push_back(1);
    }

    auto mean_dims = mean->dims();
    const_cast<Tensor*>(mean)->Resize(framework::make_ddim({new_shape}));
    const_cast<Tensor*>(variance)->Resize(framework::make_ddim({new_shape}));

    Tensor default_scale(x->type());
    if (!scale) {
      default_scale.mutable_data<T>(framework::make_ddim(axes), place);
      Tensor value(x->type());
      value.mutable_data<T>({1}, place);
      TensorFromVector(std::vector<T>{static_cast<T>(1.0)},
                       ctx.device_context(), &value);
      auto runner =
          NpuOpRunner("FillD", {value}, {default_scale}, {{"dims", axes}});
      runner.Run(stream);
      scale = &default_scale;
    } else {
      const_cast<Tensor*>(scale)->Resize(framework::make_ddim(axes));
    }

    // cast scale from LayerNormParamType to T if needed
    Tensor cast_scale(x->type());
    if (x->type() == framework::proto::VarType::FP16 &&
        scale->type() == framework::proto::VarType::FP32) {
      cast_scale.Resize(scale->dims());
      cast_scale.mutable_data<T>(ctx.GetPlace());
      auto dst_dtype = ConvertToNpuDtype(x->type());
      auto runner_cast_scale =
          NpuOpRunner("Cast", {*scale}, {cast_scale},
                      {{"dst_type", static_cast<int>(dst_dtype)}});
      runner_cast_scale.Run(stream);
    } else {
      cast_scale.ShareDataWith(*scale);
    }

    // cast mean from LayerNormParamType to T if needed
    Tensor cast_mean(x->type());
    if (x->type() == framework::proto::VarType::FP16 &&
        mean->type() == framework::proto::VarType::FP32) {
      cast_mean.Resize(mean->dims());
      cast_mean.mutable_data<T>(ctx.GetPlace());
      auto dst_dtype = ConvertToNpuDtype(x->type());
      auto runner_cast_mean =
          NpuOpRunner("Cast", {*mean}, {cast_mean},
                      {{"dst_type", static_cast<int>(dst_dtype)}});
      runner_cast_mean.Run(stream);
    } else {
      cast_mean.ShareDataWith(*mean);
    }

    // cast variance from LayerNormParamType to T if needed
    Tensor cast_variance(x->type());
    if (x->type() == framework::proto::VarType::FP16 &&
        variance->type() == framework::proto::VarType::FP32) {
      cast_variance.Resize(variance->dims());
      cast_variance.mutable_data<T>(ctx.GetPlace());
      auto dst_dtype = ConvertToNpuDtype(x->type());
      auto runner_cast_variance =
          NpuOpRunner("Cast", {*variance}, {cast_variance},
                      {{"dst_type", static_cast<int>(dst_dtype)}});
      runner_cast_variance.Run(stream);
    } else {
      cast_variance.ShareDataWith(*variance);
    }

    Tensor dx_(dy->type()), dscale_(dy->type()), dbias_(dy->type());
    dx = (dx == nullptr) ? &dx_ : dx;
    dscale = (dscale == nullptr) ? &dscale_ : dscale;
    dbias = (dbias == nullptr) ? &dbias_ : dbias;

    dx->Resize(x->dims());
    dx->mutable_data<T>(ctx.GetPlace());

    dscale->Resize(framework::make_ddim(axes));

    dbias->Resize(framework::make_ddim(axes));

    // dscale should be of  U type
    Tensor* tmp_dscale = dscale;
    Tensor cast_dscale(x->type());
    if (x->type() == framework::proto::VarType::FP16 &&
        (mean->type() == framework::proto::VarType::FP32 ||
         variance->type() == framework::proto::VarType::FP32)) {
      cast_dscale.Resize(dscale->dims());
      cast_dscale.mutable_data<T>(ctx.GetPlace());
      tmp_dscale = &cast_dscale;
      dscale->mutable_data<U>(ctx.GetPlace());
    } else {
      dscale->mutable_data<T>(ctx.GetPlace());
    }

    // same for dbias
    Tensor* tmp_dbias = dbias;
    Tensor cast_dbias(x->type());
    if (x->type() == framework::proto::VarType::FP16 &&
        (mean->type() == framework::proto::VarType::FP32 ||
         variance->type() == framework::proto::VarType::FP32)) {
      cast_dbias.Resize(dbias->dims());
      cast_dbias.mutable_data<T>(ctx.GetPlace());
      tmp_dbias = &cast_dbias;
      dbias->mutable_data<U>(ctx.GetPlace());
    } else {
      dbias->mutable_data<T>(ctx.GetPlace());
    }

    auto runner = NpuOpRunner("LayerNormGrad",
                              {*dy, *x, cast_variance, cast_mean, cast_scale},
                              {*dx, *tmp_dscale, *tmp_dbias}, {});
    runner.Run(stream);

    // cast back from FP16 to FP32
    if (x->type() == framework::proto::VarType::FP16 &&
        dscale->type() == framework::proto::VarType::FP32) {
      auto dst_dtype = ConvertToNpuDtype(dscale->type());
      auto runner_cast_dscale =
          NpuOpRunner("Cast", {*tmp_dscale}, {*dscale},
                      {{"dst_type", static_cast<int>(dst_dtype)}});
      runner_cast_dscale.Run(stream);
    }
    // same for dbias
    if (x->type() == framework::proto::VarType::FP16 &&
        dbias->type() == framework::proto::VarType::FP32) {
      auto dst_dtype = ConvertToNpuDtype(dbias->type());
      auto runner_cast_dbias =
          NpuOpRunner("Cast", {*tmp_dbias}, {*dbias},
                      {{"dst_type", static_cast<int>(dst_dtype)}});
      runner_cast_dbias.Run(stream);
    }

    const_cast<Tensor*>(mean)->Resize(mean_dims);
    const_cast<Tensor*>(variance)->Resize(mean_dims);
    const_cast<Tensor*>(scale)->Resize(framework::make_ddim({right}));
    dscale->Resize(framework::make_ddim({right}));
    dbias->Resize(framework::make_ddim({right}));
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
namespace plat = paddle::platform;

REGISTER_OP_NPU_KERNEL(layer_norm, ops::LayerNormNPUKernel<float>,
                       ops::LayerNormNPUKernel<plat::float16>);
REGISTER_OP_NPU_KERNEL(layer_norm_grad, ops::LayerNormGradNPUKernel<float>,
                       ops::LayerNormGradNPUKernel<plat::float16>);