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

namespace tmo {
namespace torch_api {
void reshape_linear_cache(const at::Tensor &key,
                          const c10::optional<at::Tensor> &value,
                          at::Tensor &key_cache,
                          const c10::optional<at::Tensor> &value_cache,
                          const at::Tensor &context_lengths,
                          const int64_t max_context_len,
                          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 device and dtype
  checkTensorSameAttr<TensorAttr::ALL>(key, key_cache, value, value_cache);

  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(1);

  const int context_bs_stride = packed ? 0 : key.stride(0);
  const int context_head_stride = packed ? key.stride(1) : key.stride(2);
  const int context_seq_stride = packed ? key.stride(0) : key.stride(1);
  const int cache_bs_stride = key_cache.stride(0);
  const int cache_head_stride = key_cache.stride(1);
  const int cache_seq_stride = key_cache.stride(2);

  TORCH_CHECK(context_lengths.scalar_type() == torch::kInt32,
              "context_lengths type need be torch::kInt32");

  if (value_cache.has_value()) {
    CHECK_SHAPE(value_cache.value(), max_bs, num_heads, cache_mem_len, head_size);
  }

  const int64_t device_id = key.get_device();
  bool has_context_seq_offset = context_seq_offset.has_value();
  const void *context_seq_offset_ptr = nullptr;
  if (has_context_seq_offset) {
    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();
  }

  bool has_cache_bs_id = cache_bs_id.has_value();
  const void *cache_bs_id_ptr = nullptr;
  if (has_cache_bs_id) {
    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();
  }

  bool has_cache_seqlen_offset = cache_seqlen_offset.has_value();
  const void *cache_seqlen_offset_ptr = nullptr;
  if (has_cache_seqlen_offset) {
    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();
  }

  // check contiguous
  TORCH_CHECK(key_cache.is_contiguous(), "key_cache tensor must be contiguous.");
  if (value_cache.has_value()) {
    TORCH_CHECK(value_cache.value().is_contiguous(), "value_cache tensor must be contiguous.");
  }

  // check shape
  if (packed) {
    CHECK_SHAPE(key, total_seqlen, num_heads, head_size);
    if (value.has_value()) {
      CHECK_SHAPE(value.value(), 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);
    if (value.has_value()) {
      CHECK_SHAPE(value.value(), batch_size, key.size(1), num_heads, head_size);
    }
    CHECK_SHAPE(context_lengths, batch_size);
  }

  if (value.has_value()) {
    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 axis ", i);
    }
  }

  TORCH_CHECK(key_cache.size(0) >= batch_size,
              "first dim of key_cache must be great than or equal to batch_size");
  TORCH_CHECK(key_cache.dim() == 4, "dim of key_cache must be 4.");

  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(invokeReshapeLinearCache(
      queue, getAtTensorPtr(key_cache), getAtTensorPtr(value_cache), 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, packed));
}
}  // namespace torch_api
}  // namespace tmo
