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

namespace fd_xpu3 {

__device__ void prefix_sum(__shared_ptr__ int* sm_seq_lens_encoder,
                           __shared_ptr__ int* sm_seq_lens_this_time,
                           __shared_ptr__ int* sm_batch_token_num,
                           __shared_ptr__ int* sm_cu_batch_token_offset,
                           __shared_ptr__ int* sm_cu_next_token_offset,
                           __global_ptr__ int* batch_token_num,
                           __global_ptr__ int* cu_batch_token_offset,
                           __global_ptr__ const int* seq_lens_this_time,
                           __global_ptr__ const int* seq_lens_encoder,
                           const int real_bsz) {
  int cid = core_id();
  int clus_id = cluster_id();

  if (clus_id < real_bsz && cid == 0) {
    GM2SM_ASYNC(seq_lens_encoder, sm_seq_lens_encoder, real_bsz * sizeof(int));
    GM2SM(seq_lens_this_time, sm_seq_lens_this_time, real_bsz * sizeof(int));
    int next_token_num_previous = 0;
    for (int bid = 0; bid < real_bsz; bid++) {
      sm_batch_token_num[bid] =
          sm_seq_lens_encoder[bid] > 0 ? 2 : sm_seq_lens_this_time[bid];
      if (bid == 0) {
        sm_cu_batch_token_offset[bid] = 0;
        sm_cu_next_token_offset[bid] = 0;
      } else {
        sm_cu_batch_token_offset[bid] =
            sm_cu_batch_token_offset[bid - 1] + sm_batch_token_num[bid - 1];
        sm_cu_next_token_offset[bid] =
            sm_cu_next_token_offset[bid - 1] + next_token_num_previous;
      }
      next_token_num_previous =
          sm_seq_lens_encoder[bid] > 0 ? 1 : sm_seq_lens_this_time[bid];
    }
    mfence_sm();
    if (clus_id == 0) {
      SM2GM_ASYNC(sm_batch_token_num, batch_token_num, real_bsz * sizeof(int));
      SM2GM_ASYNC(sm_cu_batch_token_offset,
                  cu_batch_token_offset,
                  real_bsz * sizeof(int));
    }
  }
  mfence_sm();
  sync_all();
}
__global__ void speculate_get_logits(float* draft_logits,
                                     int* next_token_num,
                                     int* batch_token_num,
                                     int* cu_next_token_offset,
                                     int* cu_batch_token_offset,
                                     const float* logits,
                                     const float* first_token_logits,
                                     const int* seq_lens_this_time,
                                     const int* seq_lens_encoder,
                                     const int real_bsz,
                                     const int vocab_size) {
  int cid = core_id();
  int ncores = core_num();
  int clus_id = cluster_id();
  int nclusters = cluster_num();

  int lm_size = 2 * 1024;
  int lm_buf_len = lm_size / sizeof(float);
  float first_token_logits_now_lm[lm_buf_len];
  float logits_now_lm[lm_buf_len];

  const int sm_size = 256 * 1024;
  __shared__ char sm[sm_size];
  int sm_max_buf_len = 256 * 1024 / sizeof(int);
  sm_max_buf_len /= 5;
  __shared_ptr__ int* sm_seq_lens_encoder = (__shared_ptr__ int*)sm;
  __shared_ptr__ int* sm_seq_lens_this_time =
      sm_seq_lens_encoder + sm_max_buf_len;
  __shared_ptr__ int* sm_batch_token_num =
      sm_seq_lens_this_time + sm_max_buf_len;
  __shared_ptr__ int* sm_cu_batch_token_offset =
      sm_batch_token_num + sm_max_buf_len;
  __shared_ptr__ int* sm_cu_next_token_offset =
      sm_cu_batch_token_offset + sm_max_buf_len;

  prefix_sum(sm_seq_lens_encoder,
             sm_seq_lens_this_time,
             sm_batch_token_num,
             sm_cu_batch_token_offset,
             sm_cu_next_token_offset,
             batch_token_num,
             cu_batch_token_offset,
             seq_lens_this_time,
             seq_lens_encoder,
             real_bsz);

  for (int bid = clus_id; bid < real_bsz; bid += nclusters) {
    auto* draft_logits_now =
        draft_logits + sm_cu_batch_token_offset[bid] * vocab_size;
    auto* logits_now = logits + sm_cu_next_token_offset[bid] * vocab_size;
    auto* first_token_logits_now = first_token_logits + bid * vocab_size;

    for (int i = cid * lm_buf_len; i < vocab_size; i += ncores * lm_buf_len) {
      int read_len = min(lm_buf_len, vocab_size - i);
      if (sm_seq_lens_encoder[bid] > 0) {
        GM2LM_ASYNC(first_token_logits_now + i,
                    first_token_logits_now_lm,
                    read_len * sizeof(float));
        GM2LM(logits_now + i, logits_now_lm, read_len * sizeof(float));
        LM2GM_ASYNC(first_token_logits_now_lm,
                    draft_logits_now + i,
                    read_len * sizeof(float));
        LM2GM(logits_now_lm,
              draft_logits_now + vocab_size + i,
              read_len * sizeof(float));
      } else {
        for (int j = 0; j < sm_seq_lens_this_time[bid]; j++) {
          GM2LM(logits_now + j * vocab_size + i,
                logits_now_lm,
                read_len * sizeof(float));
          LM2GM(logits_now_lm,
                draft_logits_now + j * vocab_size + i,
                read_len * sizeof(float));
        }
      }
    }
  }
}

}  // namespace fd_xpu3
