/*************************************************************************
 * 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 "torch_ops_api.h"
#include "utils.h"
namespace tmo {
namespace torch_api {

// q_ori, out Tensor shape is [batch, seq_len, head_num, head_size]
// if (pagedattn) k_cache, v_cache shape is [num_blocks, k_head_num, block_size, head_size]
//    else [max_bs, head_num, max_seq_len, head_size]
// if(pagedattn)  k_cache_quant_scale shape is [num_blocks, k_head_num, block_size] else [max_bs,
// head_num, max_seq_len]
// if (pagedattn) block_tables shape is [bs, max_nblock_per_seq] else [bs, 1]
// context_lens: true seqkv_len of each bacth
void single_query_cached_kv_attn(
    const torch::Tensor &q_ori,
    const torch::Tensor &k_cache,
    const torch::Tensor &v_cache,
    const torch::Tensor &output,
    const torch::Tensor &block_tables,
    const torch::Tensor &context_lens,  // [batch]
    const c10::optional<torch::Tensor> &output_lse,
    const c10::optional<torch::Tensor> &k_cache_quant_scale,
    const c10::optional<torch::Tensor> &v_cache_quant_scale,
    const c10::optional<torch::Tensor> &alibi_slopes,  // [bs, head_num] or [head_num]
    int64_t max_context_len,
    int64_t windows_size_left,
    int64_t windows_size_right,
    double softmax_scale,
    bool return_lse,
    int64_t kv_cache_quant_bit_size) {
  // Check tensor type and tensor device.
  cnnlQuantizeLayout_t cache_quant_layout = CNNL_QUANTIZE_NONE;
  bool is_kv_quant = k_cache_quant_scale.has_value();
  bool has_alibi = alibi_slopes.has_value();

  // Check Tensor Shape
  int batch = q_ori.size(0);
  int seq_q = q_ori.size(1);
  int head_num = q_ori.size(2);
  int qk_head_size = q_ori.size(3);
  int k_head_num = k_cache.size(1);
  int v_head_size = v_cache.size(-1);
  int max_num_block_per_seq = block_tables.size(1);
  int num_blocks = k_cache.size(0);
  int block_size = k_cache.size(2);
  TORCH_CHECK(windows_size_right < 0, "only support windows_size_right < 0 currently.");
  TORCH_CHECK(head_num % k_head_num == 0, "num_heads need be mutiple of num_kv_heads.");
  TORCH_CHECK(
      kv_cache_quant_bit_size == 4 || kv_cache_quant_bit_size == 8 || kv_cache_quant_bit_size == -1,
      "illegal quant bit size, only support 4, 8 or -1.");
  CHECK_SHAPE(block_tables, batch, max_num_block_per_seq);
  if (kv_cache_quant_bit_size == 4) {
    int v_cache_len = PAD_UP_DIV(block_size, 2);
    CHECK_SHAPE(k_cache, num_blocks, k_head_num, block_size, qk_head_size / 2);
    CHECK_SHAPE(v_cache, num_blocks, k_head_num, v_cache_len, v_head_size);
  } else {
    CHECK_SHAPE(k_cache, num_blocks, k_head_num, block_size, qk_head_size);
    CHECK_SHAPE(v_cache, num_blocks, k_head_num, block_size, v_head_size);
  }
  TORCH_CHECK(q_ori.stride(-1) == 1 && q_ori.stride(-2) == qk_head_size,
              "q last two dim need be contiguous.");
  TORCH_CHECK(k_cache.is_contiguous() && v_cache.is_contiguous(),
              "k_cache and v_cache need be contiguous.");
  TORCH_CHECK(isMlu(q_ori), "q_ori need be mlu tensor.");

  checkTensorSameAttr<TensorAttr::DEVICE>(q_ori, k_cache, v_cache, output, block_tables,
                                          context_lens, k_cache_quant_scale, v_cache_quant_scale,
                                          output_lse, alibi_slopes);
  if (is_kv_quant) {
    TORCH_CHECK(k_cache_quant_scale.value().dim() == 2 || k_cache_quant_scale.value().dim() == 3 ||
                    k_cache_quant_scale.value().dim() == 4,
                "k_cache_quant_scale must be 2d or 3d or 4d.");
    TORCH_CHECK(k_cache_quant_scale.value().dim() == v_cache_quant_scale.value().dim(),
                "the dim of k_cache_quant_scale and v_cache_quant_scale must be euqal.");
    if (k_cache_quant_scale.value().dim() == 2) {
      CHECK_SHAPE(k_cache_quant_scale.value(), k_head_num, qk_head_size);
      CHECK_SHAPE(v_cache_quant_scale.value(), k_head_num, v_head_size);
      cache_quant_layout = CNNL_QUANTIZE_PER_CHANNEL;
    } else if (k_cache_quant_scale.value().dim() == 3) {
      CHECK_SHAPE(k_cache_quant_scale.value(), num_blocks, k_head_num, block_size);
      CHECK_SHAPE(v_cache_quant_scale.value(), num_blocks, k_head_num, block_size);
      cache_quant_layout = CNNL_QUANTIZE_PER_TOKEN;
    } else {
      CHECK_SHAPE(k_cache_quant_scale.value(), num_blocks, k_head_num, block_size, 1);
      CHECK_SHAPE(v_cache_quant_scale.value(), num_blocks, k_head_num, block_size, 1);
      cache_quant_layout = CNNL_QUANTIZE_PER_TOKEN;
    }
  }
  TORCH_CHECK(
      block_tables.scalar_type() == torch::kInt32 || block_tables.scalar_type() == torch::kLong,
      "block_tables type need be torch::kInt32 or torch::kLong.");
  // Check context_lens
  TORCH_CHECK(context_lens.dtype() == torch::kInt32, "context_lens type need be torch::kInt32.");
  TORCH_CHECK(context_lens.is_contiguous(), "context_lens need be contiguous.");

  if (has_alibi) {
    CHECK_SHAPE(alibi_slopes.value(), batch, head_num);
  }
  if (return_lse) {
    TORCH_CHECK(seq_q == 1, "return lse only support seq_q = 1 currently.");
    CHECK_SHAPE(output_lse.value(), batch, head_num, seq_q);
  }
  if (windows_size_left >= 0) {
    windows_size_left = windows_size_left + 1;  // would be removed next version
  }

  // Convert torch tensor to tensor descs
  auto descs = createTensorDescs(
      {q_ori, k_cache, v_cache, k_cache_quant_scale.value_or(at::Tensor()),
       v_cache_quant_scale.value_or(at::Tensor()), context_lens, block_tables,
       alibi_slopes.value_or(at::Tensor()), output, output_lse.value_or(at::Tensor())});

  if (kv_cache_quant_bit_size == 4) {
    cnnlDataType_t data_type = CNNL_DTYPE_INT4X2;
    // k_cache
    CNNL_CHECK_FATAL(cnnlSetTensorDescriptor_v2(descs[1].get(), CNNL_LAYOUT_ARRAY, data_type,
                                                k_cache.sizes().size(), k_cache.sizes().data()));
    // v_cache
    CNNL_CHECK_FATAL(cnnlSetTensorDescriptor_v2(descs[2].get(), CNNL_LAYOUT_ARRAY, data_type,
                                                v_cache.sizes().size(), v_cache.sizes().data()));
  }

  // Get current handle.
  const torch_mlu::mlu::MLUGuard device_guard(q_ori.device());
  cnnlHandle_t handle = torch_mlu::getCurrentHandle();
  // Get workspace size and malloc workspace.
  size_t workspace_size = 0;
  cnnlSingleQueryCachedKVAttnDescriptor_t att_desc;
  CNNL_CHECK_FATAL(cnnlGetSingleQueryCachedKVAttnWorkspaceSize_v2(
      handle, att_desc, descs[0].get(), descs[1].get(), descs[2].get(), (int)max_context_len,
      &workspace_size));
  auto workspace =
      at::empty({static_cast<int64_t>(workspace_size)}, q_ori.options().dtype(at::kByte));

  cnnlDataType_t data_dtype = getCnnlDataType(q_ori.scalar_type());
  SingleQueryCachedKvAttnTheory obj(batch, seq_q, head_num, k_head_num, num_blocks, block_size,
                                    qk_head_size, v_head_size, kv_cache_quant_bit_size,
                                    cache_quant_layout, data_dtype);
  cnpxPush(obj);
  // call cnnl extra op.
  CNNL_CHECK_FATAL(cnnlSingleQueryCachedKVAttn_v3(
      handle, att_desc, descs[0].get(), getAtTensorPtr(q_ori), descs[1].get(),
      getAtTensorPtr(k_cache), descs[2].get(), getAtTensorPtr(v_cache), descs[3].get(),
      getAtTensorPtr(k_cache_quant_scale), descs[4].get(), getAtTensorPtr(v_cache_quant_scale),
      descs[5].get(), getAtTensorPtr(context_lens), descs[6].get(), getAtTensorPtr(block_tables),
      descs[7].get(), getAtTensorPtr(alibi_slopes), cache_quant_layout, (int)max_context_len,
      windows_size_left, windows_size_right, softmax_scale, return_lse, getAtTensorPtr(workspace),
      workspace_size, descs[8].get(), getAtTensorPtr(output), descs[9].get(),
      getAtTensorPtr(output_lse)));
  cnpxPop();
}

}  // namespace torch_api
}  // namespace tmo
