/*************************************************************************
 * 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_paged_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_paged_cache(at::Tensor &key,  // [total_seqlen, head_num, head_size]
                              const c10::optional<at::Tensor> &value,  // same as above
                              const at::Tensor &key_cache,  // [token_num, head_num, block_size]
                              const c10::optional<at::Tensor> &value_cache,  // same as above
                              // [token_num, head_num, block_size, head_size] for per-token
                              // [head_num, head_size] for per-channel
                              const at::Tensor &key_cache_quant_scale,
                              const c10::optional<at::Tensor> &value_cache_quant_scale,
                              const at::Tensor &context_lengths,  // [batch_size]
                              int64_t max_context_len,
                              const c10::optional<at::Tensor> &context_seq_offset,  // [batch_size]
                              const at::Tensor &block_tables,  // [batch_size, max_block_num]
                              int64_t quant_mode,              // 0 is per-channel, 1 is per-token
                              int64_t quant_bit) {             // quantization bit only support 8
  // check same attr for tensors
  checkTensorSameAttr<TensorAttr::DEVICE>(key, key_cache, key_cache_quant_scale, value, value_cache,
                                          value_cache_quant_scale, context_lengths,
                                          context_seq_offset, block_tables);
  checkTensorSameAttr<TensorAttr::DTYPE>(context_lengths, context_seq_offset, block_tables);

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

  /***************************************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_cache_quant_scale, torch::kFloat32);
  CHECK_TENSOR_DTYPE(block_tables, torch::kInt32);

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

  // check shape
  TORCH_CHECK(key.dim() == 3, "The dimensions of tensor key only supports 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.");
  TORCH_CHECK(block_tables.dim() == 2, "The dimensions of tensor block_tables only supports 2.");
  const int32_t token_num = key.size(0);
  const int32_t head_num = key.size(1);
  const int32_t head_size = key.size(2);
  const int32_t block_num = key_cache.size(0);
  const int32_t block_size = key_cache.size(2);
  const int32_t batch_size = context_lengths.size(0);
  const int32_t max_block_num = block_tables.size(1);
  CHECK_SHAPE(key_cache, block_num, head_num, block_size, head_size);
  int64_t kv_cache_range =
      (int64_t)block_num * head_num * block_size * head_size * key_cache.element_size();
  // kernel use uint32_t to calculate offsets
  TORCH_CHECK(kv_cache_range <= UINT32_MAX, "The addressing range of key_cache cannot exceed 4GB.");
  if (quant_mode == 0) {
    CHECK_SHAPE(key_cache_quant_scale, head_num, head_size);
  } else if (quant_mode == 1) {
    CHECK_SHAPE(key_cache_quant_scale, block_num, head_num, block_size);
  }

  /***************************************check value***************************************/
  if (value.has_value() || value_cache.has_value() || value_cache_quant_scale.has_value()) {
    TORCH_CHECK(value.has_value() && value_cache.has_value() && value_cache_quant_scale.has_value(),
                "value, value_cache, and value_cache_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_cache_quant_scale, value_cache_quant_scale);
    TORCH_CHECK(value.value().stride(-1) == 1, "Tensor value last dim must be contiguous.")
    CHECK_OPTIONAL_TENSOR_CONTIGUOUS(value_cache);
    CHECK_OPTIONAL_TENSOR_CONTIGUOUS(value_cache_quant_scale);
    CHECK_SHAPE(value.value(), token_num, head_num, head_size);
    TORCH_CHECK(value_cache_quant_scale.value().dim() == key_cache_quant_scale.dim(),
                "value_cache_quant_scale dim should keep same with key_cache_quant_scale.");
    if (quant_mode == 0) {
      CHECK_SHAPE(value_cache_quant_scale.value(), head_num, head_size);
    } else {
      CHECK_SHAPE(value_cache_quant_scale.value(), block_num, head_num, block_size);
    }

    CHECK_SHAPE(value_cache.value(), block_num, head_num, block_size, 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, ".");
    }
  }

  /*********************************check index tensor***********************************/
  CHECK_SHAPE(context_lengths, batch_size);
  CHECK_SHAPE(block_tables, batch_size, max_block_num);
  TORCH_CHECK(max_context_len <= block_size * max_block_num,
              "max_context_len should smaller than or equal to block_size * max_block_num.");

  /*********************************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.");
  }

  // 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_cache_quant_scale.dim() == 3 ? key_cache_quant_scale.stride(0) : 0;
  const size_t cache_scale_head_stride = key_cache_quant_scale.dim() == 3
                                             ? key_cache_quant_scale.stride(1)
                                             : key_cache_quant_scale.stride(0);

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

  // run forward
  TMO_KERNEL_CHECK_FATAL(invokeDequantFromPagedCache(
      queue, getAtTensorPtr(key), getAtTensorPtr(value), getAtTensorPtr(key_cache),
      getAtTensorPtr(value_cache), getAtTensorPtr(key_cache_quant_scale),
      getAtTensorPtr(value_cache_quant_scale), getAtTensorPtr(context_lengths),
      context_seq_offset_ptr, getAtTensorPtr(block_tables), (int)max_context_len, max_block_num,
      batch_size, head_num, key_group_num, value_group_num, block_size, 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
