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

namespace tmo {
namespace kernel_test_api {
void slice_cu_seq_lens_test(const at::Tensor &cu_seq_lens,
                            at::Tensor &sliced_cu_seq_lens,
                            int batch,
                            int parallel_num) {
  MLU_TENSOR_CHECK_FATAL(cu_seq_lens);
  MLU_TENSOR_CHECK_FATAL(sliced_cu_seq_lens);

  cnnlHandle_t handle = torch_mlu::getCurrentHandle();
  cnrtQueue_t queue;
  cnnlGetQueue(handle, &queue);

  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(invokeSliceCuSeqlens(queue, (int *)cu_seq_lens.data_ptr(),
                                              (int *)sliced_cu_seq_lens.data_ptr(), batch,
                                              parallel_num));
  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);
}

void generate_cu_seq_lens_test(at::Tensor &gen_cu_seq_lens,
                               int seq_len,
                               int parallel_num,
                               bool is_causal_mask,
                               bool is_kv_seq_len) {
  MLU_TENSOR_CHECK_FATAL(gen_cu_seq_lens);

  cnnlHandle_t handle = torch_mlu::getCurrentHandle();
  cnrtQueue_t queue;
  cnnlGetQueue(handle, &queue);

  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(invokeGenerateCuSeqlens(queue, (int *)gen_cu_seq_lens.data_ptr(), seq_len,
                                                 parallel_num, is_causal_mask, is_kv_seq_len));
  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);
}

}  // namespace kernel_test_api
}  // namespace tmo
