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

namespace tmo {
namespace torch_api {

void fused_rope(at::Tensor &qkv,
                at::Tensor &key_cache_hp,
                at::Tensor &value_cache_hp,
                const c10::optional<at::Tensor> &key_cache_lp,
                const c10::optional<at::Tensor> &value_cache_lp,
                const at::Tensor &sin_table,
                const at::Tensor &cos_table,
                const at::Tensor &position_ids,
                const at::Tensor &gamma,
                const at::Tensor &beta,
                const c10::optional<at::Tensor> &key_scale_hp,
                const c10::optional<at::Tensor> &value_scale_hp,
                const c10::optional<at::Tensor> &key_scale_lp,
                const c10::optional<at::Tensor> &value_scale_lp,
                const c10::optional<at::Tensor> &cache_bs_id_hp,
                const c10::optional<at::Tensor> &cache_seq_offsets_hp,
                const c10::optional<at::Tensor> &cache_bs_id_lp,
                const c10::optional<at::Tensor> &cache_seq_offsets_lp,
                const c10::optional<at::Tensor> &slot_mapping_hp,
                const c10::optional<at::Tensor> &slot_mapping_lp,
                const double eps) {
  bool paged_cache_hp = slot_mapping_hp.has_value();
  bool paged_cache_lp = slot_mapping_lp.has_value();
  bool mixed_cache = key_cache_lp.has_value() && value_cache_lp.has_value();
  const int origin_device_id = qkv.get_device();

  checkTensorSameAttr<TensorAttr::ALL>(qkv, sin_table, cos_table, gamma, beta);
  TORCH_CHECK(qkv.is_contiguous(), "qkv tensor must be contiguous.")
  TORCH_CHECK(key_cache_hp.is_contiguous(), "key_cache_hp tensor must be contiguous")
  TORCH_CHECK(value_cache_hp.is_contiguous(), "value_cache_hp tensor must be contiguous")
  if (mixed_cache) {
    TORCH_CHECK(key_cache_lp.value().is_contiguous(), "key_cache_lp tensor must be contiguous")
    TORCH_CHECK(key_scale_hp.has_value() && value_scale_hp.has_value(),
                "key_scale_hp and value_scale_hp must not be null under mixed cache.")
  }
  if (mixed_cache) {
    TORCH_CHECK(value_cache_lp.value().is_contiguous(), "value_cache_lp tensor must be contiguous")
  }
  TORCH_CHECK(position_ids.is_contiguous(), "position_ids tensor must be contiguous");
  TORCH_CHECK(gamma.is_contiguous(), "gamma tensor must be contiguous");
  TORCH_CHECK(beta.is_contiguous(), "beta tensor must be contiguous");
  if (cache_bs_id_hp.has_value()) {
    TORCH_CHECK(cache_bs_id_hp.value().is_contiguous(), "cache_bs_id_hp tensor must be contiguous.")
  }
  if (cache_seq_offsets_hp.has_value()) {
    TORCH_CHECK(cache_seq_offsets_hp.value().is_contiguous(),
                "cache_seq_offsets_hp tensor must be contiguous")
  }
  if (cache_bs_id_lp.has_value()) {
    TORCH_CHECK(cache_bs_id_lp.value().is_contiguous(), "cache_bs_id_lp tensor must be contiguous.")
  }
  if (cache_seq_offsets_lp.has_value()) {
    TORCH_CHECK(cache_seq_offsets_lp.value().is_contiguous(),
                "cache_seq_offsets_lp tensor must be contiguous")
  }
  if (paged_cache_hp) {
    TORCH_CHECK(slot_mapping_hp.value().is_contiguous(),
                "slot_mapping_hp tensor must be contiguous.")
  }
  if (paged_cache_lp) {
    TORCH_CHECK(slot_mapping_lp.value().is_contiguous(),
                "slot_mapping_lp tensor must be contiguous.")
  }

  // check qkv key_cache value_cache
  auto dtype = qkv.dtype();
  TORCH_CHECK(dtype == torch::kFloat16 || dtype == torch::kBFloat16,
              "qkv dtype must be half or bfloat16.")
  TORCH_CHECK(qkv.dim() == 4,
              "qkv tensor dim must be 4: (batch_size, 1, head_num_q + head_num_kv * 2, head_size).")
  TORCH_CHECK(qkv.size(1) == 1, "only support seq_len == 1.");
  TORCH_CHECK(key_cache_hp.dim() == 4,
              "key_cache_hp dim must be 4: (max_bs, head_num_kv, max_decode_len_hp, head_size),"
              "or (num_blocks, head_num_kv, block_size_hp, head_size).");
  TORCH_CHECK(value_cache_hp.dim() == 4,
              "value_cache_hp dim must be 4: (max_bs, head_num_kv, max_decode_len_hp, head_size),"
              "or (num_blocks, head_num_kv, block_size_hp, head_size).");
  void *key_cache_lp_ptr = nullptr;
  void *value_cache_lp_ptr = nullptr;
  if (mixed_cache) {
    TORCH_CHECK(
        key_cache_lp.value().dim() == 4,
        "key_cache_lp dim must be 4: (max_bs, head_num_kv, max_decode_len_lp, head_size / 2),"
        "or (num_blocks, head_num_kv, block_size_lp, head_size / 2).");
    TORCH_CHECK(
        value_cache_lp.value().dim() == 4,
        "value_cache_lp dim must be 4: (max_bs, head_num_kv, max_decode_len_lp / 2, head_size),"
        "or (num_blocks, head_num_kv, block_size_lp / 2, head_size).");
    TORCH_CHECK(key_cache_lp.value().get_device() == origin_device_id,
                "key_scale_hp tensor device index is not same, original index: ", origin_device_id,
                "now index is: ", key_cache_lp.value().get_device());
    TORCH_CHECK(value_cache_lp.value().get_device() == origin_device_id,
                "value_scale_hp tensor device index is not same, original index: ",
                origin_device_id, "now index is: ", value_cache_lp.value().get_device());
    key_cache_lp_ptr = key_cache_lp.value().data_ptr();
    value_cache_lp_ptr = value_cache_lp.value().data_ptr();
  }

  int batch_size = qkv.size(0);
  int head_num_qkv = qkv.size(-2);
  int head_size = qkv.size(-1);
  int head_num_kv = key_cache_hp.size(1);  // linear cache and paged cache same
  int max_bs_hp = key_cache_hp.size(0);    // for linear cache, paged cache no use
  int max_bs_lp = mixed_cache ? key_cache_lp.value().size(0) : 0;
  int max_decode_len_hp = key_cache_hp.size(2);  // for linear cache, paged cache no use
  int max_decode_len_lp = mixed_cache ? key_cache_lp.value().size(2) : 0;
  int num_blocks_hp = key_cache_hp.size(0);
  int num_blocks_lp = mixed_cache ? key_cache_lp.value().size(0) : 0;
  int block_size_hp = key_cache_hp.size(2);  // for paged cache, linear cache no use
  int block_size_lp = mixed_cache ? key_cache_lp.value().size(2) : 0;
  int head_num_q = head_num_qkv - 2 * head_num_kv;
  int group_size = head_size;

  TORCH_CHECK(head_size <= 128, "only support qkv head_size <= 128.");
  TORCH_CHECK(head_size % 2 == 0, "head_size must be divided by 2.");
  TORCH_CHECK(head_num_q <= 32, "only support head_num_q <= 32.");
  TORCH_CHECK(head_num_kv <= 32, "only support head_num_kv <= 32.");
  int kv_out_size = key_scale_hp.has_value() && value_scale_hp.has_value() ? 1 : 2;
  size_t hp_cache_total_bytes =
      paged_cache_hp
          ? (size_t)num_blocks_hp * head_num_kv * block_size_hp * head_size * kv_out_size
          : (size_t)max_bs_hp * max_decode_len_hp * head_num_kv * head_size * kv_out_size;
  TORCH_CHECK(hp_cache_total_bytes <= INT32_MAX,
              "hp_cache memory btyes must be less than or equal to INT32_MAX.");
  if (mixed_cache) {
    size_t lp_cache_total_bytes =
        paged_cache_lp ? (size_t)num_blocks_lp * head_num_kv * block_size_lp * head_size / 2
                       : (size_t)max_bs_lp * max_decode_len_lp * head_num_kv * head_size / 2;
    TORCH_CHECK(lp_cache_total_bytes <= INT32_MAX,
                "lp_cache memory btyes must be less than or equal to INT32_MAX.");
  }

  // check sin_table cos_table gamma beta
  TORCH_CHECK(sin_table.dim() == 2 && cos_table.dim() == 2,
              "sin_table and cos_table tensor dim must be 2: (rotary_seqlen, head_size).");
  TORCH_CHECK(sin_table.size(1) == head_size && cos_table.size(1) == head_size,
              "rotary_dim must be same with head_size.");
  TORCH_CHECK(sin_table.stride(0) == cos_table.stride(0),
              "sin_table first stride must be equal to cos_table first_stride.");
  int rotary_stride = sin_table.stride(0);

  CHECK_SHAPE(gamma, head_size);
  CHECK_SHAPE(beta, head_size);

  // check key value scale
  bool has_kv_scale = key_scale_hp.has_value() && value_scale_hp.has_value();
  const void *key_scale_hp_ptr = nullptr;
  const void *value_scale_hp_ptr = nullptr;
  if (has_kv_scale) {
    CHECK_SHAPE(key_scale_hp.value(), head_num_kv, head_size);
    CHECK_SHAPE(value_scale_hp.value(), head_num_kv, head_size);
    TORCH_CHECK(key_scale_hp.value().scalar_type() == torch::kFloat32,
                "key_scale_hp dtype must be float.");
    TORCH_CHECK(value_scale_hp.value().scalar_type() == torch::kFloat32,
                "value_scale_hp dtype must be float.");
    TORCH_CHECK(key_scale_hp.value().get_device() == origin_device_id,
                "key_scale_hp tensor device index is not same, original index: ", origin_device_id,
                "now index is: ", key_scale_hp.value().get_device());
    TORCH_CHECK(value_scale_hp.value().get_device() == origin_device_id,
                "value_scale_hp tensor device index is not same, original index: ",
                origin_device_id, "now index is: ", value_scale_hp.value().get_device());
    key_scale_hp_ptr = key_scale_hp.value().data_ptr();
    value_scale_hp_ptr = value_scale_hp.value().data_ptr();
  }

  void *key_scale_lp_ptr = nullptr;
  void *value_scale_lp_ptr = nullptr;
  if (mixed_cache) {
    TORCH_CHECK(key_scale_lp.value().dim() == 4,
                "key_scale_lp dim must be 4: (max_bs, head_num_kv, max_decode_len_lp, group_num),"
                "or (num_blocks, head_num_kv, block_size_lp, group_num).");
    TORCH_CHECK(key_scale_lp.value().scalar_type() == torch::kFloat32,
                "key_scale_hp dtype must be float.");
    TORCH_CHECK(value_scale_lp.value().scalar_type() == torch::kFloat32,
                "value_scale_hp dtype must be float.");
    TORCH_CHECK(key_scale_lp.value().get_device() == origin_device_id,
                "key_scale_hp tensor device index is not same, original index: ", origin_device_id,
                "now index is: ", key_scale_lp.value().get_device());
    TORCH_CHECK(value_scale_lp.value().get_device() == origin_device_id,
                "value_scale_hp tensor device index is not same, original index: ",
                origin_device_id, "now index is: ", value_scale_lp.value().get_device());
    group_size = head_size / key_scale_lp.value().size(-1);
    key_scale_lp_ptr = key_scale_lp.value().data_ptr();
    value_scale_lp_ptr = value_scale_lp.value().data_ptr();
  }

  // check cache_bs_id cache_seq_offsets
  const void *cache_bs_id_hp_ptr = nullptr;
  const void *cache_bs_id_lp_ptr = nullptr;
  bool has_cache_bs_id_hp = cache_bs_id_hp.has_value();
  bool has_cache_bs_id_lp = cache_bs_id_lp.has_value();
  if (has_cache_bs_id_hp) {
    CHECK_SHAPE(cache_bs_id_hp.value(), batch_size);
    TORCH_CHECK(cache_bs_id_hp.value().scalar_type() == torch::kInt32,
                "cache_bs_id_hp dtype need be torch::kInt32");
    TORCH_CHECK(cache_bs_id_hp.value().get_device() == origin_device_id,
                "cache_bs_id_hp tensor device index is not same, original index: ",
                origin_device_id, "now index is: ", cache_bs_id_hp.value().get_device());
    cache_bs_id_hp_ptr = cache_bs_id_hp.value().data_ptr();
  }
  if (mixed_cache && has_cache_bs_id_lp) {
    CHECK_SHAPE(cache_bs_id_lp.value(), batch_size);
    TORCH_CHECK(cache_bs_id_lp.value().scalar_type() == torch::kInt32,
                "cache_bs_id_lp dtype need be torch::kInt32");
    TORCH_CHECK(cache_bs_id_lp.value().get_device() == origin_device_id,
                "cache_bs_id_lp tensor device index is not same, original index: ",
                origin_device_id, "now index is: ", cache_bs_id_lp.value().get_device());
    cache_bs_id_lp_ptr = cache_bs_id_lp.value().data_ptr();
  }
  const void *slot_mapping_hp_ptr = nullptr;
  const void *slot_mapping_lp_ptr = nullptr;
  if (paged_cache_hp) {
    CHECK_SHAPE(slot_mapping_hp.value(), batch_size);
    TORCH_CHECK(slot_mapping_hp.value().scalar_type() == torch::kInt32,
                "slot_mapping_hp dtype need be torch::kInt32");
    TORCH_CHECK(slot_mapping_hp.value().get_device() == origin_device_id,
                "slot_mapping_hp tensor device index is not same, original index: ",
                origin_device_id, "now index is: ", slot_mapping_hp.value().get_device());
    slot_mapping_hp_ptr = slot_mapping_hp.value().data_ptr();
  }
  if (mixed_cache && paged_cache_lp) {
    CHECK_SHAPE(slot_mapping_lp.value(), batch_size);
    TORCH_CHECK(slot_mapping_lp.value().scalar_type() == torch::kInt32,
                "slot_mapping_lp dtype need be torch::kInt32");
    TORCH_CHECK(slot_mapping_lp.value().get_device() == origin_device_id,
                "slot_mapping_lp tensor device index is not same, original index: ",
                origin_device_id, "now index is: ", slot_mapping_lp.value().get_device());
    slot_mapping_lp_ptr = slot_mapping_lp.value().data_ptr();
  }
  const void *cache_seq_offsets_hp_ptr = nullptr;
  const void *cache_seq_offsets_lp_ptr = nullptr;
  if (!paged_cache_hp) {
    CHECK_SHAPE(cache_seq_offsets_hp.value(), batch_size);
    TORCH_CHECK(cache_seq_offsets_hp.value().dtype() == torch::kInt32,
                "cache_seq_offsets_hp dtype need be torch::kInt32");
    TORCH_CHECK(cache_seq_offsets_hp.value().get_device() == origin_device_id,
                "cache_seq_offsets_hp tensor device index is not same, original index:",
                origin_device_id, "now index is: ", cache_seq_offsets_hp.value().get_device());
    cache_seq_offsets_hp_ptr = cache_seq_offsets_hp.value().data_ptr();
  }
  if (mixed_cache && !paged_cache_lp) {
    CHECK_SHAPE(cache_seq_offsets_lp.value(), batch_size);
    TORCH_CHECK(cache_seq_offsets_lp.value().dtype() == torch::kInt32,
                "cache_seq_offsets_lp dtype need be torch::kInt32");
    TORCH_CHECK(cache_seq_offsets_lp.value().get_device() == origin_device_id,
                "cache_seq_offsets_lp tensor device index is not same, original index:",
                origin_device_id, "now index is: ", cache_seq_offsets_lp.value().get_device());
    cache_seq_offsets_lp_ptr = cache_seq_offsets_lp.value().data_ptr();
  }

  auto queue = torch_mlu::getCurrentMLUStream();
  auto data_type = getCnnlDataType(qkv.scalar_type());
  invokeFusedRope(queue, getAtTensorPtr(qkv), getAtTensorPtr(key_cache_hp),
                  getAtTensorPtr(value_cache_hp), key_cache_lp_ptr, value_cache_lp_ptr,
                  getAtTensorPtr(sin_table), getAtTensorPtr(cos_table),
                  getAtTensorPtr(position_ids), getAtTensorPtr(gamma), getAtTensorPtr(beta),
                  key_scale_hp_ptr, value_scale_hp_ptr, key_scale_lp_ptr, value_scale_lp_ptr,
                  cache_bs_id_hp_ptr, cache_seq_offsets_hp_ptr, cache_bs_id_lp_ptr,
                  cache_seq_offsets_lp_ptr, slot_mapping_hp_ptr, slot_mapping_lp_ptr, rotary_stride,
                  batch_size, head_num_q, head_num_kv, head_size, max_decode_len_hp,
                  max_decode_len_lp, block_size_hp, block_size_lp, group_size, data_type, eps);
}
}  // namespace torch_api
}  // namespace tmo
