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

namespace tmo {
namespace torch_api {

#define CHECK_TENSOR_DTYPE(x, expected_type)                                     \
  TORCH_CHECK(x.scalar_type() == expected_type, "Tensor " #x " type should be ", \
              torchDtype2Str(expected_type), ".");

void dequant_from_linear_cache(
    at::Tensor &key,                         // [total_seqlen, head_num, head_size]
    const c10::optional<at::Tensor> &value,  // same as above
    const at::Tensor &key_cache,             // [max_batch_size, head_num, cache_mem_len, head_size]
    const c10::optional<at::Tensor> &value_cache,  // same as above
    const at::Tensor &key_quant_scale,             // quant_mode is 0: [head_num, head_size]
                                                   // quant_mode is 1:
                                                   // [max_batch_size, head_num, cache_mem_len]
    const c10::optional<at::Tensor> &value_quant_scale,  // same as above
    const at::Tensor &context_lengths,
    const int64_t max_context_len,
    const c10::optional<at::Tensor> &context_seq_offset,
    const c10::optional<at::Tensor> &cache_bs_id,
    const c10::optional<at::Tensor> &cache_seq_offset,
    const int64_t quant_mode,   // 0:per_channel, 1:per_head
    const int64_t quant_bit) {  // 4 or 8
  // check same attr for tensors
  checkTensorSameAttr<TensorAttr::DEVICE>(key, key_cache, key_quant_scale, value, value_cache,
                                          value_quant_scale, context_lengths, context_seq_offset,
                                          cache_bs_id, cache_seq_offset);
  checkTensorSameAttr<TensorAttr::DTYPE>(context_lengths, context_seq_offset, cache_bs_id,
                                         cache_seq_offset);

  // check quant parameters first
  TORCH_CHECK(quant_mode >= 0 && quant_mode <= 1, "quantization mode support 0 and 1.");
  TORCH_CHECK(quant_bit == 4 || quant_bit == 8, "quantization bit width support 4 and 8.");

  /****************************************check key***************************************/
  // check dtype
  TORCH_CHECK(key.scalar_type() == torch::kFloat16 || key.scalar_type() == torch::kBFloat16,
              "Tensor key type should be half or bfloat16.");
  CHECK_TENSOR_DTYPE(key_cache, torch::kInt8);
  CHECK_TENSOR_DTYPE(key_quant_scale, torch::kFloat32);

  // check_contiguous
  TORCH_CHECK(key.stride(-1) == 1, "Tensor key last dim must be contiguous.");
  CHECK_TENSOR_CONTIGUOUS(key_cache);
  CHECK_TENSOR_CONTIGUOUS(key_quant_scale);

  // check shape
  TORCH_CHECK(key.dim() == 3, "The dimensions of tensor key only support 3.");
  TORCH_CHECK(key_cache.dim() == 4, "The dimensions of tensor key_cache only supports 4.");
  TORCH_CHECK(context_lengths.dim() == 1,
              "The dimensions of tensor context_lengths only supports 1.");
  const int32_t total_seqlen = key.size(0);
  const int32_t head_num = key.size(1);
  const int32_t head_size = key.size(2);
  const int32_t max_batch_size = key_cache.size(0);
  const int32_t cache_mem_len = key_cache.size(2);
  const int32_t batch_size = context_lengths.size(0);
  CHECK_SHAPE(context_lengths, batch_size);
  TORCH_CHECK(max_batch_size >= batch_size,
              "max_batch_size should be greater than or equal to batch_size.");
  TORCH_CHECK(cache_mem_len % 2 == 0, "cache_mem_len should be a multiply of 2.");
  if (quant_mode == 0) {
    CHECK_SHAPE(key_quant_scale, head_num, head_size);
  } else if (quant_mode == 1) {
    CHECK_SHAPE(key_quant_scale, max_batch_size, head_num, cache_mem_len);
  }

  if (quant_bit == 4) {
    TORCH_CHECK(head_size % 2 == 0, "head_size should be a multiply of 2 if quant_bit is 4.");
    CHECK_SHAPE(key_cache, max_batch_size, head_num, cache_mem_len, head_size >> 1);
  } else {
    CHECK_SHAPE(key_cache, max_batch_size, head_num, cache_mem_len, head_size);
  }

  /***************************************check value***************************************/
  if (value.has_value() || value_cache.has_value() || value_quant_scale.has_value()) {
    TORCH_CHECK(value.has_value() && value_cache.has_value() && value_quant_scale.has_value(),
                "value, value_cache, and value_quant_scale must all exists.")
  }

  if (value_cache.has_value()) {
    checkTensorSameAttr<TensorAttr::DTYPE>(key, value);
    checkTensorSameAttr<TensorAttr::DTYPE>(key_cache, value_cache);
    checkTensorSameAttr<TensorAttr::DTYPE>(key_quant_scale, value_quant_scale);
    CHECK_OPTIONAL_TENSOR_CONTIGUOUS(value_cache);
    CHECK_OPTIONAL_TENSOR_CONTIGUOUS(value_quant_scale);
    TORCH_CHECK(key_quant_scale.dim() == value_quant_scale.value().dim(),
                "value_cache_quant_scale dim should keep same with key_cache_quant_scale.");
    CHECK_SHAPE(value.value(), total_seqlen, head_num, head_size);
    if (quant_mode == 0) {
      CHECK_SHAPE(value_quant_scale.value(), head_num, head_size);
    } else {
      CHECK_SHAPE(value_quant_scale.value(), max_batch_size, head_num, cache_mem_len);
    }

    if (quant_bit == 4) {
      CHECK_SHAPE(value_cache.value(), max_batch_size, head_num, cache_mem_len >> 1, head_size);
    } else {
      CHECK_SHAPE(value_cache.value(), max_batch_size, head_num, cache_mem_len, head_size);
    }

    for (int i = 0; i < key.dim(); i++) {
      TORCH_CHECK(value.value().stride(i) == key.stride(i),
                  "key and value must have same stride along axi ", i, ".");
    }
  }

  TORCH_CHECK(context_lengths.scalar_type() == torch::kInt32,
              "context_lengths type need be torch::kInt32.");
  /*********************************check optional tensor***********************************/
  const void *context_seq_offset_ptr = nullptr;
  if (context_seq_offset.has_value()) {
    CHECK_SHAPE(context_seq_offset.value(), batch_size);
    context_seq_offset_ptr = context_seq_offset.value().data_ptr();
  } else {
    TORCH_CHECK(batch_size <= 1024,
                "batch_size greater than 1024 not support when context_seq_offset is None.");
  }

  const void *cache_bs_id_ptr = nullptr;
  if (cache_bs_id.has_value()) {
    CHECK_SHAPE(cache_bs_id.value(), batch_size);
    cache_bs_id_ptr = cache_bs_id.value().data_ptr();
  }

  const void *cache_seq_offset_ptr = nullptr;
  if (cache_seq_offset.has_value()) {
    CHECK_SHAPE(cache_seq_offset.value(), batch_size);
    cache_seq_offset_ptr = cache_seq_offset.value().data_ptr();
  }

  // propare parameters before calling invokeDequantFromLinearCache
  const int32_t key_group_num = 1;
  const int32_t value_group_num = 1;
  const size_t context_head_stride = key.stride(1);
  const size_t context_seq_stride = key.stride(0);
  const size_t cache_bs_stride = key_cache.stride(0);
  const size_t cache_head_stride = key_cache.stride(1);
  const size_t key_cache_seq_stride = key_cache.stride(2);
  const size_t value_cache_seq_stride = value_cache.has_value() ? value_cache.value().stride(2) : 0;
  const size_t cache_scale_bs_stride = key_quant_scale.dim() == 3 ? key_quant_scale.stride(0) : 0;
  const size_t cache_scale_head_stride =
      key_quant_scale.dim() == 3 ? key_quant_scale.stride(1) : key_quant_scale.stride(0);

  const torch_mlu::mlu::MLUGuard device_guard(key.device());
  auto data_dtype = getCnnlDataType(key.scalar_type());
  auto queue = torch_mlu::getCurMLUStream();

  // run forward
  TMO_KERNEL_CHECK_FATAL(invokeDequantFromLinearCache(
      queue, getAtTensorPtr(key), getAtTensorPtr(value), getAtTensorPtr(key_cache),
      getAtTensorPtr(value_cache), getAtTensorPtr(key_quant_scale),
      getAtTensorPtr(value_quant_scale), getAtTensorPtr(context_lengths), context_seq_offset_ptr,
      cache_bs_id_ptr, cache_seq_offset_ptr, (int)max_context_len, batch_size, head_num,
      key_group_num, value_group_num, cache_mem_len, head_size, quant_mode, quant_bit,
      context_head_stride, context_seq_stride, cache_bs_stride, cache_head_stride,
      key_cache_seq_stride, value_cache_seq_stride, cache_scale_bs_stride, cache_scale_head_stride,
      data_dtype));
}
}  // namespace torch_api
}  // namespace tmo
