#include "xpu/kernel/cluster.h"
#include "xpu/kernel/cluster_debug.h"
#include "xpu/kernel/cluster_primitive.h"
namespace xpu3 {
namespace plugin {
#define MAX_LM_SIZE 28672
// One core has 32KB LM（group LM), MAX_LM_SIZE = (32 - 4)KB / 2 = 30720, 4KB is
// the stack space
#define MAX_BATCH 512
#define ALIGNMENT 64

template <typename TX, typename TY>
static __device__ void do_memcpy_1d(_global_ptr_ TX* src,
                                    _global_ptr_ TY* dst,
                                    int64_t copy_size) {
#ifdef __XPU3__
  constexpr int buf_size = 2048;
#else
  constexpr int buf_size = 512;
#endif
  __group_shared__ __simd__ float double_lmx[2][buf_size];
  int64_t pingpong = 0;
  for (int64_t i = 0; i < copy_size; i += buf_size) {
    int real_size = min<int64_t>(buf_size, copy_size - i);
    _group_shared_ptr_ float* lmx = double_lmx[pingpong];
    GM2GSM(src + i, lmx, real_size * sizeof(TX));
    if (!xpu_std::is_same<TX, TY>::value) {
      primitive_cast_gsm<TX, float>(
          (_group_shared_ptr_ TX*)lmx, lmx, real_size);
      primitive_cast_gsm<float, TY>(
          lmx, (_group_shared_ptr_ TY*)lmx, real_size);
    }
    GSM2GM_ASYNC((_group_shared_ptr_ TY*)lmx, dst + i, real_size * sizeof(TY));
    pingpong = 1 - pingpong;
  }
  mfence();
}

template <typename TX, typename TY>
__global__ void eb_gather_next_token(TX* src,
                                     TY* dst,
                                     int* encoder_seqs_lods,
                                     int* encoder_batch_map,
                                     int* decoder_batch_map,
                                     int en_batch,
                                     int de_batch,
                                     int64_t copy_size) {
  int tid = core_id() * cluster_num() + cluster_id();
  int nthreads = core_num() * cluster_num();
  __group_shared__ int local_lods_en[MAX_BATCH + 1];
  __group_shared__ int local_map_en[MAX_BATCH];
  __group_shared__ int local_map_de[MAX_BATCH];
  GM2GSM_ASYNC(encoder_seqs_lods, local_lods_en, (en_batch + 1) * sizeof(int));
  if (en_batch > 0) {
    GM2GSM_ASYNC(encoder_batch_map, local_map_en, en_batch * sizeof(int));
  }
  if (de_batch > 0) {
    GM2GSM_ASYNC(decoder_batch_map, local_map_de, de_batch * sizeof(int));
  }
  mfence();
  int encode_len_total = en_batch > 0 ? local_lods_en[en_batch] : 0;
  int total_batch = en_batch + de_batch;
  int start = 0;
  int end = 0;
  partition(tid, nthreads, total_batch, 1, &start, &end);
  for (int i = start; i < end; i++) {
    if (i < en_batch) {
      // src encode part
      _global_ptr_ TX* cur_src = src + (local_lods_en[i + 1] - 1) * copy_size;
      _global_ptr_ TY* cur_dst = dst + local_map_en[i] * copy_size;
      do_memcpy_1d<TX, TY>(cur_src, cur_dst, copy_size);
    } else {
      // src decode part
      int de_idx = i - en_batch;
      _global_ptr_ TX* cur_src = src + (encode_len_total + de_idx) * copy_size;
      _global_ptr_ TY* cur_dst = dst + local_map_de[de_idx] * copy_size;
      do_memcpy_1d<TX, TY>(cur_src, cur_dst, copy_size);
    }
  }
}
#define _XPU_DEF__EB_GATHER_NEXT_TOKEN(TX, TY)           \
  template __global__ void eb_gather_next_token<TX, TY>( \
      TX * src,                                          \
      TY * dst,                                          \
      int* encoder_seqs_lods,                            \
      int* encoder_batch_map,                            \
      int* decoder_batch_map,                            \
      int en_batch,                                      \
      int de_batch,                                      \
      int64_t copy_size);

_XPU_DEF__EB_GATHER_NEXT_TOKEN(float16, float16);
_XPU_DEF__EB_GATHER_NEXT_TOKEN(bfloat16, bfloat16);
_XPU_DEF__EB_GATHER_NEXT_TOKEN(float, float);
_XPU_DEF__EB_GATHER_NEXT_TOKEN(float16, float);
_XPU_DEF__EB_GATHER_NEXT_TOKEN(float, float16);
_XPU_DEF__EB_GATHER_NEXT_TOKEN(bfloat16, float16);
_XPU_DEF__EB_GATHER_NEXT_TOKEN(float16, bfloat16);
_XPU_DEF__EB_GATHER_NEXT_TOKEN(bfloat16, float);
_XPU_DEF__EB_GATHER_NEXT_TOKEN(float, bfloat16);
}  // namespace plugin
}  // namespace xpu3
