/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/

#include "torch_ops_api.h"

namespace tmo {
namespace torch_api {

std::vector<at::Tensor> attention_project(const at::Tensor &input,
                                          const at::Tensor &q_weight,
                                          const c10::optional<at::Tensor> &q_bias,
                                          const c10::optional<at::Tensor> &k_weight,
                                          const c10::optional<at::Tensor> &k_bias,
                                          const c10::optional<at::Tensor> &v_weight,
                                          const c10::optional<at::Tensor> &v_bias,
                                          const c10::optional<at::Tensor> &norm_weight,
                                          const c10::optional<at::Tensor> &norm_bias,
                                          const c10::optional<at::Tensor> &residual,
                                          const std::string &out_layout,
                                          int64_t head_size,
                                          double eps,
                                          double alpha,
                                          double beta,
                                          bool norm_out) {
  // check device and dtype
  checkTensorSameAttr<TensorAttr::ALL>(input, q_weight, q_bias, k_weight, k_bias, v_weight, v_bias,
                                       norm_weight, norm_bias, residual);

  // check contiguous
  CHECK_TENSOR_CONTIGUOUS(input)
  CHECK_OPTIONAL_TENSOR_CONTIGUOUS(residual)

  // check size
  const int64_t nDim = input.dim();
  at::Tensor input_view = input;
  if (nDim == 2) {
    input_view = input.unsqueeze(0);
  }

  const bool has_k = k_weight.has_value();
  const bool has_v = v_weight.has_value();
  const int64_t n = input_view.size(0);
  const int64_t t = input_view.size(1);
  const int64_t hidden_size_q = q_weight.size(0);
  const int64_t hidden_size_k = has_k ? k_weight.value().size(0) : 0;
  const int64_t hidden_size_v = has_v ? v_weight.value().size(0) : 0;
  // check params
  bool has_bias = q_bias.has_value();
  bool has_ln = norm_weight.has_value();
  bool has_residual = residual.has_value();
  TORCH_CHECK(!(has_ln && has_residual), "cannot support layernorm and residual at the same time.")
  TORCH_CHECK(out_layout == "nthc" || (out_layout == "nhtc" && nDim == 3),
              "input must be 3-D if out_layout is 'nhtc'")
  bool trans_out = (out_layout == "nhtc");
  const int64_t head_num_q = hidden_size_q / head_size;
  const int64_t head_num_k = hidden_size_k / head_size;
  const int64_t head_num_v = hidden_size_v / head_size;

  const torch_mlu::mlu::MLUGuard device_guard(input_view.device());
  auto q_shape = trans_out ? std::vector<int64_t>({n, head_num_q, t, head_size})
                           : std::vector<int64_t>({n, t, hidden_size_q});
  auto k_shape = trans_out ? std::vector<int64_t>({n, head_num_k, t, head_size})
                           : std::vector<int64_t>({n, t, hidden_size_k});
  auto v_shape = trans_out ? std::vector<int64_t>({n, head_num_v, t, head_size})
                           : std::vector<int64_t>({n, t, hidden_size_v});
  auto out_q = at::empty(q_shape, input_view.options());
  auto out_k = has_k ? at::empty(k_shape, input_view.options()) : at::Tensor();
  auto out_v = has_v ? at::empty(q_shape, input_view.options()) : at::Tensor();
  auto out_ln = norm_out ? at::empty(input.sizes(), input_view.options()) : at::Tensor();

  // create tensor descs
  auto descs =
      createTensorDescs({input_view, q_weight, q_bias.value_or(at::Tensor()),
                         k_weight.value_or(at::Tensor()), k_bias.value_or(at::Tensor()),
                         v_weight.value_or(at::Tensor()), v_bias.value_or(at::Tensor()),
                         norm_weight.value_or(at::Tensor()), norm_bias.value_or(at::Tensor()),
                         residual.value_or(at::Tensor()), out_q, out_k, out_v, out_ln});

  // create and set attn_proj_desc
  cnnlTransformerLayernormResidualStructure_t layernorm_residual_mode =
      has_ln ? CNNL_TRANSFORMER_PRE_LAYERNORM_NO_RESIDUAL
             : (has_residual ? CNNL_TRANSFORMER_NO_LAYERNORM_WITH_RESIDUAL
                             : CNNL_TRANSFORMER_NO_LAYERNORM_NO_RESIDUAL);
  auto compute_dtype = getCnnlDataType(input_view.scalar_type());
  auto attn_proj_desc = tmo::op_desc::AttnProjDesc();
  attn_proj_desc.setDesc(layernorm_residual_mode, /*layernorm_residual_mode*/
                         compute_dtype, true,     /*has_q*/
                         has_k,                   /*has_k*/
                         has_v,                   /*has_v*/
                         has_bias, false, 0,      /*no packed*/
                         trans_out,               /*trans_out*/
                         norm_out,                /*store_layernorm out*/
                         alpha, beta,             /*alpha && beta */
                         eps /*eps*/);

  auto handle = torch_mlu::getCurrentHandle();
  // get workspace size
  size_t workspace_size = 0;
  CNNL_CHECK_FATAL(cnnlGetTransformerAttnProjWorkspaceSize(handle, attn_proj_desc, nullptr,
                                                           descs[0].get(), descs[1].get(),
                                                           descs[10].get(), &workspace_size));
  auto workspace =
      at::empty({static_cast<int64_t>(workspace_size)}, input.options().dtype(at::kByte));

  // run forward
  CNNL_CHECK_FATAL(cnnlTransformerAttnProj(
      handle, attn_proj_desc, nullptr,             /*quant_desc*/
      descs[0].get(), getAtTensorPtr(input_view),  /* input */
      descs[9].get(), getAtTensorPtr(residual),    /* residual */
      descs[1].get(), getAtTensorPtr(q_weight),    /* q weight */
      descs[3].get(), getAtTensorPtr(k_weight),    /* k weight */
      descs[5].get(), getAtTensorPtr(v_weight),    /* v weight */
      descs[2].get(), getAtTensorPtr(q_bias),      /* q bias */
      descs[4].get(), getAtTensorPtr(k_bias),      /* k bias */
      descs[6].get(), getAtTensorPtr(v_bias),      /* v bias */
      nullptr, nullptr,                            /*no valid token*/
      descs[7].get(), getAtTensorPtr(norm_weight), /* layernorm weight */
      descs[8].get(), getAtTensorPtr(norm_bias),   /* layernorm bias */
      getAtTensorPtr(workspace), workspace_size,   /* workspace */
      descs[10].get(), getAtTensorPtr(out_q),      /* q out */
      descs[11].get(), getAtTensorPtr(out_k),      /* k out */
      descs[12].get(), getAtTensorPtr(out_v),      /* v out */
      descs[13].get(), getAtTensorPtr(out_ln)      /* layernorm out */
      ));

  // return
  if (nDim == 2) {
    out_q.squeeze_(0);
    if (has_k) out_k.squeeze_(0);
    if (has_v) out_v.squeeze_(0);
    if (norm_out) out_ln.squeeze_(0);
  }
  std::vector<at::Tensor> output_list;
  output_list.emplace_back(out_q);
  if (has_k) output_list.emplace_back(out_k);
  if (has_v) output_list.emplace_back(out_v);
  if (norm_out) output_list.emplace_back(out_ln);
  return output_list;
}

}  //  namespace torch_api
}  //  namespace tmo
