// Copyright (c) 2024 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/backends/xpu/enforce_xpu.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/utils/optional.h"

namespace phi {

template <typename T, typename Context>
void ResNetUnitXPUKernel(const Context &dev_ctx,
                         const DenseTensor &x_in,
                         const DenseTensor &filter_x_in,
                         const DenseTensor &scale_x_in,
                         const DenseTensor &bias_x_in,
                         const DenseTensor &mean_x_in,
                         const DenseTensor &var_x_in,
                         const optional<DenseTensor> &z_in,
                         const optional<DenseTensor> &filter_z_in,
                         const optional<DenseTensor> &scale_z_in,
                         const optional<DenseTensor> &bias_z_in,
                         const optional<DenseTensor> &mean_z_in,
                         const optional<DenseTensor> &var_z_in,
                         int stride,
                         int stride_z,
                         int padding,
                         int dilation,
                         int group,
                         float momentum_in,
                         float epsilon,
                         const std::string &data_format,
                         bool fuse_add,
                         bool has_shortcut,
                         bool use_global_stats,
                         bool is_test,
                         bool use_addto,
                         const std::string &act_type,
                         DenseTensor *out,
                         DenseTensor *bit_mask,
                         DenseTensor *conv_x,
                         DenseTensor *saved_mean_x,
                         DenseTensor *saved_invstd_x,
                         DenseTensor *running_mean_x,
                         DenseTensor *running_var_x,
                         DenseTensor *conv_z,
                         DenseTensor *saved_mean_z,
                         DenseTensor *saved_invstd_z,
                         DenseTensor *running_mean_z,
                         DenseTensor *running_var_z) {
  using XPUType = typename XPUTypeTrait<T>::Type;

  bool is_nchw = (data_format == "NCHW");
  // input x
  const DenseTensor *input_x = &x_in;
  const DenseTensor *filter_x = &filter_x_in;
  const DenseTensor *scale_x = &scale_x_in;
  const DenseTensor *bias_x = &bias_x_in;

  // output x
  DenseTensor *conv_out_x = conv_x;

  DenseTensor *output = out;

  //  attrs
  float eps = epsilon;
  float momentum = momentum_in;
  bool is_train = !is_test && !use_global_stats;

  std::vector<const XPUType *> x_list = {
      reinterpret_cast<const XPUType *>(input_x->data<T>())};
  std::vector<const XPUType *> w_list = {
      reinterpret_cast<const XPUType *>(filter_x->data<T>())};
  std::vector<XPUType *> conv_y_list = {
      reinterpret_cast<XPUType *>(dev_ctx.template Alloc<T>(conv_out_x))};

  std::vector<std::vector<int64_t>> x_shape_list = {
      vectorize<int64_t>(input_x->dims())};

  auto filter_x_shape = vectorize<int64_t>(filter_x->dims());
  std::vector<int64_t> ksize = {filter_x_shape[2], filter_x_shape[3]};
  if (!is_nchw) {
    ksize[0] = filter_x_shape[1];
    ksize[1] = filter_x_shape[2];
  }
  std::vector<int64_t> strides = {stride, stride};
  std::vector<std::vector<int64_t>> ksize_list = {ksize};
  std::vector<std::vector<int64_t>> stride_list = {strides};
  std::vector<int64_t> paddings = {padding, padding};
  std::vector<int64_t> dilations = {dilation, dilation};
  std::vector<const float *> scale_list = {scale_x->data<float>()};
  std::vector<const float *> bias_list = {bias_x->data<float>()};
  std::vector<float *> batch_mean_list = {
      dev_ctx.template Alloc<float>(saved_mean_x)};
  std::vector<float *> batch_invstd_list = {
      dev_ctx.template Alloc<float>(saved_invstd_x)};
  std::vector<float *> global_mean_list = {
      dev_ctx.template Alloc<float>(running_mean_x)};
  std::vector<float *> global_var_list = {
      dev_ctx.template Alloc<float>(running_var_x)};

  std::vector<const float *> x_maxlist = {nullptr};
  std::vector<const float *> w_maxlist = {nullptr};
  if (has_shortcut) {
    // input z
    const DenseTensor *input_z = z_in.get_ptr();
    const DenseTensor *filter_z = filter_z_in.get_ptr();
    const DenseTensor *scale_z = scale_z_in.get_ptr();
    const DenseTensor *bias_z = bias_z_in.get_ptr();

    DenseTensor *conv_out_z = conv_z;

    x_list.push_back(reinterpret_cast<const XPUType *>(input_z->data<T>()));
    w_list.push_back(reinterpret_cast<const XPUType *>(filter_z->data<T>()));
    conv_y_list.push_back(
        reinterpret_cast<XPUType *>(dev_ctx.template Alloc<T>(conv_out_z)));

    x_shape_list.push_back(vectorize<int64_t>(input_z->dims()));

    auto filter_z_shape = vectorize<int64_t>(filter_z->dims());
    std::vector<int64_t> ksize_z = {filter_z_shape[2], filter_z_shape[3]};
    if (!is_nchw) {
      ksize_z[0] = filter_z_shape[1];
      ksize_z[1] = filter_z_shape[2];
    }
    ksize_list.push_back(ksize_z);
    stride_list.push_back({stride_z, stride_z});
    scale_list.push_back(scale_z->data<float>());
    bias_list.push_back(bias_z->data<float>());
    batch_mean_list.push_back(dev_ctx.template Alloc<float>(saved_mean_z));
    batch_invstd_list.push_back(dev_ctx.template Alloc<float>(saved_invstd_z));
    global_mean_list.push_back(dev_ctx.template Alloc<float>(running_mean_z));
    global_var_list.push_back(dev_ctx.template Alloc<float>(running_var_z));
    x_maxlist.push_back(nullptr);
    w_maxlist.push_back(nullptr);
  } else {
    if (fuse_add) {
      const DenseTensor *input_z = z_in.get_ptr();
      auto input_z_shape = vectorize<int64_t>(input_z->dims());
      x_list.push_back(reinterpret_cast<const XPUType *>(input_z->data<T>()));
      x_shape_list.push_back(input_z_shape);
      x_maxlist.push_back(nullptr);
    }
  }
  int r = xpu::resnet_unit_fusion<XPUType, XPUType, XPUType, int16_t>(
      dev_ctx.x_context(),
      x_list,
      w_list,
      conv_y_list,
      reinterpret_cast<XPUType *>(dev_ctx.template Alloc<T>(output)),
      x_shape_list,
      filter_x_shape[0],
      ksize_list,
      stride_list,
      paddings,
      dilations,
      group,
      eps,
      momentum,
      x_maxlist,
      w_maxlist,
      scale_list,
      bias_list,
      batch_mean_list,
      batch_invstd_list,
      global_mean_list,
      global_var_list,
      xpu::Activation_t::RELU,
      is_nchw,
      has_shortcut,
      fuse_add,
      is_train);
  PADDLE_ENFORCE_XDNN_SUCCESS(r, "resnet_unit_fusion");
}

}  // namespace phi

PD_REGISTER_KERNEL(resnet_unit,
                   XPU,
                   ALL_LAYOUT,
                   phi::ResNetUnitXPUKernel,
                   phi::float16,
                   float) {}