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

namespace xpu3 {
namespace plugin {

static __device__ void atomic_add(_shared_ptr_ int *ptr, int v) {
  bool fail = true;
  while (fail) {
    int a;
    __asm__ __volatile__("loada.w %0,%1" : "=&r"(a) : "r"(ptr));
    a += v;
    __asm__ __volatile__("storea.w %0,%1,%2" : "=&r"(fail) : "r"(a), "r"(ptr));
  }
}

// original version
__device__ void speculate_update_repeat_times_normal(
    char *lm,
    __shared_ptr__ char *sm,
    __global_ptr__ const int64_t *pre_ids,
    __global_ptr__ const int64_t *cur_len,
    __global_ptr__ int *repeat_times,
    __global_ptr__ const int *output_padding_offset,
    const int64_t bs,
    const int64_t length,
    const int64_t length_id,
    const int64_t token_num,
    const int64_t max_seq_len) {
  int cid = core_id();
  int ncores = core_num();
  int clusterid = cluster_id();
  int nclusters = cluster_num();
  int tid = clusterid * ncores + cid;

  const int max_sm_len = 256 * 1024 / sizeof(int);
  __shared_ptr__ int *repeated_times_sm = (__shared_ptr__ int *)sm;
  int64_t pre_id_lm;
  int n_length = (length + max_sm_len - 1) / max_sm_len;

  int64_t *cur_len_lm = (int64_t *)lm;
  int output_padding_offset_now;
  GM2LM(cur_len, cur_len_lm, bs * sizeof(int64_t));

  for (int nli = 0; nli < n_length; nli++) {
    int step = nli * max_sm_len;
    int cur_length = min(max_sm_len, length - step);
    for (int64_t i = clusterid; i < token_num; i += nclusters) {
      GM2LM(output_padding_offset + i, &output_padding_offset_now, sizeof(int));
      int64_t bi = (i + output_padding_offset_now) / max_seq_len;
      if (bi >= bs || cur_len_lm[bi] < 0) {
        continue;
      }
      if (cid == 0) {
        GM2SM_ASYNC(repeat_times + i * length + step,
                    repeated_times_sm,
                    sizeof(int) * cur_length);
      }
      mfence();
      sync_cluster();
      for (int j = cid; j < length_id; j += ncores) {
        GM2LM(pre_ids + bi * length_id + j, &pre_id_lm, sizeof(int64_t));
        if (pre_id_lm < 0) {
          break;
        }
        if (pre_id_lm >= step && pre_id_lm < step + cur_length) {
          atomic_add(repeated_times_sm + pre_id_lm - step, 1);
        }
        mfence();
      }
      sync_cluster();
      if (cid == 0) {
        SM2GM_ASYNC(repeated_times_sm,
                    repeat_times + i * length + step,
                    sizeof(int) * cur_length);
      }
      mfence();
      sync_cluster();
    }
  }
}

// best optimized version
// about 49000+ ns
__device__ void speculate_update_repeat_times_optimized(
    char *lm,
    __shared_ptr__ char *sm,
    __global_ptr__ const int64_t *pre_ids,            // {bs, length_id}
    __global_ptr__ const int64_t *cur_len,            // {bs}
    __global_ptr__ int *repeat_times,                 // {token_num, length}
    __global_ptr__ const int *output_padding_offset,  // {token_num}
    const int64_t bs,
    const int64_t length,
    const int64_t length_id,
    const int64_t token_num,
    const int64_t max_seq_len) {
  int cid = core_id();
  int ncores = core_num();
  int clusterid = cluster_id();
  int nclusters = cluster_num();
  int tid = clusterid * ncores + cid;

  const int repeat_times_sm_len = 250 * 1024 / sizeof(int);
  __shared_ptr__ int *repeat_times_sm = (__shared_ptr__ int *)sm;

  // assert bs <= 640
  int cur_len_sm_len = 640;
  __shared_ptr__ int64_t *cur_len_sm =
      (__shared_ptr__ int64_t *)(repeat_times_sm + repeat_times_sm_len);
  __shared_ptr__ int *output_padding_offset_sm =
      (__shared_ptr__ int *)(cur_len_sm + cur_len_sm_len);
  DoublePtr<1, SmPtr<int>> buffer_ptr_output_padding_offset(
      (SmPtr<int>(output_padding_offset_sm)));
  int pre_ids_lm_len = 4;
  int64_t *pre_ids_lm = (int64_t *)lm;
  DoublePtr<4, LmPtr<int64_t>> buffer_ptr_pre_ids((LmPtr<int64_t>(pre_ids_lm)));

  int64_t i = clusterid;
  if (i < token_num && cid == 0) {
    GM2SM_ASYNC(cur_len, cur_len_sm, bs * sizeof(int64_t));
    buffer_ptr_output_padding_offset.gm_load_async(output_padding_offset + i,
                                                   1);
    mfence_sm();
  }
  sync_all();
  for (; i < token_num; i += nclusters) {
    if (cid == 0 && i + nclusters < token_num) {
      buffer_ptr_output_padding_offset.next().gm_load_async(
          output_padding_offset + i + nclusters, 1);
    }
    int64_t bi = (i + (buffer_ptr_output_padding_offset.ptr[0])) / max_seq_len;
    buffer_ptr_output_padding_offset.toggle();
    if (bi >= bs || cur_len_sm[bi] < 0) {
      mfence_sm();
      sync_all();
      continue;
    }
    int64_t boundary = -1;
    for (int64_t repeat_times_start = 0; repeat_times_start < length;
         repeat_times_start += repeat_times_sm_len) {
      int64_t repeat_times_read_size =
          min(length - repeat_times_start, repeat_times_sm_len);
      int64_t start, end;
      partition(cid, ncores, repeat_times_read_size, 1, &start, &end);
      int64_t load_start = repeat_times_start + start;
      int64_t repeat_times_read_size_per_core = end - start;
      if (repeat_times_read_size_per_core > 0) {
        GM2SM(repeat_times + i * length + load_start,
              repeat_times_sm + start,
              repeat_times_read_size_per_core * sizeof(int));
      }
      sync_all();
      // each core loads pre_ids step by step and record the index of pre_ids
      // which is less than zero, and store the index to boundary
      if (repeat_times_start == 0) {
        bool do_prone = false;
        int64_t j = cid * pre_ids_lm_len;
        int64_t pre_ids_read_size =
            min(static_cast<int64_t>(pre_ids_lm_len), length_id - j);
        buffer_ptr_pre_ids.gm_load(pre_ids + bi * length_id + j,
                                   pre_ids_read_size);
        for (; j < length_id && !do_prone; j += ncores * pre_ids_lm_len) {
          int64_t pre_ids_read_size_next =
              min(static_cast<int64_t>(pre_ids_lm_len),
                  length_id - (j + ncores * pre_ids_lm_len));
          if (buffer_ptr_pre_ids.ptr[pre_ids_read_size - 1] >= 0 &&
              pre_ids_read_size_next > 0) {
            buffer_ptr_pre_ids.next().gm_load_async(
                pre_ids + bi * length_id + j + ncores * pre_ids_lm_len,
                pre_ids_read_size_next);
          }
          for (int k = 0; k < pre_ids_read_size; k++) {
            if (buffer_ptr_pre_ids.ptr[k] < 0) {
              do_prone = true;
              boundary = j + k;
              break;
            }
            if (buffer_ptr_pre_ids.ptr[k] >= repeat_times_start &&
                buffer_ptr_pre_ids.ptr[k] <
                    repeat_times_start + repeat_times_read_size) {
              atomic_add(repeat_times_sm + buffer_ptr_pre_ids.ptr[k] -
                             repeat_times_start,
                         1);
            }
          }
          mfence_lm();
          pre_ids_read_size = pre_ids_read_size_next;
          buffer_ptr_pre_ids.toggle();
        }
      }
      // each core loads all the needed pre_ids into lm without mfence inbetween
      // according to the index recorded by previous iteration
      else {
        int cnt = -1;
        int64_t pre_ids_read_size = 0;
        for (int64_t j = cid * pre_ids_lm_len; j < boundary;
             j += ncores * pre_ids_lm_len) {
          cnt++;
          pre_ids_read_size =
              min(static_cast<int64_t>(pre_ids_lm_len), boundary - j);
          GM2LM_ASYNC(pre_ids + bi * length_id + j,
                      pre_ids_lm + cnt * pre_ids_lm_len,
                      pre_ids_read_size * sizeof(int64_t));
        }
        mfence_lm();
        cnt = max(0, cnt);
        for (int k = 0; k < cnt * pre_ids_lm_len + pre_ids_read_size; k++) {
          if (pre_ids_lm[k] >= repeat_times_start &&
              pre_ids_lm[k] < repeat_times_start + repeat_times_read_size) {
            atomic_add(repeat_times_sm + pre_ids_lm[k] - repeat_times_start, 1);
          }
        }
      }
      mfence_sm();
      sync_cluster();
      if (repeat_times_read_size_per_core > 0) {
        SM2GM(repeat_times_sm + start,
              repeat_times + i * length + load_start,
              repeat_times_read_size_per_core * sizeof(int));
      }
      sync_all();
    }
  }
}

__global__ void speculate_update_repeat_times(const int64_t *pre_ids,
                                              const int64_t *cur_len,
                                              int *repeat_times,
                                              const int *output_padding_offset,
                                              const int64_t bs,
                                              const int64_t length,
                                              const int64_t length_id,
                                              const int64_t token_num,
                                              const int64_t max_seq_len) {
  char lm[6 * 1024];
  __shared__ char sm[256 * 1024];

  if (length_id <= 6 * 1024 * 64 / sizeof(int64_t)) {
    speculate_update_repeat_times_optimized(lm,
                                            sm,
                                            pre_ids,
                                            cur_len,
                                            repeat_times,
                                            output_padding_offset,
                                            bs,
                                            length,
                                            length_id,
                                            token_num,
                                            max_seq_len);
  } else {
    speculate_update_repeat_times_normal(lm,
                                         sm,
                                         pre_ids,
                                         cur_len,
                                         repeat_times,
                                         output_padding_offset,
                                         bs,
                                         length,
                                         length_id,
                                         token_num,
                                         max_seq_len);
  }
}

}  // namespace plugin
}  // namespace xpu3
