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

namespace tmo {
namespace torch_api {

void offline_quant_to_linear_cache(
    const at::Tensor &key,  // [total_seqlen, num_heads, head_size]
                            // or [bs, max_context_len, num_heads, head_size]
    const c10::optional<at::Tensor> &value,
    at::Tensor &key_cache,  // [max_bs, num_heads, cache_memory_len, head_size]
    const c10::optional<at::Tensor> &value_cache,
    const at::Tensor &key_cache_scale,  // [num_heads, cache_memory_len] or [num_heads, head_size]
    const c10::optional<at::Tensor> &value_cache_scale,
    const at::Tensor &context_lengths,
    const int64_t max_context_len,
    const int64_t quant_mode,  // 0:per_channel, others:per_head
    const bool packed,
    const c10::optional<at::Tensor> &context_seq_offset,
    const c10::optional<at::Tensor> &cache_bs_id,
    const c10::optional<at::Tensor> &cache_seqlen_offset) {
  /****************************************check key***************************************/
  // check dtype
  TORCH_CHECK(key.scalar_type() == torch::kFloat32 || key.scalar_type() == torch::kFloat16 ||
                  key.scalar_type() == torch::kBFloat16,
              "key type need be torch::kFloat32, torch::kFloat16 or torch::kBFloat16");
  TORCH_CHECK(key_cache.scalar_type() == torch::kInt8, "key_cache type need be torch::kInt8");
  TORCH_CHECK(key_cache_scale.scalar_type() == torch::kFloat32,
              "key_cache_scale type need be torch::kFloat32");

  // check shape
  const int batch_size = packed ? context_lengths.size(0) - 1 : context_lengths.size(0);
  const int max_bs = key_cache.size(0);
  const int num_heads = key_cache.size(1);
  const int cache_mem_len = key_cache.size(2);
  const int head_size = key_cache.size(3);
  const int total_seqlen = packed ? key.size(0) : batch_size * key.size(-2);

  TORCH_CHECK(key_cache.size(0) >= batch_size,
              "first dim of key_cache must be great than or equal to batch_size");
  if (packed) {
    CHECK_SHAPE(key, total_seqlen, num_heads, head_size);
    CHECK_SHAPE(context_lengths, batch_size + 1);
  } else {
    CHECK_SHAPE(key, batch_size, key.size(1), num_heads, head_size);
    CHECK_SHAPE(context_lengths, batch_size);
  }
  if (quant_mode == 0) {
    CHECK_SHAPE(key_cache_scale, num_heads, head_size);
  } else {
    CHECK_SHAPE(key_cache_scale, num_heads, cache_mem_len);
  }

  // check contiguous
  TORCH_CHECK(key_cache.is_contiguous(), "key_cache tensor must be contiguous.");
  TORCH_CHECK(key_cache_scale.is_contiguous(), "key_cache_scale tensor must be contiguous.");

  // check device
  const int64_t device_id = key.get_device();
  TORCH_CHECK(key_cache.get_device() == device_id,
              "key_cache tensor device index is not same, original index: ", device_id,
              "now index is: ", key_cache.get_device());
  TORCH_CHECK(key_cache_scale.get_device() == device_id,
              "key_cache_scale tensor device index is not same, original index: ", device_id,
              "now index is: ", key_cache_scale.get_device());

  /***************************************check value***************************************/
  if (value.has_value() || value_cache.has_value() || value_cache_scale.has_value()) {
    TORCH_CHECK(value.has_value() && value_cache.has_value() && value_cache_scale.has_value(),
                "value_cache, value_cache_scale, value must all have value.")
  }
  if (value_cache.has_value()) {
    TORCH_CHECK(value_cache.value().scalar_type() == torch::kInt8,
                "value_cache type need be torch::kInt8");
    TORCH_CHECK(value_cache_scale.value().scalar_type() == torch::kFloat32,
                "value_cache_scale type need be torch::kFloat32");
    TORCH_CHECK(value.value().scalar_type() == key.scalar_type(),
                "value type need be same with key");
    if (packed) {
      CHECK_SHAPE(value.value(), total_seqlen, num_heads, head_size);
    } else {
      CHECK_SHAPE(value.value(), batch_size, key.size(1), num_heads, head_size);
    }
    if (quant_mode == 0) {
      CHECK_SHAPE(value_cache_scale.value(), num_heads, head_size);
    } else {
      CHECK_SHAPE(value_cache_scale.value(), num_heads, cache_mem_len);
    }
    CHECK_SHAPE(value_cache.value(), max_bs, num_heads, 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(value_cache.value().is_contiguous(), "value_cache tensor must be contiguous.");
    TORCH_CHECK(value_cache_scale.value().is_contiguous(),
                "value_cache_scale tensor must be contiguous.");
    TORCH_CHECK(value_cache.value().get_device() == device_id,
                "value_cache tensor device index is not same, original index: ", device_id,
                "now index is: ", value_cache.value().get_device());
    TORCH_CHECK(value_cache_scale.value().get_device() == device_id,
                "value_cache_scale tensor device index is not same, original index: ", device_id,
                "now index is: ", value_cache_scale.value().get_device());
  }

  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);
    TORCH_CHECK(context_seq_offset.value().scalar_type() == torch::kInt32,
                "context_seq_offset type need be torch::kInt32");
    TORCH_CHECK(context_seq_offset.value().get_device() == device_id,
                "context_seq_offset tensor device index is not same, original index: ", device_id,
                "now index is: ", context_seq_offset.value().get_device());
    context_seq_offset_ptr = context_seq_offset.value().data_ptr();
  }

  const void *cache_bs_id_ptr = nullptr;
  if (cache_bs_id.has_value()) {
    CHECK_SHAPE(cache_bs_id.value(), batch_size);
    TORCH_CHECK(cache_bs_id.value().scalar_type() == torch::kInt32,
                "cache_bs_id type need be torch::kInt32");
    TORCH_CHECK(cache_bs_id.value().get_device() == device_id,
                "cache_bs_id tensor device index is not same, original index: ", device_id,
                "now index is: ", cache_bs_id.value().get_device());
    cache_bs_id_ptr = cache_bs_id.value().data_ptr();
  }

  const void *cache_seqlen_offset_ptr = nullptr;
  if (cache_seqlen_offset.has_value()) {
    CHECK_SHAPE(cache_seqlen_offset.value(), batch_size);
    TORCH_CHECK(cache_seqlen_offset.value().scalar_type() == torch::kInt32,
                "cache_seqlen_offset type need be torch::kInt32");
    TORCH_CHECK(cache_seqlen_offset.value().get_device() == device_id,
                "cache_seqlen_offset tensor device index is not same, original index: ", device_id,
                "now index is: ", cache_seqlen_offset.value().get_device());
    cache_seqlen_offset_ptr = cache_seqlen_offset.value().data_ptr();
  }

  const size_t context_bs_stride = packed ? 0 : key.stride(0);
  const size_t context_head_stride = packed ? key.stride(1) : key.stride(2);
  const size_t context_seq_stride = packed ? key.stride(0) : key.stride(1);
  const size_t cache_bs_stride = key_cache.stride(0);
  const size_t cache_head_stride = key_cache.stride(1);
  const size_t cache_seq_stride = key_cache.stride(2);
  const size_t cache_scale_head_stride = key_cache_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(invokeOfflineQuantToLinearCache(
      queue, getAtTensorPtr(key_cache), getAtTensorPtr(value_cache),
      getAtTensorPtr(key_cache_scale), getAtTensorPtr(value_cache_scale), cache_bs_id_ptr,
      cache_seqlen_offset_ptr, getAtTensorPtr(key), getAtTensorPtr(value), context_seq_offset_ptr,
      getAtTensorPtr(context_lengths), data_dtype, batch_size, num_heads, head_size,
      (int)max_context_len, cache_mem_len, context_bs_stride, context_head_stride,
      context_seq_stride, cache_bs_stride, cache_head_stride, cache_seq_stride,
      cache_scale_head_stride, packed, quant_mode));
}
}  // namespace torch_api
}  // namespace tmo
