/*************************************************************************
 * 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 "kernels/update_out_and_lse.mluh"
#include "torch_ops_api.h"
#include "utils.h"

namespace tmo {
namespace torch_api {

void update_out_and_lse(at::Tensor &out,
                        at::Tensor &lse,
                        const at::Tensor &block_out,
                        const at::Tensor &block_lse,
                        const c10::optional<at::Tensor> &seq_offsets,
                        const c10::optional<at::Tensor> &cu_seqs,
                        const c10::optional<at::Tensor> &block_cu_seqs) {
  const torch_mlu::mlu::MLUGuard device_guard(out.device());
  auto queue = torch_mlu::getCurMLUStream();
  checkTensorSameAttr<TensorAttr::DEVICE>(out, lse, block_out, block_lse, seq_offsets, cu_seqs,
                                          block_cu_seqs);
  checkTensorSameAttr<TensorAttr::DTYPE>(seq_offsets, cu_seqs, block_cu_seqs);
  checkTensorSameAttr<TensorAttr::DTYPE>(out, block_out);
  checkTensorSameAttr<TensorAttr::DTYPE>(lse, block_lse);

  TORCH_CHECK(out.dim() == 3 || out.dim() == 4, "the dim of out must be 3 or 4");
  TORCH_CHECK(block_out.dim() == 3 || block_out.dim() == 4, "the dim of block_out must be 3 or 4");
  TORCH_CHECK(lse.dim() == 3, "the dim of lse must be 3");
  TORCH_CHECK(block_lse.dim() == 3, "the dim of block_lse must be 3");

  bool packed = out.dim() == 3;

  int32_t batch{0}, head_num{0}, head_size{0}, max_seq_len{0}, block_seq_len{0};
  int64_t bs_stride{0}, seq_stride{0}, head_stride{0}, block_bs_stride{0}, block_seq_stride{0},
      block_head_stride{0};
  auto data_type = getCnnlDataType(out.scalar_type());

  batch = lse.size(0);
  head_num = lse.size(1);
  max_seq_len = lse.size(2);
  block_seq_len = block_lse.size(2);
  head_size = out.size(-1);

  TORCH_CHECK(block_seq_len <= max_seq_len,
              "the max block_seq_len of block_lse can not be greater than the max_seq_len of lse. "
              "block_seq_len: ",
              block_seq_len, " max_seq_len: ", max_seq_len);
  TORCH_CHECK(data_type == CNNL_DTYPE_FLOAT || data_type == CNNL_DTYPE_BFLOAT16 ||
                  data_type == CNNL_DTYPE_HALF,
              "the data_type of out only support float, bfloat16 and half");

  if (packed) {
    seq_stride = out.stride(0);
    head_stride = out.stride(1);
    block_seq_stride = block_out.stride(0);
    block_head_stride = block_out.stride(1);

    if (max_seq_len == block_seq_len && block_seq_len == 1) {
      bs_stride = seq_stride;
      block_bs_stride = block_seq_stride;
    }

    CHECK_SHAPE(out, out.size(0), head_num, head_size);
    CHECK_SHAPE(block_out, block_out.size(0), head_num, head_size);
  } else {
    bs_stride = out.stride(0);
    seq_stride = out.stride(1);
    head_stride = out.stride(2);
    block_bs_stride = block_out.stride(0);
    block_seq_stride = block_out.stride(1);
    block_head_stride = block_out.stride(2);
    CHECK_SHAPE(out, batch, max_seq_len, head_num, head_size);
    CHECK_SHAPE(block_out, batch, block_seq_len, head_num, head_size);
  }

  TORCH_CHECK(head_size <= 1024, "the head_size can not be greater than 1024.");

  if (seq_offsets.has_value()) {
    TORCH_CHECK(seq_offsets.value().is_contiguous(), "seq_offsets must be contiguous");
    TORCH_CHECK(seq_offsets.value().dtype() == torch::kInt32,
                "the dtype of seq_offsets must be int32");
    TORCH_CHECK(seq_offsets.value().dim() == 1, "the dim of seq_offsets must be 1");
    CHECK_SHAPE(seq_offsets.value(), batch);
  }

  if (cu_seqs.has_value()) {
    TORCH_CHECK(cu_seqs.value().is_contiguous(), "cu_seqs must be contiguous");
    TORCH_CHECK(cu_seqs.value().dim() == 1, "the dim of cu_seqs must be 1");
    CHECK_SHAPE(cu_seqs.value(), batch + 1);
  }
  if (block_cu_seqs.has_value()) {
    TORCH_CHECK(block_cu_seqs.value().is_contiguous(), "block_cu_seqs must be contiguous");
    TORCH_CHECK(block_cu_seqs.value().dim() == 1, "the dim of block_cu_seqs must be 1");
    CHECK_SHAPE(block_cu_seqs.value(), batch + 1);
  }

  invokeUpdateOutAndLse(queue, getAtTensorPtr(out), (float *)getAtTensorPtr(lse),
                        getAtTensorPtr(block_out), (float *)getAtTensorPtr(block_lse),
                        (int32_t *)getAtTensorPtr(seq_offsets), (int32_t *)getAtTensorPtr(cu_seqs),
                        (int32_t *)getAtTensorPtr(block_cu_seqs), batch, head_num, head_size,
                        max_seq_len, block_seq_len, bs_stride, seq_stride, head_stride,
                        block_bs_stride, block_seq_stride, block_head_stride, packed, data_type);
}

}  //  namespace torch_api
}  //  namespace tmo
