// 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/common/flags.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/cast_kernel.h"
#include "paddle/phi/kernels/full_kernel.h"
#include "paddle/phi/kernels/funcs/layer_norm_impl.cu.h"
#include "paddle/phi/kernels/funcs/layer_norm_util.h"
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && !defined(_WIN32)
#include "paddle/phi/kernels/funcs/fast_ln_v2.h"
#endif
#ifdef PADDLE_WITH_CUDA
#include "paddle/phi/kernels/gpu/rms_norm_cuda_kernel.h"
#endif
COMMON_DECLARE_bool(use_apex_layer_norm_kernel);

namespace phi {
enum class LayerNormGadKernelVariant { FAST_LN_V2, GENERIC };
static inline LayerNormGadKernelVariant LayerNormGradKernelDispatch(
    const DataType weight_type,
    const DataType input_type,
    const DataType output_type,
    const DataType compute_type,
    const uint32_t hidden_size,
    const int64_t x_numel,
    const DenseTensor* scale,
    const DenseTensor* bias) {
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && !defined(_WIN32)
  if (FLAGS_use_apex_layer_norm_kernel) {
    if (funcs::fast_ln_v2::has_fast_ln_v2_bwd_kernel(
            weight_type, input_type, output_type, compute_type, hidden_size)) {
      return LayerNormGadKernelVariant::FAST_LN_V2;
    }
    PADDLE_THROW(common::errors::InvalidArgument(
        "FLAGS_use_apex_layer_norm_kernel requires inputs supported by "
        "fast_ln_v2 backward kernel."));
  }
  if (FLAGS_use_accuracy_compatible_kernel) {
    return LayerNormGadKernelVariant::GENERIC;
  }
  if (scale != nullptr && bias != nullptr && input_type != DataType::FLOAT32 &&
      hidden_size != 4096 && hidden_size > 1024 && hidden_size <= 10240 &&
      x_numel <= std::numeric_limits<uint32_t>::max()) {
    // using fast_ln_v2 only sm > 70 and x_numel <= uint32_max
    auto prop = funcs::fast_ln_v2::GetDeviceProp();
    if (prop->major > 7 &&
        funcs::fast_ln_v2::has_fast_ln_v2_bwd_kernel(
            weight_type, input_type, output_type, compute_type, hidden_size)) {
      return LayerNormGadKernelVariant::FAST_LN_V2;
    }
  }
#endif
  return LayerNormGadKernelVariant::GENERIC;
}

template <typename T, typename Context>
void LayerNormGradKernel(const Context& dev_ctx,
                         const DenseTensor& x,
                         const optional<DenseTensor>& scale_opt,
                         const optional<DenseTensor>& bias_opt,
                         const DenseTensor& mean,
                         const DenseTensor& variance,
                         const DenseTensor& out_grad,
                         double epsilon,
                         int begin_norm_axis,
                         DenseTensor* x_grad,
                         DenseTensor* scale_grad,
                         DenseTensor* bias_grad) {
  if (x.numel() == 0) {
    dev_ctx.template Alloc<T>(x_grad);
    if (scale_grad) {
      Full<T, Context>(dev_ctx, scale_grad->dims(), 0, scale_grad);
      if (scale_opt.get_ptr() && x.dtype() != scale_opt.get().dtype()) {
        CastKernel<T, Context>(
            dev_ctx, *scale_grad, scale_opt.get().dtype(), scale_grad);
      }
    }
    if (bias_grad) {
      Full<T, Context>(dev_ctx, bias_grad->dims(), 0, bias_grad);
      if (bias_opt.get_ptr() && x.dtype() != bias_opt.get().dtype()) {
        CastKernel<T, Context>(
            dev_ctx, *bias_grad, bias_opt.get().dtype(), bias_grad);
      }
    }
    return;
  }
  using U = funcs::LayerNormParamType<T>;
  // d_x, d_scale, d_bias may be nullptr
  auto* d_x = x_grad;
  auto* d_scale = scale_grad;
  auto* d_bias = bias_grad;

  auto* scale = scale_opt.get_ptr();
  auto* bias = bias_opt.get_ptr();
  auto* d_y = &out_grad;

  const auto& x_dims = x.dims();
  auto matrix_dim = common::flatten_to_2d(x_dims, begin_norm_axis);
  int64_t batch_size = static_cast<int64_t>(matrix_dim[0]);
  int64_t feature_size = static_cast<int64_t>(matrix_dim[1]);

  auto* x_data = x.data<T>();
  auto* d_y_data = d_y->data<T>();

  auto* mean_data = mean.data<U>();
  auto* var_data = variance.data<U>();

  auto* d_x_data = (d_x == nullptr ? nullptr : dev_ctx.template Alloc<T>(d_x));

  auto x_dtype = x.dtype();

  DataType scale_bias_dtype;
  if (scale != nullptr) {
    scale_bias_dtype = scale->dtype();
  } else {
    // FIXME(zengjinle): do not find a better way to get the right
    // data type of the d_scale and d_bias if scale == nullptr.
    if (bias != nullptr) {
      scale_bias_dtype = bias->dtype();
    } else {
      scale_bias_dtype = x_dtype;
    }
  }

#define PADDLE_LAUNCH_LAYERNORM_BWD(ScaleBiasT, IsScaleBiasSameDTypeWithX)  \
  do {                                                                      \
    auto* scale_data =                                                      \
        (scale == nullptr ? nullptr : scale->data<ScaleBiasT>());           \
    auto* d_scale_data =                                                    \
        (d_scale == nullptr ? nullptr                                       \
                            : dev_ctx.template Alloc<ScaleBiasT>(d_scale)); \
    auto* d_bias_data =                                                     \
        (d_bias == nullptr ? nullptr                                        \
                           : dev_ctx.template Alloc<ScaleBiasT>(d_bias));   \
    auto* d_x_data =                                                        \
        (d_x == nullptr ? nullptr : dev_ctx.template Alloc<T>(d_x));        \
    funcs::LayerNormBackward<T, U, IsScaleBiasSameDTypeWithX>(x_data,       \
                                                              d_y_data,     \
                                                              scale_data,   \
                                                              mean_data,    \
                                                              var_data,     \
                                                              d_x_data,     \
                                                              d_scale_data, \
                                                              d_bias_data,  \
                                                              epsilon,      \
                                                              batch_size,   \
                                                              feature_size, \
                                                              dev_ctx);     \
  } while (0)

#define PADDLE_LAUNCH_FAST_LAYERNORM_V2_BWD(ScaleBiasT)                     \
  do {                                                                      \
    auto stream = dev_ctx.stream();                                         \
    auto place = x.place();                                                 \
    auto* scale_data =                                                      \
        (scale == nullptr ? nullptr : scale->data<ScaleBiasT>());           \
    auto* d_scale_data =                                                    \
        (d_scale == nullptr ? nullptr                                       \
                            : dev_ctx.template Alloc<ScaleBiasT>(d_scale)); \
    auto* d_bias_data =                                                     \
        (d_bias == nullptr ? nullptr                                        \
                           : dev_ctx.template Alloc<ScaleBiasT>(d_bias));   \
    auto* d_x_data =                                                        \
        (d_x == nullptr ? nullptr : dev_ctx.template Alloc<T>(d_x));        \
    funcs::fast_ln_v2::LaunchNormBwd<T, Context>(dev_ctx,                   \
                                                 stream,                    \
                                                 place,                     \
                                                 x_data,                    \
                                                 scale_data,                \
                                                 mean_data,                 \
                                                 var_data,                  \
                                                 d_y_data,                  \
                                                 d_x_data,                  \
                                                 d_scale_data,              \
                                                 d_bias_data,               \
                                                 scale_bias_dtype,          \
                                                 x_dtype,                   \
                                                 x_grad->dtype(),           \
                                                 compute_dtype,             \
                                                 feature_size,              \
                                                 batch_size,                \
                                                 feature_size,              \
                                                 epsilon);                  \
  } while (0)

  auto compute_dtype = CppTypeToDataType<U>::Type();
  auto kernel_variant = LayerNormGradKernelDispatch(scale_bias_dtype,
                                                    x_dtype,
                                                    x_dtype,
                                                    compute_dtype,
                                                    feature_size,
                                                    x.numel(),
                                                    scale,
                                                    bias);
  switch (kernel_variant) {
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && !defined(_WIN32)
    case LayerNormGadKernelVariant::FAST_LN_V2:
      if (scale_bias_dtype == x_dtype) {
        PADDLE_LAUNCH_FAST_LAYERNORM_V2_BWD(T);
      } else {
        PADDLE_LAUNCH_FAST_LAYERNORM_V2_BWD(U);
      }
      break;
#endif
    case LayerNormGadKernelVariant::GENERIC:
    default:
#ifdef PADDLE_WITH_CUDA
      if ((FLAGS_use_accuracy_compatible_kernel ||
           (!isPowerOfTwo(feature_size) && feature_size > 1024)) &&
          scale_bias_dtype == x_dtype) {
        auto* scale_data = (scale == nullptr ? nullptr : scale->data<T>());
        auto* d_scale_data =
            (d_scale == nullptr ? nullptr : dev_ctx.template Alloc<T>(d_scale));
        auto* d_bias_data =
            (d_bias == nullptr ? nullptr : dev_ctx.template Alloc<T>(d_bias));
        auto* d_x_data =
            (d_x == nullptr ? nullptr : dev_ctx.template Alloc<T>(d_x));
        LayerNormBwdCompatKernel<T, Context>(dev_ctx,
                                             d_y_data,
                                             x_data,
                                             scale_data,
                                             mean_data,
                                             var_data,
                                             d_x_data,
                                             d_scale_data,
                                             d_bias_data,
                                             epsilon,
                                             batch_size,
                                             feature_size);
      } else {
#endif
        if (scale_bias_dtype == x_dtype) {
          PADDLE_LAUNCH_LAYERNORM_BWD(T, true);
        } else {
          PADDLE_LAUNCH_LAYERNORM_BWD(U, false);
        }
#ifdef PADDLE_WITH_CUDA
      }
#endif
  }

#undef PADDLE_LAUNCH_LAYERNORM_BWD
#undef PADDLE_LAUNCH_FAST_LAYERNORM_V2_BWD
}

}  // namespace phi

#ifdef PADDLE_WITH_HIP
// MIOPEN do not support double
PD_REGISTER_KERNEL(layer_norm_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::LayerNormGradKernel,
                   float,
                   phi::float16) {
  if (kernel_key.dtype() == phi::DataType::FLOAT16) {
    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
  }
}
#elif CUDNN_VERSION_MIN(8, 1, 0)
PD_REGISTER_KERNEL(layer_norm_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::LayerNormGradKernel,
                   float,
                   double,
                   phi::float16,
                   phi::bfloat16) {
  if (kernel_key.dtype() == phi::DataType::FLOAT16) {
    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
  }
}
#else
PD_REGISTER_KERNEL(layer_norm_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::LayerNormGradKernel,
                   float,
                   double,
                   phi::float16) {
  if (kernel_key.dtype() == phi::DataType::FLOAT16) {
    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
  }
}
#endif