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

namespace fd_xpu3 {

template <typename T>
__global__ void min_length_logits_process(T* logits,
                                          const int64_t* cur_len,
                                          const int64_t* min_len,
                                          const int64_t* eos_token_id,
                                          const int64_t bs,
                                          const int64_t length,
                                          const int64_t length_id,
                                          const int64_t end_length) {
  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;
  __simd__ float float32logits_now[32];

  for (int64_t i = tid; i < bs * end_length; i += nthreads) {
    bi = i / end_length;
    end_num = i % end_length;
    mfence_lm();
    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();
    if (cur_len_now >= 0 && cur_len_now < min_len_now) {
      GM2LM(
          eos_token_id + end_num, (void*)&(eos_token_id_now), sizeof(int64_t));
      GM2LM(logits + bi * length + eos_token_id_now,
            (void*)float32logits_now,
            sizeof(T));
      primitive_cast<T, float>(
          (const T*)(float32logits_now), float32logits_now, 1);
      float32logits_now[0] = -1e4;
      mfence_lm();
      primitive_cast<float, T>(float32logits_now, (T*)float32logits_now, 1);
      LM2GM((void*)float32logits_now,
            logits + bi * length + eos_token_id_now,
            sizeof(T));
    }
  }
}

#define _XPU_DEF__UPDATE_LOGITS_REPEAT_TIMES_(DATA_TYPE)         \
  template __global__ void 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 int64_t bs,                                          \
      const int64_t length,                                      \
      const int64_t length_id,                                   \
      const int64_t end_length);
_XPU_DEF__UPDATE_LOGITS_REPEAT_TIMES_(float);
_XPU_DEF__UPDATE_LOGITS_REPEAT_TIMES_(float16);

}  // namespace fd_xpu3
