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

namespace tmo {
namespace torch_api {
void apply_rotary(const torch::Tensor &input,
                  const torch::Tensor &sin_cache,
                  const torch::Tensor &cos_cache,
                  const c10::optional<torch::Tensor> &position_ids,
                  const c10::optional<torch::Tensor> &cu_seqlens,
                  bool interleaved,
                  bool discrete,
                  bool dynamic_ntk,
                  int64_t max_seqlen) {
  auto output = input;
  // 1.check device and tensor type
  checkTensorSameAttr<TensorAttr::ALL>(input, sin_cache, cos_cache);

  const bool has_position_ids = position_ids.has_value();
  const bool has_cu_seqlens = cu_seqlens.has_value();
  const int origin_device_id = input.get_device();
  const void *position_ids_ptr = nullptr;
  const void *cu_seqlens_ptr = nullptr;

  if (has_position_ids) {
    TORCH_CHECK(position_ids.value().dtype() == torch::kInt32,
                "position_ids type need be torch::kInt32");
    TORCH_CHECK(position_ids.value().get_device() == origin_device_id,
                "Tensor device index is not the same, original index: ", origin_device_id,
                "now index is: ", position_ids.value().get_device());
    position_ids_ptr = position_ids.value().data_ptr();
  }

  if (has_cu_seqlens) {
    TORCH_CHECK(cu_seqlens.value().dtype() == torch::kInt32,
                "cu_seqlens type need be torch::kInt32");
    TORCH_CHECK(cu_seqlens.value().get_device() == origin_device_id,
                "Tensor device index is not the same, original index: ", origin_device_id,
                "now index is: ", cu_seqlens.value().get_device());
    cu_seqlens_ptr = cu_seqlens.value().data_ptr();
  }

  // 2. check shape
  int total_seqlen = 0;
  int batch_size = 0;
  int head_size = input.size(-1);
  if (input.dim() == 3) {  // pack mode
    TORCH_CHECK(has_cu_seqlens,
                "input has 3 dims: (total_seq_len, head_num, head_size),"
                " which means pack mode, cu_seqlens should not be None");
    total_seqlen = input.size(0);
    batch_size = cu_seqlens.value().size(0) - 1;
  } else if (input.dim() == 4) {
    TORCH_CHECK(!has_cu_seqlens,
                "input has 4 dims: (batch_size, seq_len, head_num, head_size),"
                " which means pad mode, cu_seqlens should be None");
    TORCH_CHECK(max_seqlen == input.size(1),
                "input has 4 dims: (batch_size, seq_len, head_num, head_size),"
                " which means pad mode, max_seqlen must be equtals to input.size(1)");
    batch_size = input.size(0);
    total_seqlen = batch_size * input.size(1);
  } else {
    TORCH_CHECK(false, "input only support 3 or 4 dims");
  }
  TORCH_CHECK(head_size <= 256, "only support input head_size <= 256");

  const int rope_seqlen = dynamic_ntk ? sin_cache.size(1) : sin_cache.size(0);
  const int rope_dim = dynamic_ntk ? sin_cache.size(2) : sin_cache.size(1);

  if (has_position_ids) {
    if (discrete) {
      CHECK_SHAPE(position_ids.value(), total_seqlen);
    } else {
      CHECK_SHAPE(position_ids.value(), batch_size);
    }
  } else {
    TORCH_CHECK(!discrete, "discrete must be false if position ids is null.")
  }

  if (!(has_position_ids && discrete)) {
    TORCH_CHECK(max_seqlen <= rope_seqlen, "max_seqlen must less than or equal to rope_seqlen.")
  }

  if (dynamic_ntk) {
    CHECK_SHAPE(sin_cache, batch_size, rope_seqlen, rope_dim);
    CHECK_SHAPE(cos_cache, batch_size, rope_seqlen, rope_dim);
  } else {
    CHECK_SHAPE(sin_cache, rope_seqlen, rope_dim);
    CHECK_SHAPE(cos_cache, rope_seqlen, rope_dim);
  }

  // 3. check strides
  TORCH_CHECK(input.stride(-1) == 1, "input last dim must be contiguous");

  if (dynamic_ntk) {
    TORCH_CHECK(sin_cache.stride(1) == cos_cache.stride(1),
                "sin_cache second stride must be equal to cos_cache second stride");
  } else {
    TORCH_CHECK(sin_cache.stride(0) == cos_cache.stride(0),
                "sin_cache first stride must be equal to cos_cache second stride");
  }

  if (has_position_ids) {
    TORCH_CHECK(position_ids.value().is_contiguous(), "position_ids must be contiguous");
  }

  if (has_cu_seqlens) {
    TORCH_CHECK(cu_seqlens.value().is_contiguous(), "cu_seqlens must be contiguous");
  }

  // prepare inputs
  auto dims = input.dim();
  const int64_t num_heads = input.size(dims - 2);
  const int64_t head_dim = input.size(dims - 1);
  const int64_t input_seq_stride = input.strides()[dims - 3];
  const int64_t input_head_stride = input.strides()[dims - 2];
  const int64_t output_seq_stride = output.strides()[dims - 3];
  const int64_t output_head_stride = output.strides()[dims - 2];

  const torch_mlu::mlu::MLUGuard device_guard(input.device());
  auto queue = torch_mlu::getCurMLUStream();
  auto data_type = getCnnlDataType(input.scalar_type());

  invokeRotaryEmbedding(queue, output.data_ptr(), input.data_ptr(), sin_cache.data_ptr(),
                        cos_cache.data_ptr(), (int *)position_ids_ptr, (int *)cu_seqlens_ptr,
                        batch_size, max_seqlen, num_heads, head_dim, rope_seqlen, rope_dim,
                        dynamic_ntk ? sin_cache.strides()[1] : sin_cache.strides()[0],
                        input_seq_stride, input_head_stride, output_seq_stride, output_head_stride,
                        interleaved, discrete, dynamic_ntk, data_type);
}
}  // namespace torch_api
}  // namespace tmo
