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

namespace fd_xpu3 {

static __device__ inline int loada_float(_shared_ptr_ const int *ptr) {
  int ret;
  __asm__ __volatile__("loada.w %0,%1" : "=&r"(ret) : "r"(ptr));
  return ret;
}

static __device__ inline bool storea_float(_shared_ptr_ int *ptr, int value) {
  bool ret;
  __asm__ __volatile__("storea.w %0,%1,%2" : "=&r"(ret) : "r"(value), "r"(ptr));
  return ret;
}

static __device__ int atomic_add(_shared_ptr_ int *ptr, int value) {
  bool fail = true;
  int old_value;
  while (fail) {
    old_value = loada_float(ptr);
    int new_value = old_value + value;
    fail = storea_float(ptr, new_value);
  }
  return old_value;
}

__global__ void mtp_free_and_dispatch_block(bool *base_model_stop_flags,
                                            bool *stop_flags,
                                            bool *batch_drop,
                                            int *seq_lens_this_time,
                                            int *seq_lens_decoder,
                                            int *block_tables,
                                            int *encoder_block_lens,
                                            int *used_list_len,
                                            int *free_list,
                                            int *free_list_len,
                                            const int bsz,
                                            const int block_size,
                                            const int block_num_per_seq,
                                            const int max_draft_tokens) {
  int cid = core_id();
  int ncores = core_num();
  int clusterid = cluster_id();
  if (clusterid != 0 || cid >= bsz) return;

  // assert bsz <= 640
  const int max_bs = 640;
  int value_zero = 0;
  bool flag_true = true;

  __shared__ int free_list_len_sm;
  // 每次最多处理free_list数量为block_table_now_len
  const int block_table_now_len = 128;
  int block_table_now[block_table_now_len];
  for (int i = 0; i < block_table_now_len; i++) {
    block_table_now[i] = -1;
  }
  __shared__ bool base_model_stop_flags_sm[max_bs];
  __shared__ bool batch_drop_sm[max_bs];
  __shared__ int encoder_block_lens_sm[max_bs];
  __shared__ int seq_lens_decoder_sm[max_bs];

  int free_list_now[block_table_now_len];
  __shared__ int need_block_len_sm;
  __shared__ int need_block_list_sm[max_bs];
  __shared__ int used_list_len_sm[max_bs];
  __shared__ bool step_max_block_flag;

  if (cid == 0) {
    // len = 1
    need_block_len_sm = 0;
    GM2SM_ASYNC(free_list_len, &free_list_len_sm, sizeof(int));
    // len = bsz
    GM2SM_ASYNC(
        base_model_stop_flags, &base_model_stop_flags_sm, bsz * sizeof(bool));
    GM2SM_ASYNC(batch_drop, &batch_drop_sm, bsz * sizeof(bool));
    GM2SM_ASYNC(encoder_block_lens, &encoder_block_lens_sm, bsz * sizeof(int));
    GM2SM_ASYNC(used_list_len, used_list_len_sm, bsz * sizeof(int));
    GM2SM_ASYNC(seq_lens_decoder, seq_lens_decoder_sm, bsz * sizeof(int));
  }
  for (int tid = cid; tid < bsz; tid += ncores) {
    need_block_list_sm[tid] = 0;
  }
  mfence();
  sync_all();

  for (int tid = cid; tid < bsz; tid += ncores) {
    int64_t first_token_id_lm = -1;
    if (base_model_stop_flags_sm[tid] || batch_drop_sm[tid]) {
      // 回收block块
      const int encoder_block_len_lm = encoder_block_lens_sm[tid];
      const int decoder_used_len_lm = used_list_len_sm[tid];
      if (decoder_used_len_lm > 0) {
        const int ori_free_list_len =
            atomic_add(&free_list_len_sm, decoder_used_len_lm);
        for (int i = 0; i < decoder_used_len_lm; i += block_table_now_len) {
          int process_len = min(block_table_now_len, decoder_used_len_lm - i);
          GM2LM(
              block_tables + tid * block_num_per_seq + encoder_block_len_lm + i,
              free_list_now,
              process_len * sizeof(int));
          LM2GM(free_list_now,
                free_list + ori_free_list_len + i,
                process_len * sizeof(int));
          LM2GM(
              block_table_now,
              block_tables + tid * block_num_per_seq + encoder_block_len_lm + i,
              process_len * sizeof(int));
        }
        encoder_block_lens_sm[tid] = 0;
        used_list_len_sm[tid] = 0;
      }
      mfence();
    }
    int max_possible_block_idx =
        (seq_lens_decoder_sm[tid] + max_draft_tokens + 1) / block_size;
    int next_block_id;
    GM2LM(block_tables + tid * block_num_per_seq + max_possible_block_idx,
          &next_block_id,
          sizeof(int));

    if (!base_model_stop_flags[tid] && !batch_drop[tid] &&
        max_possible_block_idx < block_num_per_seq && next_block_id == -1) {
      // 统计需要分配block的位置和总数
      const int ori_need_block_len = atomic_add(&need_block_len_sm, 1);
      need_block_list_sm[ori_need_block_len] = tid;
      mfence();
    }

  }  // for
  sync_cluster();

  if (cid == 0) {
    while (need_block_len_sm > free_list_len_sm) {
      // 调度block，根据used_list_len从大到小回收block，直到满足need_block_len
      int max_used_list_len_id = 0;
      int max_used_list_len = 0;
      for (int i = 0; i < bsz; i++) {
        if ((!base_model_stop_flags_sm[i]) &&
            (used_list_len_sm[i] > max_used_list_len)) {
          max_used_list_len_id = i;
          max_used_list_len = used_list_len_sm[i];
        }
      }
      const int encoder_block_len_lm =
          encoder_block_lens_sm[max_used_list_len_id];
      for (int i = 0; i < max_used_list_len; i += block_table_now_len) {
        int process_len = min(block_table_now_len, max_used_list_len - i);
        GM2LM(block_tables + max_used_list_len_id * block_num_per_seq +
                  encoder_block_len_lm + i,
              free_list_now,
              process_len * sizeof(int));
        LM2GM(free_list_now,
              free_list + free_list_len_sm + i,
              process_len * sizeof(int));
        LM2GM(block_table_now,
              block_tables + max_used_list_len_id * block_num_per_seq +
                  encoder_block_len_lm + i,
              process_len * sizeof(int));
      }
      free_list_len_sm += max_used_list_len;
      LM2GM_ASYNC(&flag_true, stop_flags + max_used_list_len_id, sizeof(bool));
      LM2GM_ASYNC(
          &value_zero, seq_lens_this_time + max_used_list_len_id, sizeof(int));

      // 后面还要用，所以先放到sm中，用完在写回GM
      batch_drop_sm[max_used_list_len_id] = true;
      seq_lens_decoder_sm[max_used_list_len_id] = 0;
      used_list_len_sm[max_used_list_len_id] = 0;
      mfence();
    }
  }
  sync_cluster();

  int need_block_len_all = need_block_len_sm;
  for (int tid = cid; tid < need_block_len_all; tid += ncores) {
    // 为需要block的位置分配block，每个位置分配一个block
    const int need_block_id = need_block_list_sm[tid];
    if (!batch_drop_sm[need_block_id]) {
      used_list_len_sm[need_block_id]++;
      const int ori_free_list_len = atomic_add(&free_list_len_sm, -1);
      int free_block_id;
      GM2LM(free_list + ori_free_list_len - 1, &free_block_id, sizeof(int));
      LM2GM(&free_block_id,
            block_tables + need_block_id * block_num_per_seq +
                (seq_lens_decoder_sm[need_block_id] + max_draft_tokens + 1) /
                    block_size,
            sizeof(int));
    }
  }
  sync_cluster();

  if (cid == 0) {
    mfence();
    SM2GM_ASYNC(&free_list_len_sm, free_list_len, sizeof(int));
    SM2GM_ASYNC(used_list_len_sm, used_list_len, sizeof(int) * bsz);
    SM2GM_ASYNC(seq_lens_decoder_sm, seq_lens_decoder, sizeof(int) * bsz);
    SM2GM_ASYNC(batch_drop_sm, batch_drop, sizeof(bool) * bsz);
    SM2GM_ASYNC(encoder_block_lens_sm, encoder_block_lens, sizeof(int) * bsz);
    mfence();
  }
}

}  // namespace fd_xpu3
