/*************************************************************************
 * 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/create_cos_sin_table.mluh"
#include <cnnl.h>
#include <torch/extension.h>
#include "register_api.h"
#include "utils.h"

namespace {
constexpr int DYNAMIC_NTK_SCALING = 2;
}
namespace tmo {
namespace kernel_test_api {
at::Tensor create_cos_sin_table_test(at::Tensor &rotary_emb_alpha_cached,
                                     const at::Tensor &seq_lens,
                                     const int max_position_embeddings,
                                     const int batch_stride,
                                     const int rotary_seq_len,
                                     const int rotary_dim,
                                     const int rotary_stride,
                                     const float rotary_scaling,
                                     const int rotary_scaling_type,
                                     const float rotary_base,
                                     const bool interleaved) {
  MLU_TENSOR_CHECK_FATAL(rotary_emb_alpha_cached);
  MLU_TENSOR_CHECK_FATAL(seq_lens);

  cnnlHandle_t handle = torch_mlu::getCurrentHandle();
  cnrtQueue_t queue;
  cnnlGetQueue(handle, &queue);
  cnnlDataType_t dtype =
      torch2cnnlDataType(torch::typeMetaToScalarType(rotary_emb_alpha_cached.dtype()));

  std::shared_ptr<torch::Device> dev =
      std::make_shared<torch::Device>(rotary_emb_alpha_cached.device());
  int batch = seq_lens.size(0);

  torch::Tensor cos_sin_table =
      rotary_scaling_type == DYNAMIC_NTK_SCALING
          ? torch::zeros({batch, rotary_seq_len, rotary_stride},
                         torch::dtype(torch::typeMetaToScalarType(rotary_emb_alpha_cached.dtype()))
                             .device(*dev)
                             .requires_grad(false))
          : torch::zeros({rotary_seq_len, rotary_stride},
                         torch::dtype(torch::typeMetaToScalarType(rotary_emb_alpha_cached.dtype()))
                             .device(*dev)
                             .requires_grad(false));

  size_t io_bytes = 0;

  cnrtNotifier_t time_begin;
  cnrtNotifier_t time_end;

  CNRT_CHECK(cnrtNotifierCreate(&time_begin));
  CNRT_CHECK(cnrtNotifierCreate(&time_end));

  CNRT_CHECK(cnrtPlaceNotifier(time_begin, queue));
  TMO_KERNEL_CHECK_FATAL(invokeCreateCosSinTable(
      queue, cos_sin_table.data_ptr(),
      reinterpret_cast<float *>(rotary_emb_alpha_cached.data_ptr()),
      reinterpret_cast<int *>(seq_lens.data_ptr()), max_position_embeddings, batch, batch_stride,
      rotary_seq_len, rotary_dim, rotary_stride, rotary_base, rotary_scaling, rotary_scaling_type,
      interleaved, dtype));
  CNRT_CHECK(cnrtPlaceNotifier(time_end, queue));
  cnrtQueueSync(queue);
  float usec = 0.0f;
  CNRT_CHECK(cnrtNotifierDuration(time_begin, time_end, &usec));
  print_info(usec, io_bytes);

  return cos_sin_table;
}

}  // namespace kernel_test_api
}  // namespace tmo
