// Copyright (c) 2023 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/phi/kernels/assign_kernel.h"
#include "paddle/phi/kernels/xpu/xpu_api_wrapper.h"

namespace phi {
namespace fusion {

#define TRANSFORMER_ENCODER_KERNEL_IMPL(x_dtype_, w_dtype_, gemm_dtype_) \
  int r = xpu::transformer_encoder<x_dtype_, w_dtype_, gemm_dtype_>(     \
      dev_ctx.x_context(),                                               \
      x_fp16_data,                                                       \
      fc_weight_data_##w_dtype_,                                         \
      out_fp16_data,                                                     \
      fc_input_max_data,                                                 \
      fc_weight_max_data,                                                \
      fc_bias_data,                                                      \
      ln_scale_data,                                                     \
      ln_bias_data,                                                      \
      qkv_attn_param,                                                    \
      mask_data);                                                        \
  PADDLE_ENFORCE_XDNN_SUCCESS(r, "multi_encoder_xpu");

template <typename T, typename Context>
void MultiEncoderXPUKernel(
    const Context& dev_ctx,
    const DenseTensor& x,
    const std::vector<const DenseTensor*>& fc_input_max,
    const std::vector<const DenseTensor*>& fc_weight,
    const std::vector<const DenseTensor*>& fc_weight_max,
    const std::vector<const DenseTensor*>& fc_bias,
    const std::vector<const DenseTensor*>& ln_scale,
    const std::vector<const DenseTensor*>& ln_bias,
    const std::vector<const DenseTensor*>& smooth_scale_weight,
    const std::vector<const DenseTensor*>& roformer_embedding,
    const optional<DenseTensor>& mask,
    const optional<DenseTensor>& seq_lod,
    const optional<DenseTensor>& max_seq_len,
    int layer_num,
    bool norm_before,
    int hidden_dim,
    int head_num,
    int size_per_head,
    int ffn_hidden_dim_scale,
    int act_type,
    int relative_type,
    int slice_idx,
    bool is_per_channel,
    int max_pos_len,
    const std::vector<float>& softmax_max_value,
    const std::vector<std::string>& quant_types,
    DenseTensor* out,
    DenseTensor* x_fp16,
    DenseTensor* out_fp16) {
  const int* seq_lod_data =
      seq_lod.get_ptr() == nullptr ? nullptr : seq_lod.get_ptr()->data<int>();
  const int* max_seq_len_data = max_seq_len.get_ptr() == nullptr
                                    ? nullptr
                                    : max_seq_len.get_ptr()->data<int>();
  int64_t batch_size = x.dims()[0];

  // TODO(large-tensor): XPU multi_encoder API not support int64
  PADDLE_ENFORCE_LE_INT_MAX(batch_size, "batch_size");

  int64_t seq_len = 1;
  int head_dim;
  if (x.dims().size() == 2) {
    head_dim = x.dims()[1];
  } else if (x.dims().size() == 3) {
    seq_len = x.dims()[1];
    head_dim = x.dims()[2];
  } else {
    PADDLE_ENFORCE(
        false,
        common::errors::PreconditionNotMet(
            "x.dims().size() MUST be 2 or 3, but get [%d].", x.dims().size()));
  }
  DDim out_dims;
  if (seq_lod_data) {
    batch_size = seq_lod.get_ptr()->numel() - 1;
    seq_len = max_seq_len_data[0];
  }
  out_dims = {batch_size, seq_len, head_dim};
  if (slice_idx != -1) {
    out_dims = {batch_size, head_dim};
  }
  out->Resize(out_dims);
  out_fp16->Resize(out_dims);
  // XPU2 only support fp16 input/output.
  auto x_dtype = x.dtype();
  const XPUTypeFP16* x_fp16_data = nullptr;
  XPUTypeFP16* out_fp16_data = nullptr;
  if (x_dtype == phi::DataType::FLOAT32) {
    auto* x_fp16_data_t = reinterpret_cast<XPUTypeFP16*>(
        dev_ctx.template Alloc<phi::float16>(x_fp16));
    int r_cast_x = xpu::cast<float, XPUTypeFP16>(
        dev_ctx.x_context(), x.data<float>(), x_fp16_data_t, x.numel());
    PADDLE_ENFORCE_XDNN_SUCCESS(r_cast_x,
                                "multi_encoder_xpu(cast x from fp32 to fp16)");
    x_fp16_data = x_fp16_data_t;
    out_fp16_data = reinterpret_cast<XPUTypeFP16*>(
        dev_ctx.template Alloc<phi::float16>(out_fp16));
  } else {
    x_fp16_data = reinterpret_cast<const XPUTypeFP16*>(x.data<phi::float16>());
    out_fp16_data = reinterpret_cast<XPUTypeFP16*>(
        dev_ctx.template Alloc<phi::float16>(out));
  }

  // q,k,v weight are fused.
  // Each encoder's weight should be: w0, null, null, w3, w4, w5
  auto enable_int8 = fc_weight[0]->dtype() == phi::DataType::INT8;
  auto local_quant = fc_weight[0]->dtype() == phi::DataType::FLOAT16;
  std::vector<xpu::QuantType> set_quant_types(8 * layer_num,
                                              xpu::QuantType::NOT_QUANT);
  if (enable_int8) {
    for (size_t i = 0; i < quant_types.size(); i++) {
      if (quant_types[i] == "enable_int8") {
        set_quant_types[i] = xpu::QuantType::QUANT_INT8;
      }
    }
  }
  std::vector<const float*> fc_input_max_data;
  std::vector<const int16_t*> fc_weight_data_int16_t;
  std::vector<const XPUTypeFP16*> fc_weight_data_XPUTypeFP16;
  std::vector<const float*> fc_weight_max_data;
  std::vector<const float*> fc_bias_data;
  for (size_t i = 0; i < fc_weight.size(); i++) {
    if (!enable_int8 && local_quant) {
      fc_weight_data_XPUTypeFP16.push_back(
          reinterpret_cast<const XPUTypeFP16*>(fc_weight[i]->data()));
    } else {
      // Int8 weight also convert to int16_t* for temporary storage.
      // The kernel dtype of int8 is chosen by quant_type in
      // xpu::transformer_encoder
      fc_weight_data_int16_t.push_back(
          reinterpret_cast<const int16_t*>(fc_weight[i]->data()));
    }
    fc_weight_max_data.push_back(fc_weight_max[i]->data<float>());
    fc_bias_data.push_back(fc_bias[i]->data<float>());
    if (i % 4 == 0) {
      fc_weight_data_int16_t.push_back(nullptr);
      fc_weight_data_int16_t.push_back(nullptr);
      fc_weight_data_XPUTypeFP16.push_back(nullptr);
      fc_weight_data_XPUTypeFP16.push_back(nullptr);
      fc_weight_max_data.push_back(nullptr);
      fc_weight_max_data.push_back(nullptr);
      fc_bias_data.push_back(nullptr);
      fc_bias_data.push_back(nullptr);
    }
  }

  for (size_t i = 0; i < fc_input_max.size(); i++) {
    fc_input_max_data.push_back(fc_input_max[i]->data<float>());
  }

  std::vector<const float*> ln_scale_data;
  std::vector<const float*> ln_bias_data;
  for (size_t i = 0; i < ln_scale.size(); i++) {
    ln_scale_data.push_back(ln_scale[i]->data<float>());
    ln_bias_data.push_back(ln_bias[i]->data<float>());
  }
  const float* mask_data =
      mask.get_ptr() == nullptr ? nullptr : mask.get_ptr()->data<float>();

  xpu::Activation_t qkv_act(static_cast<xpu::Activation_t::act_enum>(act_type));

  int64_t batch = x.dims()[0];

  // TODO(large-tensor): XPU multi_encoder QKVAttnParam not support int64
  PADDLE_ENFORCE_LE_INT_MAX(batch, "batch");

  // matmul_size * layer_num
  if (seq_lod_data) {
    xpu::VectorParam<int> query_lod = {
        seq_lod_data, seq_lod.get_ptr()->numel(), nullptr};
    int max_seq_len_value = slice_idx == -1 ? max_seq_len_data[0] : -1;
    xpu::QKVAttnParam qkv_attn_param(query_lod,
                                     head_num,
                                     size_per_head,
                                     qkv_act,
                                     slice_idx,
                                     true,
                                     max_seq_len_value,
                                     hidden_dim,
                                     norm_before,
                                     is_per_channel);
    if (!softmax_max_value.empty()) {
      qkv_attn_param.ptq_max_value = softmax_max_value;
    }
    if (!smooth_scale_weight.empty()) {
      qkv_attn_param.is_smooth_quant = true;
      std::vector<const XPUTypeFP16*> smooth_scale_weight_ptr;
      for (const auto& weight : smooth_scale_weight) {
        auto tmp_ptr =
            reinterpret_cast<const XPUTypeFP16*>(weight->data<phi::float16>());
        smooth_scale_weight_ptr.push_back(tmp_ptr);
      }
      qkv_attn_param.smooth_scale.assign(smooth_scale_weight_ptr.begin(),
                                         smooth_scale_weight_ptr.end());
    }
    qkv_attn_param.quant_type_.assign(set_quant_types.begin(),
                                      set_quant_types.end());
    qkv_attn_param.scale_of_hidden_units = ffn_hidden_dim_scale;
    if (!roformer_embedding.empty()) {
      std::vector<const float*> roformer_embedding_data;
      for (size_t i = 0; i < roformer_embedding.size(); i++) {
        roformer_embedding_data.push_back(roformer_embedding[i]->data<float>());
      }
      qkv_attn_param.relative_type = relative_type;
      qkv_attn_param.max_pos_len = max_pos_len;
      qkv_attn_param.relative_pos.assign(roformer_embedding_data.begin(),
                                         roformer_embedding_data.end());
    }
    if (!enable_int8 && local_quant) {
      TRANSFORMER_ENCODER_KERNEL_IMPL(XPUTypeFP16, XPUTypeFP16, float)
    } else {
      // The kernel dtype of int8 is chosen by quant_type in
      // xpu::transformer_encoder This template args, int16_t, is only for skip
      // quant fc
      TRANSFORMER_ENCODER_KERNEL_IMPL(XPUTypeFP16, int16_t, int16_t)
    }
  } else if (mask_data) {
    auto mask_dims = mask.get_ptr()->dims();
    std::vector<int> mask_shape(mask_dims.Get(),
                                mask_dims.Get() + mask_dims.size());
    int64_t max_seq_len_value = x.dims()[1];

    // TODO(large-tensor): XPU QKVAttnParam not support int64
    PADDLE_ENFORCE_LE_INT_MAX(max_seq_len_value, "max_seq_len_value");

    xpu::QKVAttnParam qkv_attn_param(static_cast<int>(batch),
                                     static_cast<int>(max_seq_len_value),
                                     head_num,
                                     size_per_head,
                                     mask_shape,
                                     qkv_act,
                                     slice_idx,
                                     true,
                                     hidden_dim,
                                     norm_before,
                                     is_per_channel);
    if (!softmax_max_value.empty()) {
      qkv_attn_param.ptq_max_value = softmax_max_value;
    }
    if (!smooth_scale_weight.empty()) {
      qkv_attn_param.is_smooth_quant = true;
      std::vector<const XPUTypeFP16*> smooth_scale_weight_ptr;
      for (const auto& weight : smooth_scale_weight) {
        auto tmp_ptr =
            reinterpret_cast<const XPUTypeFP16*>(weight->data<phi::float16>());
        smooth_scale_weight_ptr.push_back(tmp_ptr);
      }
      qkv_attn_param.smooth_scale.assign(smooth_scale_weight_ptr.begin(),
                                         smooth_scale_weight_ptr.end());
    }
    qkv_attn_param.quant_type_.assign(set_quant_types.begin(),
                                      set_quant_types.end());
    qkv_attn_param.scale_of_hidden_units = ffn_hidden_dim_scale;
    if (!roformer_embedding.empty()) {
      std::vector<const float*> roformer_embedding_data;
      for (size_t i = 0; i < roformer_embedding.size(); i++) {
        roformer_embedding_data.push_back(roformer_embedding[i]->data<float>());
      }
      qkv_attn_param.relative_type = relative_type;
      qkv_attn_param.max_pos_len = max_pos_len;
      qkv_attn_param.relative_pos.assign(roformer_embedding_data.begin(),
                                         roformer_embedding_data.end());
    }
    if (!enable_int8 && local_quant) {
      TRANSFORMER_ENCODER_KERNEL_IMPL(XPUTypeFP16, XPUTypeFP16, float)
    } else {
      TRANSFORMER_ENCODER_KERNEL_IMPL(XPUTypeFP16, int16_t, int16_t)
    }
  } else {
    // When no mask input, like VIT, create LOD to act as vsl.
    int64_t max_seq_len_value = x.dims()[1];

    // TODO(large-tensor): XPU QKVAttnParam not support int64
    PADDLE_ENFORCE_LE_INT_MAX(max_seq_len_value * batch,
                              "max_seq_len_value*batch");

    std::vector<int> lod;
    for (int64_t i = 0; i < batch + 1; i++) {
      lod.push_back(static_cast<int>(i * max_seq_len_value));
    }
    xpu::VectorParam<int> query_lod = {
        lod.data(), static_cast<int>(lod.size()), nullptr};
    // No need to pad, no matter slice or not
    xpu::QKVAttnParam qkv_attn_param(query_lod,
                                     head_num,
                                     size_per_head,
                                     qkv_act,
                                     slice_idx,
                                     true,
                                     -1,
                                     hidden_dim,
                                     norm_before,
                                     is_per_channel);
    if (!softmax_max_value.empty()) {
      qkv_attn_param.ptq_max_value = softmax_max_value;
    }
    if (!smooth_scale_weight.empty()) {
      qkv_attn_param.is_smooth_quant = true;
      std::vector<const XPUTypeFP16*> smooth_scale_weight_ptr;
      for (const auto& weight : smooth_scale_weight) {
        auto tmp_ptr =
            reinterpret_cast<const XPUTypeFP16*>(weight->data<phi::float16>());
        smooth_scale_weight_ptr.push_back(tmp_ptr);
      }
      qkv_attn_param.smooth_scale.assign(smooth_scale_weight_ptr.begin(),
                                         smooth_scale_weight_ptr.end());
    }
    qkv_attn_param.quant_type_.assign(set_quant_types.begin(),
                                      set_quant_types.end());
    qkv_attn_param.scale_of_hidden_units = ffn_hidden_dim_scale;
    if (!roformer_embedding.empty()) {
      std::vector<const float*> roformer_embedding_data;
      for (size_t i = 0; i < roformer_embedding.size(); i++) {
        roformer_embedding_data.push_back(roformer_embedding[i]->data<float>());
      }
      qkv_attn_param.relative_type = relative_type;
      qkv_attn_param.max_pos_len = max_pos_len;
      qkv_attn_param.relative_pos.assign(roformer_embedding_data.begin(),
                                         roformer_embedding_data.end());
    }
    if (!enable_int8 && local_quant) {
      TRANSFORMER_ENCODER_KERNEL_IMPL(XPUTypeFP16, XPUTypeFP16, float)
    } else {
      TRANSFORMER_ENCODER_KERNEL_IMPL(XPUTypeFP16, int16_t, int16_t)
    }
  }

  if (x_dtype == phi::DataType::FLOAT32) {
    int r_cast_out =
        xpu::cast<XPUTypeFP16, float>(dev_ctx.x_context(),
                                      out_fp16_data,
                                      dev_ctx.template Alloc<float>(out),
                                      out->numel());
    PADDLE_ENFORCE_XDNN_SUCCESS(
        r_cast_out, "multi_encoder_xpu(cast out from fp16 to fp32)");
  }
}

}  // namespace fusion
}  // namespace phi

PD_REGISTER_KERNEL(multi_encoder_xpu,
                   XPU,
                   ALL_LAYOUT,
                   phi::fusion::MultiEncoderXPUKernel,
                   float,
                   phi::float16) {
  kernel->InputAt(10).SetBackend(phi::Backend::CPU);
  kernel->InputAt(11).SetBackend(phi::Backend::CPU);
}