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

#include "xpu/kernel/cluster_debug.h"

namespace fd_xpu3 {

#define MAX_BATCH_SIZE 1024

static inline __device__ int v_reduce_sum_int32(int32x16_t& v0) {
  auto v1 = vsrlp_int32x16(1 << 8, v0);
  v0 = vvadd_int32x16(v0, v1);
  v1 = vsrlp_int32x16(1 << 7, v0);
  v0 = vvadd_int32x16(v0, v1);
  v1 = vsrlp_int32x16(1 << 6, v0);
  v0 = vvadd_int32x16(v0, v1);
  v1 = vsrlp_int32x16(1 << 5, v0);
  v0 = vvadd_int32x16(v0, v1);
  return vextract_int32x16(v0, 1);
}

inline __device__ int primitive_reduce_sum_sm(__shared_ptr__ const int* x,
                                              int64_t len) {
  int32x16_t x_l, x_h;
  int32x16_t sum = vset_zero_int();
  const auto rounddown_len = rounddown32(len);

  for (int64_t i = 0; i < rounddown_len; i += 32) {
    vload2_sm(x + i, x_l, x_h);
    sum = vvadd_int32x16(sum, x_l);
    sum = vvadd_int32x16(sum, x_h);
  }

  if (rounddown_len < len) {
    const auto mask = ~(-1 << (len - rounddown_len));
    vload2_sm_mz(x + rounddown_len, x_l, x_h, mask);
    sum = vvadd_int32x16(sum, x_l);
    sum = vvadd_int32x16(sum, x_h);
  }
  return v_reduce_sum_int32(sum);
}

__global__ void speculate_preprocess_kernel(
    int64_t* ids_remove_padding,
    int* batch_id_per_token,
    int* cu_seqlens_q,
    int* cu_seqlens_k,
    int* seq_lens_output,
    int* cu_seq_lens_q_output,
    int* batch_id_per_token_output,
    int* real_output_token_num,
    const int64_t* input_data,
    const int* seq_lens,
    const int64_t* draft_tokens,
    const int* seq_lens_encoder,
    const int max_seq_len,
    const int max_draft_tokens_per_batch,
    const int real_bs) {
  int cid = core_id();
  int ncores = core_num();
  int clusterid = cluster_id();
  int nclusters = cluster_num();
  __shared__ int sm_seq_lens[MAX_BATCH_SIZE];
  __shared__ int sm_seq_lens_output[MAX_BATCH_SIZE];
  __shared__ int sm_seq_lens_encoder[MAX_BATCH_SIZE];
  __shared__ int sm_cum_seq_len, sm_cum_seq_len_output;
  __simd__ __shared__ int buffer_cu_seqlens[64];
  __simd__ __shared__ int buffer_cu_seqlens_output[64];

  if (cid == 0) {
    GM2SM_ASYNC(seq_lens, sm_seq_lens, sizeof(int) * real_bs);
    GM2SM(seq_lens_encoder, sm_seq_lens_encoder, sizeof(int) * real_bs);
  }
  sync_all();
  for (int bid = cid; bid < real_bs; bid += ncores) {
    if (sm_seq_lens[bid] == 0) {
      sm_seq_lens_output[bid] = 0;
    } else if (sm_seq_lens[bid] == 1) {
      sm_seq_lens_output[bid] = 1;
    } else if (sm_seq_lens_encoder[bid] != 0) {
      sm_seq_lens_output[bid] = 1;
    } else {
      sm_seq_lens_output[bid] = sm_seq_lens[bid];
    }
  }
  mfence_sm();
  sync_all();

  for (int bi = clusterid; bi < real_bs; bi += nclusters) {
    int cum_seq_len = 0;
    int cum_seq_len_output = 0;
    for (int i = cid; i < bi + 1; i += ncores) {
      cum_seq_len += sm_seq_lens[i];
      cum_seq_len_output += sm_seq_lens_output[i];
    }
    buffer_cu_seqlens[cid] = cum_seq_len;
    buffer_cu_seqlens_output[cid] = cum_seq_len_output;
    mfence();
    sync_all();
    if (cid == 0) {
      cum_seq_len =
          primitive_reduce_sum_sm(buffer_cu_seqlens, min(bi + 1, ncores));
      cum_seq_len_output = primitive_reduce_sum_sm(buffer_cu_seqlens_output,
                                                   min(bi + 1, ncores));
      LM2GM_ASYNC(&cum_seq_len, cu_seqlens_q + bi + 1, sizeof(int));
      LM2GM_ASYNC(&cum_seq_len, cu_seqlens_k + bi + 1, sizeof(int));
      LM2GM_ASYNC(
          &cum_seq_len_output, cu_seq_lens_q_output + bi + 1, sizeof(int));
      if (bi == real_bs - 1) {
        LM2GM_ASYNC(&cum_seq_len_output, real_output_token_num, sizeof(int));
      }
      sm_cum_seq_len = cum_seq_len;
      sm_cum_seq_len_output = cum_seq_len_output;
    }
    mfence();
    sync_all();

    const int lm_seq_lens = sm_seq_lens[bi];
    const int lm_seq_lens_encoder = sm_seq_lens_encoder[bi];
    for (int i = cid; i < lm_seq_lens; i += ncores) {
      const int tgt_seq_id = sm_cum_seq_len - lm_seq_lens + i;
      if (max_draft_tokens_per_batch > 0 && lm_seq_lens_encoder <= 0) {
        // speculative decoding
        const int src_seq_id = bi * max_draft_tokens_per_batch + i;
        int64_t lm_draft_tokens;
        GM2LM(draft_tokens + src_seq_id, &lm_draft_tokens, sizeof(int64_t));
        LM2GM(
            &lm_draft_tokens, ids_remove_padding + tgt_seq_id, sizeof(int64_t));
      } else {
        // Non-speculative decoding
        const int src_seq_id = bi * max_seq_len + i;
        int64_t lm_input_data;
        GM2LM(input_data + src_seq_id, &lm_input_data, sizeof(int64_t));
        LM2GM(&lm_input_data, ids_remove_padding + tgt_seq_id, sizeof(int64_t));
      }
      LM2GM(&bi, batch_id_per_token + tgt_seq_id, sizeof(int));
    }

    const int lm_seq_lens_output = sm_seq_lens_output[bi];
    for (int i = cid; i < lm_seq_lens_output; i += ncores) {
      const int tgt_seq_id_output =
          sm_cum_seq_len_output - lm_seq_lens_output + i;
      LM2GM(&bi, batch_id_per_token_output + tgt_seq_id_output, sizeof(int));
    }
    mfence();
    sync_all();
  }

  if (cid == 0 && clusterid == 0) {
    const int lm_zero = 0;
    LM2GM_ASYNC(&lm_zero, cu_seqlens_q, sizeof(int));
    LM2GM_ASYNC(&lm_zero, cu_seqlens_k, sizeof(int));
    LM2GM(&lm_zero, cu_seq_lens_q_output, sizeof(int));
  }
}
}  // namespace fd_xpu3
