#include <stdexcept>
#include "cnrt.h"
#include "get_glm_position_id.mluh"
// clang-format off
#include <mlu.h>
// clang-format on

namespace tmo {
namespace kernels {
__nram__ int nram_buffer[__MLU_NRAM_SIZE__ * 3 / 4 * 1024 / sizeof(int)];

__mlu_global__ void MLUBlockSliceIndividualPosId(int *context_pos_id,
                                                 int *generate_pos_id,
                                                 int batch,
                                                 int context_seq_len,
                                                 int pos_id_dimension /* 1 for 1D, 2 for 2D */) {
  if (taskId != 0) return;
  __memcpy(nram_buffer, context_pos_id + context_seq_len - 1, sizeof(int), GDRAM2NRAM, sizeof(int),
           context_seq_len * sizeof(int), pos_id_dimension * batch - 1);
  if (pos_id_dimension == 2) {
    for (int i = 1; i < 2 * batch; i += 2) {
      nram_buffer[i] += 1;
    }
  }
  __memcpy(generate_pos_id, nram_buffer, pos_id_dimension * batch * sizeof(int), NRAM2GDRAM);
}

__mlu_global__ void MLUBlockIncrement2DPosId(int *generate_pos_id, int batch) {
  if (taskId != 0) return;
  __memcpy(nram_buffer, generate_pos_id, 2 * batch * sizeof(int), GDRAM2NRAM);
  for (int i = 1; i < 2 * batch; i += 2) {
    nram_buffer[i] += 1;
  }
  __memcpy(generate_pos_id, nram_buffer, 2 * batch * sizeof(int), NRAM2GDRAM);
}
}  // namespace kernels
KernelStatus invokeSliceIndividualPosId(cnrtQueue_t queue,
                                        int *context_pos_id,
                                        int *generate_pos_id,
                                        int batch,
                                        int context_seq_len,
                                        int pos_id_dimension /* 1 for 1D, 2 for 2D */) {
  if (pos_id_dimension != 1 && pos_id_dimension != 2) {
    std::cerr << "[invokeSliceIndividualPosId]: pos_id_dimension must be 1 or 2, but got "
              << pos_id_dimension << std::endl;
    return KernelStatus::KERNEL_STATUS_FAILED;
  }

  cnrtDim3_t dim{4, 1, 1};
  kernels::MLUBlockSliceIndividualPosId<<<dim, cnrtFuncTypeUnion1, queue>>>(
      context_pos_id, generate_pos_id, batch, context_seq_len, pos_id_dimension);
  return KernelStatus::KERNEL_STATUS_SUCCESS;
}

KernelStatus invokeIncrement2DPosId(cnrtQueue_t queue, int *generate_pos_id, int batch) {
  cnrtDim3_t dim{4, 1, 1};
  kernels::MLUBlockIncrement2DPosId<<<dim, cnrtFuncTypeUnion1, queue>>>(generate_pos_id, batch);
  return KernelStatus::KERNEL_STATUS_SUCCESS;
}
}  // namespace tmo
