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

namespace tmo {
namespace torch_api {
void offline_quant_to_paged_cache(
    const torch::Tensor &k,                             // [num_tokens, num_heads, head_size]
    const c10::optional<torch::Tensor> &v,              // [num_tokens, num_heads, head_size]
    const torch::Tensor &k_cache_scale,                 // [num_heads, head_size]
    const c10::optional<torch::Tensor> &v_cache_scale,  // [num_heads, head_size]
    const torch::Tensor &slot_mapping,
    torch::Tensor &k_cache,  // [num_blocks, num_heads, block_size, head_size]
    const c10::optional<torch::Tensor>
        &v_cache) {  // [num_blocks, num_heads, block_size, head_size]
  // 1. check device and tensor type
  TORCH_CHECK(slot_mapping.dtype() == torch::kInt32, "slot_mapping type need be torch::kInt32");
  checkTensorSameAttr<TensorAttr::DEVICE>(k, v, k_cache_scale, v_cache_scale, slot_mapping);

  // 2. check dim and shape
  TORCH_CHECK(k.dim() == 3, "dim of k must be 3");
  TORCH_CHECK(k_cache.dim() == 4, "dim of k_cache must be 4");
  TORCH_CHECK(k_cache_scale.dim() == 2, "dim of k_cache_scale must be 2");
  TORCH_CHECK(
      v.has_value() == v_cache.has_value() && v.has_value() == v_cache_scale.has_value(),
      "v.has_value() == v_cache.has_value() && v.has_value() == v_cache_scale.has_value().");
  if (v.has_value()) {
    TORCH_CHECK(v.value().dim() == 3, "dim of v must be 3");
    TORCH_CHECK(v_cache.value().dim() == 4, "dim of v_cache must be 4");
    TORCH_CHECK(v_cache_scale.value().dim() == 2, "dim of v_cache_scale must be 2");
  }
  TORCH_CHECK(slot_mapping.dim() == 1, "dim of slot_mapping must be 1");

  const int num_tokens = k.size(0);
  const int num_heads = k.size(1);
  const int head_size = k.size(2);
  const int block_size = k_cache.size(2);
  const int num_blocks = k_cache.size(0);

  CHECK_SHAPE(k, num_tokens, num_heads, head_size);
  CHECK_SHAPE(k_cache_scale, num_heads, head_size);
  CHECK_SHAPE(k_cache, num_blocks, num_heads, block_size, head_size);
  if (v.has_value()) {
    CHECK_SHAPE(v.value(), num_tokens, num_heads, head_size);
    CHECK_SHAPE(v_cache_scale.value(), num_heads, head_size);
    CHECK_SHAPE(v_cache.value(), num_blocks, num_heads, block_size, head_size);
  }
  CHECK_SHAPE(slot_mapping, num_tokens);

  // 3. check strides
  TORCH_CHECK(slot_mapping.is_contiguous(), "slot_mapping need be contiguous.");
  TORCH_CHECK(k.stride(-1) == 1, "k last dim must be contiguous.");
  TORCH_CHECK(k.stride(-2) == head_size, "k second dim must be contiguous.");
  if (v.has_value()) {
    TORCH_CHECK(v.value().stride(-1) == 1, "v last dim must be contiguous.");
    TORCH_CHECK(v.value().stride(-2) == head_size, "v second dim must be contiguous.");
  }

  int64_t kv_cache_range =
      (int64_t)num_blocks * num_heads * block_size * head_size * k_cache.element_size();
  TORCH_CHECK(kv_cache_range <= UINT32_MAX, "The addressing range of kv_cache cannot exceed 4G.");

  const torch_mlu::mlu::MLUGuard device_guard(k.device());
  auto queue = torch_mlu::getCurMLUStream();
  cnnlDataType_t dtype = getCnnlDataType(k.scalar_type());
  int32_t key_stride0 = static_cast<int32_t>(k.stride(0));
  int32_t value_stride0 = 1;
  if (v.has_value()) {
    value_stride0 = static_cast<int32_t>(v.value().stride(0));
  }
  TMO_KERNEL_CHECK_FATAL(tmo::invokeOfflineQuantToPagedCache(
      queue, dtype, getAtTensorPtr(k), getAtTensorPtr(v), getAtTensorPtr(k_cache),
      getAtTensorPtr(v_cache), getAtTensorPtr(k_cache_scale), getAtTensorPtr(v_cache_scale),
      getAtTensorPtr(slot_mapping), key_stride0, value_stride0, num_tokens, num_heads, num_blocks,
      block_size, head_size));
}
}  // namespace torch_api
}  // namespace tmo
