#include <algorithm>
#include "cnrt.h"
#include "operate_cu_seq_lens.mluh"

namespace {
constexpr int pair_elem_num = 2;
}

namespace tmo {
namespace kernels {
#define ONCHIP_DATA_NUM ((int)((__MLU_NRAM_SIZE__ * 1024 - 32 * 1024) / sizeof(int)))
__nram__ int nram_buffer[ONCHIP_DATA_NUM];
__nram__ const int acc_seq_lens[16] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};

__mlu_func__ void genSeqLens(int *seq_len_nram, int start, int multi, int elem_count) {
  constexpr int acc_seq_lens_size = 16;
  int count = std::min(acc_seq_lens_size, elem_count);
  int add_on = multi * acc_seq_lens_size;
  __bang_mul_scalar(seq_len_nram, acc_seq_lens, multi, count);
  __bang_add_scalar(seq_len_nram, seq_len_nram, start, count);
  while (count < elem_count) {
    __bang_add_scalar(seq_len_nram + count, seq_len_nram, add_on,
                      std::min(count, elem_count - count));
    count *= 2;
    add_on *= 2;
  }
}

__mlu_global__ void MLUSliceCuSeqlens(int *cu_seq_lens,
                                      int *sliced_cu_seq_lens,
                                      int batch,
                                      int every,
                                      int remain,
                                      int loop) {
  int cu_seq_lens_elem_count = batch + 1;
  int sliced_cu_seq_lens_elem_count = batch + loop;

  int *cu_seq_lens_narm = nram_buffer;
  int *sliced_cu_seq_lens_narm = cu_seq_lens_narm + cu_seq_lens_elem_count;
  int *sliced_cu_seq_lens_narm_start = sliced_cu_seq_lens_narm;

  __memcpy(cu_seq_lens_narm, cu_seq_lens, cu_seq_lens_elem_count * sizeof(int), GDRAM2NRAM);
  __bang_write_zero(sliced_cu_seq_lens_narm, sliced_cu_seq_lens_elem_count);
  for (int i = 0; i < loop; ++i) {
    int elem_num = 1 + (i == loop - 1 && remain != 0 ? remain : every);
    __bang_sub_scalar(sliced_cu_seq_lens_narm, cu_seq_lens_narm, cu_seq_lens_narm[0], elem_num);
    cu_seq_lens_narm += elem_num - 1;
    sliced_cu_seq_lens_narm += elem_num;
  }
  __memcpy(sliced_cu_seq_lens, sliced_cu_seq_lens_narm_start,
           sliced_cu_seq_lens_elem_count * sizeof(int), NRAM2GDRAM);
}

__mlu_global__ void MLUGenerateKVCuSeqlens(int *gen_cu_seq_lens,
                                           int every,
                                           int remain,
                                           int loop,
                                           int seq_len,
                                           bool is_causal_mask,
                                           int seg_data_num,
                                           int task_num) {
  int offset = seg_data_num * taskIdX;
  int total_elem_num = std::min(seg_data_num, loop * pair_elem_num - offset);
  int seq_len_elem_num = total_elem_num / pair_elem_num;

  int *gen_cu_seq_lens_narm = nram_buffer;
  __bang_write_zero(gen_cu_seq_lens_narm, total_elem_num);

  if (is_causal_mask) {
    int *seq_lens_narm = gen_cu_seq_lens_narm + total_elem_num;
    genSeqLens(seq_lens_narm, every * offset / pair_elem_num, every, seq_len_elem_num);
    __memcpy(gen_cu_seq_lens_narm + 1, seq_lens_narm, sizeof(int), NRAM2NRAM,
             pair_elem_num * sizeof(int), sizeof(int), seq_len_elem_num - 1);
    if (remain != 0 && taskIdX == task_num - 1) {
      gen_cu_seq_lens_narm[total_elem_num - 1] -= (every - remain);
    }
  } else {
    __bang_write_value(gen_cu_seq_lens_narm + 1, 1, seq_len, pair_elem_num * sizeof(int),
                       seq_len_elem_num - 1, seq_len_elem_num * pair_elem_num * sizeof(int), 0);
  }
  __memcpy(gen_cu_seq_lens + offset, gen_cu_seq_lens_narm, total_elem_num * sizeof(int),
           NRAM2GDRAM);
}

__mlu_global__ void MLUGenerateQCuSeqlens(int *gen_cu_seq_lens,
                                          int every,
                                          int remain,
                                          int loop,
                                          int seg_data_num,
                                          int task_num) {
  int offset = seg_data_num * taskIdX;
  int total_elem_num = std::min(seg_data_num, loop * pair_elem_num - offset);
  int seq_len_elem_num = total_elem_num / pair_elem_num;

  int *gen_cu_seq_lens_narm = nram_buffer;

  __bang_write_zero(gen_cu_seq_lens_narm, total_elem_num);
  __bang_write_value(gen_cu_seq_lens_narm + 1, 1, every, pair_elem_num * sizeof(int),
                     seq_len_elem_num - 1, seq_len_elem_num * pair_elem_num * sizeof(int), 0);
  if (remain != 0 && taskIdX == task_num - 1) {
    gen_cu_seq_lens_narm[total_elem_num - 1] = remain;
  }
  __memcpy(gen_cu_seq_lens + offset, gen_cu_seq_lens_narm, total_elem_num * sizeof(int),
           NRAM2GDRAM);
}
}  // namespace kernels

KernelStatus invokeSliceCuSeqlens(cnrtQueue_t queue,
                                  int *cu_seq_lens,
                                  int *sliced_cu_seq_lens,
                                  int batch,
                                  int parallel_num) {
  int every = (batch + parallel_num - 1) / parallel_num;
  int repeat = batch / every;
  int remain = batch % every;
  int loop = repeat + (remain != 0);

  cnrtDim3_t dim{1, 1, 1};
  kernels::MLUSliceCuSeqlens<<<dim, cnrtFuncTypeBlock, queue>>>(cu_seq_lens, sliced_cu_seq_lens,
                                                                batch, every, remain, loop);
  return KernelStatus::KERNEL_STATUS_SUCCESS;
}

KernelStatus invokeGenerateCuSeqlens(cnrtQueue_t queue,
                                     int *gen_cu_seq_lens,
                                     int seq_len,
                                     int parallel_num,
                                     bool is_causal_mask,
                                     bool is_kv_seq_len) {
  int every = (seq_len + parallel_num - 1) / parallel_num;
  int repeat = seq_len / every;
  int remain = seq_len % every;
  int loop = repeat + (remain != 0);

  int seg_data_num = ONCHIP_DATA_NUM / 2;
  if (is_kv_seq_len && is_causal_mask) {
    // max segnum for 2d memcpy is 64k
    seg_data_num = std::min(pair_elem_num * 64 * 1024, seg_data_num);
  }
  int total_elem_num = loop * pair_elem_num;
  int task_num = (total_elem_num + seg_data_num - 1) / seg_data_num;

  cnrtDim3_t dim{(unsigned int)task_num, 1, 1};

  if (is_kv_seq_len) {
    kernels::MLUGenerateKVCuSeqlens<<<dim, cnrtFuncTypeBlock, queue>>>(
        gen_cu_seq_lens, every, remain, loop, seq_len, is_causal_mask, seg_data_num, task_num);
  } else {
    kernels::MLUGenerateQCuSeqlens<<<dim, cnrtFuncTypeBlock, queue>>>(
        gen_cu_seq_lens, every, remain, loop, seg_data_num, task_num);
  }

  return KernelStatus::KERNEL_STATUS_SUCCESS;
}

}  // namespace tmo
