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

namespace fd_xpu3 {

template <typename T>
__global__ void speculate_min_length_logits_process(
    T* logits,
    const int64_t* cur_len,
    const int64_t* min_len,
    const int64_t* eos_token_id,
    const int* batch_id_per_token_output,
    const int* cu_seqlens_q_output,
    const int64_t bs,
    const int64_t length,
    const int64_t length_id,
    const int64_t end_length,
    const int64_t token_num,
    const int64_t max_seq_len) {
  int ncores = core_num();
  int cid = core_id();
  int tid = cluster_num() * cid + cluster_id();
  int nthreads = cluster_num() * ncores;

  int64_t cur_len_now;
  int64_t min_len_now;
  int64_t eos_token_id_now;
  int64_t bi;
  int64_t end_num;
  int batch_id_per_token_output_now;
  int cu_seqlens_q_output_now;
  __simd__ float float32logits_now[32];

  for (int64_t i = tid; i < token_num * end_length; i += nthreads) {
    int64_t token_idx = i / end_length;
    GM2LM(batch_id_per_token_output + token_idx,
          &batch_id_per_token_output_now,
          sizeof(int));
    bi = batch_id_per_token_output[token_idx];
    if (bi >= bs) {
      continue;
    }
    end_num = i % end_length;
    GM2LM_ASYNC(
        cu_seqlens_q_output + bi, (void*)&cu_seqlens_q_output_now, sizeof(int));
    GM2LM_ASYNC(cur_len + bi, (void*)&(cur_len_now), sizeof(int64_t));
    GM2LM_ASYNC(min_len + bi, (void*)&(min_len_now), sizeof(int64_t));
    mfence();
    int query_start_token_idx = cu_seqlens_q_output_now;
    if (cur_len_now >= 0 &&
        (cur_len_now + (token_idx - query_start_token_idx) < min_len_now)) {
      GM2LM(
          eos_token_id + end_num, (void*)&(eos_token_id_now), sizeof(int64_t));
      GM2LM(logits + token_idx * length + eos_token_id_now,
            (void*)float32logits_now,
            sizeof(T));
      primitive_cast<T, float>(
          (const T*)(float32logits_now), float32logits_now, 1);
      float32logits_now[0] = std::is_same<T, float16>::value ? -1e4 : -1e10;
      mfence_lm();
      primitive_cast<float, T>(float32logits_now, (T*)float32logits_now, 1);
      LM2GM((void*)float32logits_now,
            logits + token_idx * length + eos_token_id_now,
            sizeof(T));
    }
  }
}

#define _XPU_DEF__UPDATE_LOGITS_REPEAT_TIMES_(DATA_TYPE)                   \
  template __global__ void speculate_min_length_logits_process<DATA_TYPE>( \
      DATA_TYPE * logits,                                                  \
      const int64_t* cur_len,                                              \
      const int64_t* min_len,                                              \
      const int64_t* eos_token_id,                                         \
      const int* batch_id_per_token_output,                                \
      const int* cu_seqlens_q_output,                                      \
      const int64_t bs,                                                    \
      const int64_t length,                                                \
      const int64_t length_id,                                             \
      const int64_t end_length,                                            \
      const int64_t token_num,                                             \
      const int64_t max_seq_len);
_XPU_DEF__UPDATE_LOGITS_REPEAT_TIMES_(float);
_XPU_DEF__UPDATE_LOGITS_REPEAT_TIMES_(float16);
_XPU_DEF__UPDATE_LOGITS_REPEAT_TIMES_(bfloat16);

}  // namespace fd_xpu3
