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

namespace xpu3 {
namespace plugin {

__global__ void ComputeSelfOrderKernel(const int* last_seq_lens_this_time,
                                       const int* seq_lens_this_time,
                                       const int64_t* step_idx,
                                       int* src_map,
                                       int* output_token_num,
                                       int bsz) {
  int tid = core_id() * cluster_num() + cluster_id();
  if (tid != 0) {
    return;
  }

  char lm[6 * 1024];
  int buf_size = 256;
  int* lm_last_seq_lens_this_time = (int*)lm;
  int* lm_seq_lens_this_time = lm_last_seq_lens_this_time + buf_size;
  int64_t* lm_step_idx = (int64_t*)(lm_seq_lens_this_time + buf_size);
  int* lm_src_map = (int*)(lm_step_idx + buf_size);

  int in_offset = 0;
  int out_offset = 0;
  int previous_out_offset = out_offset;
  for (int i = 0; i < bsz; i += buf_size) {
    int64_t read_size = min(static_cast<int64_t>(bsz - i), buf_size);
    GM2LM_ASYNC(last_seq_lens_this_time + i,
                lm_last_seq_lens_this_time,
                read_size * sizeof(int));
    GM2LM_ASYNC(
        seq_lens_this_time + i, lm_seq_lens_this_time, read_size * sizeof(int));
    GM2LM(step_idx + i, lm_step_idx, read_size * sizeof(int64_t));
    for (int j = 0; j < read_size; j++) {
      int cur_seq_lens_this_time = lm_seq_lens_this_time[j];
      int cur_last_seq_lens_this_time = lm_last_seq_lens_this_time[j];
      int64_t cur_step_idx = lm_step_idx[j];
      // 1. encoder
      if (cur_step_idx == 1 && cur_seq_lens_this_time > 0) {
        in_offset += 1;
        lm_src_map[j] = in_offset - 1;
        out_offset++;
        // 2. decoder
      } else if (cur_seq_lens_this_time > 0) /* =1 */ {
        in_offset += cur_last_seq_lens_this_time;
        lm_src_map[j] = in_offset - 1;
        out_offset++;
        // 3. stop
      } else {
        // first token end
        if (cur_step_idx == 1) {
          in_offset += cur_last_seq_lens_this_time > 0 ? 1 : 0;
          // normal end
        } else {
          in_offset += cur_last_seq_lens_this_time;
        }
      }
    }
    mfence_lm();
    if (out_offset > previous_out_offset) {
      LM2GM_ASYNC(lm_src_map,
                  src_map + previous_out_offset,
                  (out_offset - previous_out_offset) * sizeof(int));
    }
    previous_out_offset = out_offset;
  }
  mfence_lm();
  LM2GM(&out_offset, output_token_num, sizeof(int));
}

}  // namespace plugin
}  // namespace xpu3
