#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 rebuildHiddenStatesKernel(const T* input,
                                          const int* position_map,
                                          T* output,
                                          int dim_embed,
                                          int elem_cnt) {
  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_cnt / dim_embed,
                       dim_embed,
                       &mstart,
                       &mend,
                       &nstart,
                       &nend);

  const int64_t BUFFER_LEN = 6144 / sizeof(T);
  T lm_input[BUFFER_LEN];

  for (int64_t _m = mstart; _m < mend; _m++) {
    int ori_token_idx = _m;
    int token_idx;
    GM2LM(position_map + _m, &token_idx, sizeof(int));
    if (token_idx >= 0) {
      for (int64_t _n = nstart; _n < nend; _n += BUFFER_LEN) {
        int64_t read_size = min(BUFFER_LEN, nend - _n);
        GM2LM(input + ori_token_idx * dim_embed + _n,
              lm_input,
              read_size * sizeof(T));
        LM2GM(lm_input,
              output + token_idx * dim_embed + _n,
              read_size * sizeof(T));
      }
    }
  }
}

#define _XPU_DEF_REBUILD_HIDDEN_STATES_KERNEL(T)         \
  template __global__ void rebuildHiddenStatesKernel<T>( \
      const T* input,                                    \
      const int* position_map,                           \
      T* output,                                         \
      int dim_embed,                                     \
      int elem_cnt);

_XPU_DEF_REBUILD_HIDDEN_STATES_KERNEL(bfloat16);
_XPU_DEF_REBUILD_HIDDEN_STATES_KERNEL(float);
_XPU_DEF_REBUILD_HIDDEN_STATES_KERNEL(float16);

}  // namespace plugin
}  // namespace xpu3
