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

namespace tmo {
namespace torch_api {
void reshape_paged_cache(
    const torch::Tensor &k,                       // [num_tokens, num_heads, head_dim]
    const c10::optional<torch::Tensor> &v,        // [num_tokens, num_heads, head_dim]
    torch::Tensor &k_cache,                       // [num_blocks, num_heads, block_size, head_dim]
    const c10::optional<torch::Tensor> &v_cache,  // [num_blocks, num_heads, block_size, head_dim]
    const torch::Tensor &slot_mapping) {
  // 1. check device and tensor type
  checkTensorSameAttr<TensorAttr::ALL>(k, v, k_cache, v_cache);
  TORCH_CHECK(slot_mapping.dtype() == torch::kInt32 || slot_mapping.dtype() == torch::kLong,
              "slot_mapping type need be torch::kInt32 or torch::Long");
  TORCH_CHECK(slot_mapping.get_device() == k.get_device(),
              "Tensor device index is not same, original index: ", k.get_device(),
              "now index is: ", slot_mapping.get_device());

  // 2. check shape
  const int num_tokens = k.size(0);
  const int num_heads = k.size(1);
  const int head_dim = k.size(2);
  const int block_size = k_cache.size(2);
  const int head_size = k_cache.size(3);
  const int num_blocks = k_cache.size(0);

  CHECK_SHAPE(k, num_tokens, num_heads, head_dim);
  CHECK_SHAPE(k_cache, num_blocks, num_heads, block_size, head_dim);
  TORCH_CHECK(v.has_value() == v_cache.has_value(), "v.has_value() == v_cache.has_value().")
  if (v.has_value()) {
    CHECK_SHAPE(v.value(), num_tokens, num_heads, head_dim);
    CHECK_SHAPE(v_cache.value(), num_blocks, num_heads, block_size, head_dim);
  }
  CHECK_SHAPE(slot_mapping, num_tokens);

  // 3. check strides
  TORCH_CHECK(k_cache.is_contiguous(), "k_cache need be contiguous.");
  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_dim, "k second dim must be contiguous.");
  if (v.has_value()) {
    TORCH_CHECK(v_cache.value().is_contiguous(), "v_cache need be contiguous.");
    TORCH_CHECK(v.value().stride(-1) == 1, "v last dim must be contiguous.");
    TORCH_CHECK(v.value().stride(-2) == head_dim, "v second dim must be contiguous.");
  }

  // check large tensor
  int64_t kv_cache_range =
      (int64_t)num_blocks * num_heads * block_size * head_dim * 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::invokeReshapePagedCache(
      queue, dtype, getAtTensorPtr(k), getAtTensorPtr(v), getAtTensorPtr(k_cache),
      getAtTensorPtr(v_cache), getAtTensorPtr(slot_mapping), key_stride0, value_stride0, num_tokens,
      num_heads, num_blocks, block_size, head_size));
}
}  // namespace torch_api
}  // namespace tmo
