#include "xpu/kernel/cluster.h"
#include "xpu/kernel/cluster_partition.h"
#include "xpu/kernel/cluster_primitive.h"

namespace xpu3 {
namespace plugin {

template <typename T>
__global__ void RebuildAppendPaddingKernel(const T *full_hidden_states,
                                           const int *cum_offset,
                                           const int *seq_len_encoder,
                                           const int *seq_len_decoder,
                                           const int *output_padding_offset,
                                           int max_seq_len,
                                           int dim_embed,
                                           int elem_nums,
                                           T *out) {
  int ncores = core_num();
  int cid = core_id();
  int tid = cid * cluster_num() + cluster_id();
  int nthreads = cluster_num() * ncores;
  int64_t mstart = -1;
  int64_t mend = -1;
  int64_t nstart = -1;
  int64_t nend = -1;
  partition2d<int64_t>(tid,
                       nthreads,
                       elem_nums / dim_embed,
                       dim_embed,
                       &mstart,
                       &mend,
                       &nstart,
                       &nend);

  const int64_t BUFFER_LEN = rounddown(6144 / sizeof(T), 64);
  __simd__ T lm_full_hidden_states[BUFFER_LEN];
  int output_padding_offset_val, cum_offset_val, seq_len_encoder_val,
      seq_len_decoder_val;

  for (int64_t _m = mstart; _m < mend; _m++) {
    int out_token_id = _m;
    GM2LM(output_padding_offset + out_token_id,
          &output_padding_offset_val,
          sizeof(int));
    int ori_token_id = out_token_id + output_padding_offset_val;
    int bi = ori_token_id / max_seq_len;
    GM2LM_ASYNC(seq_len_encoder + bi, &seq_len_encoder_val, sizeof(int));
    GM2LM(seq_len_decoder + bi, &seq_len_decoder_val, sizeof(int));
    int seq_id = 0;
    if (seq_len_encoder_val == 0 and seq_len_decoder_val == 0) {
      continue;
    } else if (seq_len_encoder_val != 0) {
      seq_id = seq_len_encoder_val - 1;
    }
    GM2LM(cum_offset + bi, &cum_offset_val, sizeof(int));
    int input_token_id = ori_token_id - cum_offset_val + seq_id;
    for (int64_t _n = nstart; _n < nend; _n += BUFFER_LEN) {
      int64_t read_size = min(BUFFER_LEN, nend - _n);
      // out[i] = full_hidden_states[(i / dim_embed +
      // output_padding_offset[i / dim_embed] - cum_offset[(i / dim_embed
      // + output_padding_offset[i / dim_embed]) / max_seq_len] + seq_id)
      // * dim_embed + i % dim_embed]
      GM2LM(full_hidden_states + input_token_id * dim_embed + _n,
            lm_full_hidden_states,
            read_size * sizeof(T));
      LM2GM(lm_full_hidden_states,
            out + _m * dim_embed + _n,
            read_size * sizeof(T));
    }
  }
}

#define _XPU_DEF_REBUILD_APPEND_PADDING_KERNEL(T)         \
  template __global__ void RebuildAppendPaddingKernel<T>( \
      const T *full_hidden_states,                        \
      const int *cum_offset,                              \
      const int *seq_len_encoder,                         \
      const int *seq_len_decoder,                         \
      const int *output_padding_offset,                   \
      int max_seq_len,                                    \
      int dim_embed,                                      \
      int elem_nums,                                      \
      T *out);

_XPU_DEF_REBUILD_APPEND_PADDING_KERNEL(bfloat16);
_XPU_DEF_REBUILD_APPEND_PADDING_KERNEL(float16);
_XPU_DEF_REBUILD_APPEND_PADDING_KERNEL(float);

}  // namespace plugin
}  // namespace xpu3
