// clang-format off
#include <mlu.h>
// clang-format on
#include "kernel_utils.h"
#include "update_out_and_lse.mluh"

namespace tmo {
namespace kernels {
#define NRAM_REMAIN_SIZE (32 * 1024)
#define NRAM_BUFFER_SIZE (__MLU_NRAM_SIZE__ * 1024 - NRAM_REMAIN_SIZE)
#define INF (2139095040)
__nram__ char nram_buffer[NRAM_BUFFER_SIZE];

__mlu_func__ void splitTask(int32_t total_task, int32_t &task_length, int32_t &task_offset) {
  task_length = (total_task + taskDimX - 1) / taskDimX;
  task_offset = taskIdX * task_length;
  task_length = std::min(total_task - task_offset, task_length);
}

template <typename T>
__mlu_func__ void toFloat(float *dst, T *src, int32_t num) {
  if (std::is_same<T, half>::value) {
    __bang_half2float(dst, (half *)src, num);
  } else if (std::is_same<T, bfloat16_t>::value) {
    __bang_bfloat162float(dst, (bfloat16_t *)src, num);
  } else if (std::is_same<T, float>::value) {
    __bang_move(dst, src, num * sizeof(float));
  }
}

template <typename T>
__mlu_func__ void floatTo(T *dst, float *src, int32_t num) {
  if (std::is_same<T, half>::value) {
    __bang_float2half_rn((half *)dst, src, num);
  } else if (std::is_same<T, bfloat16_t>::value) {
    __bang_float2bfloat16_rn((bfloat16_t *)dst, src, num);
  } else if (std::is_same<T, float>::value) {
    __bang_move(dst, src, num * sizeof(float));
  }
}

template <typename T>
__mlu_func__ void subCvt(T *dst, float *src0, float *src1, int32_t num) {
#if __BANG_ARCH__ >= 500
  if (std::is_same<T, half>::value) {
    __asm__("sub.nram.crn.f16.f32 [%[dst]], [%[src0]], [%[src1]], %[num];" ::[dst] "r"(dst),
            [src0] "r"(src0), [src1] "r"(src1), [num] "r"(num));
  } else if (std::is_same<T, bfloat16_t>::value) {
    __asm__("sub.nram.crn.bf16.f32 [%[dst]], [%[src0]], [%[src1]], %[num];" ::[dst] "r"(dst),
            [src0] "r"(src0), [src1] "r"(src1), [num] "r"(num));
  } else {
    __asm__("sub.nram.crn.f32 [%[dst]], [%[src0]], [%[src1]], %[num];" ::[dst] "r"(dst),
            [src0] "r"(src0), [src1] "r"(src1), [num] "r"(num));
  }
#else
  __bang_sub((float *)dst, src0, src1, num);
  floatTo(dst, (float *)dst, num);
#endif
}

// task_length在非decoder模式下，为每个实际的block_seq_len的长度；在decoder模式下，是batch*head_num的长度。
// out_task_stride/block_task_stride在非deocder模式下，为seq_stride；在decoder模式下，是head_stride。
// out = out - sigmoid(block_lse - lse) * (out - block_out)
// lse = lse - logsigmoid(lse - block_lse)
// sigmoid: y = 1 / (1 + e ^ (-x))
// logsigmoid: y = log(1 / (1 + e ^ (-x))) = -log(1 + e ^ (-x))
template <typename T>
__mlu_func__ void updateOutAndLse(T *out,
                                  float *lse,
                                  const T *block_out,
                                  const float *block_lse,
                                  int32_t head_size,
                                  int32_t task_length,
                                  int64_t out_task_stride,
                                  int64_t block_task_stride) {
  constexpr int32_t tmp = 0x3fb8aa3b;
  const float log2e = *(float *)&tmp;
  const float neg_log2e = (-1) * log2e;
  const float recip_log2e = 1 / log2e;

  constexpr bool is_fp32 = std::is_same<T, float>::value;
  constexpr int32_t buffer_num = is_fp32 ? 5 : 3;
  constexpr int32_t pad_length = is_fp32 ? 16 : 32;

  int32_t task_once =
      PAD_DOWN(NRAM_BUFFER_SIZE / ((head_size * buffer_num + 3) * sizeof(float)), pad_length);
  int32_t task_loop = (task_length + task_once - 1) / task_once;
  if (task_loop > 1) {
    task_once = PAD_UP((task_length + task_loop - 1) / task_loop, pad_length);
  }
  float *nram_out = (float *)nram_buffer;
  float *nram_block_out = nram_out + head_size * task_once * (is_fp32 + 1);
  float *nram_lse = nram_block_out + head_size * task_once * (is_fp32 + 1);
  float *nram_block_lse = nram_lse + task_once;
  float *nram_sigmoid_result = nram_block_lse + task_once;
  __attribute__((unused)) float *nram_end = nram_sigmoid_result + task_once;

  int32_t task_deal{task_once};

  for (int32_t task_i = 0; task_i < task_loop; ++task_i) {
    task_deal = std::min(task_once, task_length - task_i * task_once);
    /*
      读取 lse block_lse
    */
    __memcpy_async(nram_lse, lse, task_deal * sizeof(float), GDRAM2NRAM);
    __memcpy_async(nram_block_lse, block_lse, task_deal * sizeof(float), GDRAM2NRAM);
    __sync_io();

    /*
      计算 lse + 读取 out
      block_lse      = (block_lse - lse) * -log2e
      sigmoid_result = 1 / (pow2(block_lse) + 1)
    */
    __bang_fusion(FUSION_FSM, nram_block_lse, nram_block_lse, nram_lse, neg_log2e, task_deal,
                  task_deal);
    __bang_pow2(nram_sigmoid_result, nram_block_lse, task_deal);
    __bang_add_scalar(nram_sigmoid_result, nram_sigmoid_result, (float)1, task_deal);
    __bang_recip(nram_sigmoid_result, nram_sigmoid_result, task_deal);
    __memcpy_async(nram_out, out, head_size * sizeof(T), GDRAM2NRAM, head_size * sizeof(T),
                   out_task_stride * sizeof(T), task_deal - 1);

    __sync_io();
    /*
      转置和升位宽 out + 读取 block_out
    */
    __bang_transpose((T *)nram_out + task_deal * head_size, (T *)nram_out, task_deal, head_size);
    toFloat(nram_out, (T *)nram_out + task_deal * head_size, task_deal * head_size);
    __memcpy_async(nram_block_out, block_out, head_size * sizeof(T), GDRAM2NRAM,
                   head_size * sizeof(T), block_task_stride * sizeof(T), task_deal - 1);
    __sync_io();
    /*
      转置和升位宽 block_out
    */
    __bang_transpose((T *)nram_block_out + task_deal * head_size, (T *)nram_block_out, task_deal,
                     head_size);
    toFloat(nram_block_out, (T *)nram_block_out + task_deal * head_size, task_deal * head_size);
    /*
      bang_fusor的计算流：
        ((out - block_out) * sigmoid_result * -1 + out).tofp16()
    */
    __bang_sub(nram_block_out, nram_out, nram_block_out, task_deal * head_size);
    __bang_cycle_mul(nram_block_out, nram_block_out, nram_sigmoid_result, task_deal * head_size,
                     task_deal);
    subCvt((T *)nram_out, nram_out, nram_block_out, task_deal * head_size);
    __bang_transpose((T *)nram_out + task_deal * head_size, (T *)nram_out, head_size, task_deal);
    __sync_compute();

    __memcpy_async(out, (T *)nram_out + task_deal * head_size, head_size * sizeof(T), NRAM2GDRAM,
                   out_task_stride * sizeof(T), head_size * sizeof(T), task_deal - 1);

    /*
      算法上: lse    = lse - logsigmoid(lse - block_lse)
                    = lse - (-log(1 + e ^ (-(lse - block_lse))))
                    = lse + log(1 + e ^ (block_lse - lse))
      之前block_lse  = (block_lse - lse) * -log2e
      实际逻辑如下：
      block_lse = block_lse * -1
                = (block_lse - lse) * log2e
      block_lse = log2(pow2(block_lse) + 1) / log2e + lse
                = lse + log(1 + e ^ (block_lse - lse))
    */
    __bang_mul_scalar(nram_block_lse, nram_block_lse, -1, task_deal);
    __bang_pow2(nram_sigmoid_result, nram_block_lse, task_deal);
    __bang_add_scalar(nram_sigmoid_result, nram_sigmoid_result, (float)1.0f, task_deal);
    __bang_log2(nram_sigmoid_result, nram_sigmoid_result, task_deal);
    __bang_mul_scalar(nram_sigmoid_result, nram_sigmoid_result, recip_log2e, task_deal);

    /*
      nram_sigmoid_result 中的值为log(1 + e ^ (block_lse - lse))
      在一些数值分布场景，例如block_lse - lse大于85左右时，这个值会出现inf。
      gpu采用的是std::log1p,
      在原始公式中，log(1 / (1 + e ^ (block_lse - lse)))中的log里的数值会极限接近0，
      采用log1p会比普通log在靠近0时有更高精度。
      mlu这里由于做了对数倒数外提*-1，所以log里的数值会变为inf(如果不外提，log2(0)同样会出现inf)。

      logsigmoid在大数值场景下，是等于原值的，例如logsigmoid(-100) = -100
      所以以下逻辑用于 识别inf的值，对inf的位置进行写入原值。
    */

    __bang_ne_scalar((uint32_t *)nram_block_out, (uint32_t *)nram_sigmoid_result, (uint32_t)INF,
                     task_deal);
    __bang_mul((uint32_t *)nram_sigmoid_result, (uint32_t *)nram_sigmoid_result,
               (uint32_t *)nram_block_out, task_deal);
    __bang_not((uint32_t *)nram_block_out, (uint32_t *)nram_block_out, task_deal);
    __bang_mul((uint32_t *)nram_block_lse, (uint32_t *)nram_block_lse, (uint32_t *)nram_block_out,
               task_deal);
    __bang_mul_scalar(nram_block_lse, nram_block_lse, recip_log2e,
                      task_deal);  // block_lse里的原值是*了log2e的，这里需要除回去。
    __bang_fusion(FUSION_FAA, nram_lse, nram_lse, nram_sigmoid_result, nram_block_lse, task_deal,
                  task_deal);

    __sync_compute();
    __memcpy_async(lse, nram_lse, task_deal * sizeof(float), NRAM2GDRAM);

    lse += task_deal;
    block_lse += task_deal;
    out += task_deal * out_task_stride;
    block_out += task_deal * out_task_stride;
  }
}

// 非decoder模式，采用taskDimZ拆分batch，taskDimY拆分head，每个task内部进行block_seq_len的循环
template <typename T>
__mlu_global__ void MluUpdateOutAndLse(void *out,
                                       float *lse,
                                       const void *block_out,
                                       const float *block_lse,
                                       const int32_t *seq_offsets,
                                       const int32_t *cu_seqs,
                                       const int32_t *block_cu_seqs,
                                       const int32_t batch,
                                       const int32_t head_num,
                                       const int32_t head_size,
                                       const int32_t max_seq_len,
                                       const int32_t block_seq_len,
                                       const bool packed,
                                       const int64_t bs_stride,
                                       const int64_t seq_stride,
                                       const int64_t head_stride,
                                       const int64_t block_bs_stride,
                                       const int64_t block_seq_stride,
                                       const int64_t block_head_stride) {
#if __BANG_ARCH__ > 300
  if (!(std::is_same<T, bfloat16_t>::value && __BANG_ARCH__ < 500)) {
    int64_t kernel_out_offset = 0;
    int64_t kernel_block_out_offset = 0;
    int64_t kernel_lse_offset = 0;
    int64_t kernel_block_lse_offset = 0;
    int64_t kernel_seq_offset = 0;

    int32_t block_seq_len_real{block_seq_len};

    if (seq_offsets != nullptr) {
      kernel_seq_offset = __load_gdram(seq_offsets + taskIdZ);
    }

    if (!packed) {
      kernel_out_offset =
          taskIdZ * bs_stride + taskIdY * head_stride + kernel_seq_offset * seq_stride;
      kernel_block_out_offset = taskIdZ * block_bs_stride + taskIdY * block_head_stride;
      kernel_lse_offset =
          taskIdZ * max_seq_len * head_num + taskIdY * max_seq_len + kernel_seq_offset;
      kernel_block_lse_offset = taskIdZ * block_seq_len * head_num + taskIdY * block_seq_len;
    } else {
      int32_t block_seq_begin = __load_gdram(block_cu_seqs + taskIdZ);
      int32_t block_seq_end = __load_gdram(block_cu_seqs + taskIdZ + 1);
      int32_t out_seq_begin = __load_gdram(cu_seqs + taskIdZ);
      block_seq_len_real = block_seq_end - block_seq_begin;
      kernel_out_offset = (out_seq_begin + kernel_seq_offset) * seq_stride + taskIdY * head_stride;
      kernel_block_out_offset = block_seq_begin * block_seq_stride + taskIdY * head_stride;
      kernel_lse_offset =
          taskIdZ * max_seq_len * head_num + taskIdY * max_seq_len + kernel_seq_offset;
      kernel_block_lse_offset = taskIdZ * block_seq_len * head_num + taskIdY * block_seq_len;
    }

    auto kernel_out = (T *)out + kernel_out_offset;
    auto kernel_lse = lse + kernel_lse_offset;
    auto kernel_block_out = (T *)block_out + kernel_block_out_offset;
    auto kernel_block_lse = block_lse + kernel_block_lse_offset;

    updateOutAndLse<T>(kernel_out, kernel_lse, kernel_block_out, kernel_block_lse, head_size,
                       block_seq_len_real, seq_stride, block_seq_stride);
  }
#endif
}

// decoder模式下，采用launch 所有core，内部拆分batch*head_num维度，
// 每个core处理 batch*head_num/taskDimX
template <typename T>
__mlu_global__ void MluUpdateOutAndLseDecoder(void *out,
                                              float *lse,
                                              const void *block_out,
                                              const float *block_lse,
                                              const int32_t *seq_offsets,
                                              const int32_t *cu_seqs,
                                              const int32_t *block_cu_seqs,
                                              const int32_t batch,
                                              const int32_t head_num,
                                              const int32_t head_size,
                                              const int32_t max_seq_len,
                                              const int32_t block_seq_len,
                                              const bool packed,
                                              const int64_t bs_stride,
                                              const int64_t seq_stride,
                                              const int64_t head_stride,
                                              const int64_t block_bs_stride,
                                              const int64_t block_seq_stride,
                                              const int64_t block_head_stride) {
#if __BANG_ARCH__ > 300
  if (!(std::is_same<T, bfloat16_t>::value && __BANG_ARCH__ < 500)) {
    int32_t task_length{0}, task_begin{0};
    splitTask(batch * head_num, task_length, task_begin);
    if (task_length <= 0) {
      return;
    }

    int32_t batch_idx = task_begin / head_num;
    int32_t head_idx = task_begin % head_num;

    auto kernel_out_offset = batch_idx * bs_stride + head_idx * head_stride;
    auto kernel_block_out_offset = batch_idx * block_bs_stride + head_idx * block_head_stride;
    auto kernel_lse_offset = batch_idx * max_seq_len * head_num + head_idx * max_seq_len;
    auto kernel_block_lse_offset = batch_idx * block_seq_len * head_num + head_idx * block_seq_len;

    auto kernel_out = (T *)out + kernel_out_offset;
    auto kernel_lse = lse + kernel_lse_offset;
    auto kernel_block_out = (T *)block_out + kernel_block_out_offset;
    auto kernel_block_lse = block_lse + kernel_block_lse_offset;

    updateOutAndLse<T>(kernel_out, kernel_lse, kernel_block_out, kernel_block_lse, head_size,
                       task_length, head_stride, block_head_stride);
  }
#endif
}

#if __BANG_ARCH__ < 500
template <>
__mlu_global__ void MluUpdateOutAndLseDecoder<bfloat16_t>(void *out,
                                                          float *lse,
                                                          const void *block_out,
                                                          const float *block_lse,
                                                          const int32_t *seq_offsets,
                                                          const int32_t *cu_seqs,
                                                          const int32_t *block_cu_seqs,
                                                          const int32_t batch,
                                                          const int32_t head_num,
                                                          const int32_t head_size,
                                                          const int32_t max_seq_len,
                                                          const int32_t block_seq_len,
                                                          const bool packed,
                                                          const int64_t bs_stride,
                                                          const int64_t seq_stride,
                                                          const int64_t head_stride,
                                                          const int64_t block_bs_stride,
                                                          const int64_t block_seq_stride,
                                                          const int64_t block_head_stride) {}

template <>
__mlu_global__ void MluUpdateOutAndLse<bfloat16_t>(void *out,
                                                   float *lse,
                                                   const void *block_out,
                                                   const float *block_lse,
                                                   const int32_t *seq_offsets,
                                                   const int32_t *cu_seqs,
                                                   const int32_t *block_cu_seqs,
                                                   const int32_t batch,
                                                   const int32_t head_num,
                                                   const int32_t head_size,
                                                   const int32_t max_seq_len,
                                                   const int32_t block_seq_len,
                                                   const bool packed,
                                                   const int64_t bs_stride,
                                                   const int64_t seq_stride,
                                                   const int64_t head_stride,
                                                   const int64_t block_bs_stride,
                                                   const int64_t block_seq_stride,
                                                   const int64_t block_head_stride) {}
#endif
}  // namespace kernels

inline int32_t dtype_index(cnnlDataType_t dtype) {
  switch (dtype) {
    case CNNL_DTYPE_HALF:
      return 0;
      break;
    case CNNL_DTYPE_BFLOAT16:
      return 1;
      break;
    case CNNL_DTYPE_FLOAT:
      return 2;
      break;
    default:
      return 0;
      break;
  }
}

KernelStatus invokeUpdateOutAndLse(cnrtQueue_t queue,
                                   void *out,
                                   float *lse,
                                   const void *block_out,
                                   const float *block_lse,
                                   const int32_t *seq_offsets,
                                   const int32_t *cu_seqs,
                                   const int32_t *block_cu_seqs,
                                   const int32_t batch,
                                   const int32_t head_num,
                                   const int32_t head_size,
                                   const int32_t max_seq_len,
                                   const int32_t block_seq_len,
                                   const int64_t bs_stride,
                                   const int64_t seq_stride,
                                   const int64_t head_stride,
                                   const int64_t block_bs_stride,
                                   const int64_t block_seq_stride,
                                   const int64_t block_head_stride,
                                   const bool packed,
                                   const cnnlDataType_t dtype) {
  void (*update_out_and_lse_kernels[])(
      void *, float *, const void *, const float *, const int32_t *, const int32_t *,
      const int32_t *, const int32_t, const int32_t, const int32_t, const int32_t, const int32_t,
      const bool, const int64_t, const int64_t, const int64_t, const int64_t, const int64_t,
      const int64_t) = {kernels::MluUpdateOutAndLse<half>,
                        kernels::MluUpdateOutAndLse<bfloat16_t>,
                        kernels::MluUpdateOutAndLse<float>,
                        kernels::MluUpdateOutAndLseDecoder<half>,
                        kernels::MluUpdateOutAndLseDecoder<bfloat16_t>,
                        kernels::MluUpdateOutAndLseDecoder<float>};

  // 非decoder模式，采用taskDimZ拆分batch，taskDimY拆分head，每个task内部进行block_seq_len的循环
  uint32_t task_dimx = 1;
  uint32_t task_dimy = head_num;
  uint32_t task_dimz = batch;

  bool decoder_mode = (block_seq_len == 1 && max_seq_len == 1);
  if (decoder_mode) {
    // decoder模式下，采用launch 所有core，内部拆分batch*head_num维度，
    // 每个core处理 batch*head_num/taskDimX
    CNdev dev;
    cnCtxGetDevice(&dev);
    int32_t cluster_num;
    CNRT_CHECK(cnrtDeviceGetAttribute(&cluster_num, cnrtAttrClusterCount, dev));
    int32_t core_num;
    CNRT_CHECK(cnrtDeviceGetAttribute(&core_num, cnrtAttrMcorePerCluster, dev));

    task_dimx = std::min(PAD_UP(batch * head_num, 16) / 16, cluster_num * core_num);
    task_dimy = 1;
    task_dimz = 1;
  }

  cnrtDim3_t task_dim = {task_dimx, task_dimy, task_dimz};
  cnrtFunctionType_t func_type = cnrtFuncTypeBlock;
  int32_t kernel_index = dtype_index(dtype) + decoder_mode * 3;

  if (dtype == CNNL_DTYPE_BFLOAT16 && !isBf16Supported()) {
    std::cerr << "[invokeUpdateOutAndLse]: MLU300 devices do not support bfloat16." << std::endl;
    return KernelStatus::KERNEL_STATUS_FAILED;
  }

  update_out_and_lse_kernels[kernel_index]<<<task_dim, func_type, queue>>>(
      out, lse, block_out, block_lse, seq_offsets, cu_seqs, block_cu_seqs, batch, head_num,
      head_size, max_seq_len, block_seq_len, packed, bs_stride, seq_stride, head_stride,
      block_bs_stride, block_seq_stride, block_head_stride);

  return KernelStatus::KERNEL_STATUS_SUCCESS;
}
}  // namespace tmo
