#include <algorithm>
#include <cassert>
#include <iostream>
#include <map>
#include <ostream>
#include "add_bias_activation.mluh"
#include "cnnl.h"
#include "cnrt.h"
// clang-format off
#include <mlu.h>
// clang-format on

namespace tmo {

namespace kernels {
#define PAD_UP(x, y) (((x) / (y) + (int)((x) % (y) > 0)) * (y))
#define USE_NRAM_SIZE (__MLU_NRAM_SIZE__ * 1024 - 20 * 1024)
#define USE_SRAM_SIZE (__MLU_SRAM_SIZE__ * 1024 - 20 * 1024)

__nram__ int8_t nram_buffer[USE_NRAM_SIZE];
__mlu_shared__ int8_t sram_buffer[USE_SRAM_SIZE];

__mlu_func__ void get_expert_info(int *nram_count,
                                  int *count_sram,
                                  uint32_t *gather_offset,
                                  int real_inner,
                                  int tokens_start,
                                  int tokens_end,
                                  int tokens_load,
                                  int expert_deal_start,
                                  int expert_deal_end,
                                  uint32_t &expert_start,
                                  uint32_t &expert_end,
                                  uint32_t &tokens_deal_first,
                                  int dtype_size) {
  bool record_start = false;
  // loop expert to find first and last deal expert in current core
  for (int expert_id = expert_deal_start; expert_id <= expert_deal_end; expert_id++) {
    if (__load_nram(nram_count + expert_id + 1) > tokens_start && !record_start) {
      expert_start = expert_id;
      tokens_deal_first =
          std::min(__load_nram(nram_count + expert_id + 1) - 1, tokens_end) - tokens_start + 1;
      record_start = true;
    }
    if (__load_nram(nram_count + expert_id + 1) > tokens_end) {
      expert_end = expert_id;
      break;
    }
  }

  // record expert offset to gather bias
  __bang_write_zero(gather_offset, tokens_load);
  int tokens_load_total = 0;
  for (int expert_id = expert_start; expert_id <= expert_end; expert_id++) {
    int tokens_expand = __load_sram((int *)count_sram + expert_id);
    if (expert_id == expert_start) {
      tokens_expand = tokens_deal_first;
    } else if (expert_id == expert_end) {
      tokens_expand = tokens_load - tokens_load_total;
    }
    if (tokens_expand == 0) {
      continue;
    }
    __bang_write_value(gather_offset + tokens_load_total, tokens_expand,
                       (int)((expert_id - expert_deal_start) * real_inner * dtype_size));
    tokens_load_total += tokens_expand;
  }
}

/*************** functions for compute basic operation ***************/

template <typename T>
__mlu_func__ void add_bias(T *dst_src, T *bias, int number) {
  // cycle add bias
  __bang_add((T *)dst_src, (T *)dst_src, (T *)bias, number);
}

template <>
__mlu_func__ void add_bias(bfloat16_t *dst_src, bfloat16_t *bias, int number) {
#if __BANG_ARCH__ > 500
  __bang_add((bfloat16_t *)dst_src, (bfloat16_t *)dst_src, (bfloat16_t *)bias, number);
#endif
}

template <typename T>
__mlu_func__ void mul_left_right(T *left, T *right, int number) {
  __bang_mul((T *)left, (T *)left, (T *)right, number);
}

template <>
__mlu_func__ void mul_left_right(bfloat16_t *left, bfloat16_t *right, int number) {
#if __BANG_ARCH__ > 500
  __bang_mul((bfloat16_t *)left, (bfloat16_t *)left, (bfloat16_t *)right, number);
#endif
}

__mlu_func__ void do_activation(float *input_left,
                                float *act_space,
                                int number,
                                float active_coef,
                                cnnlActivationMode_t act_type) {
  if (act_type == CNNL_ACTIVATION_GELU) {
    __bang_active_gelu((float *)input_left, (float *)input_left, number);
  } else if (act_type == CNNL_ACTIVATION_SWISH) {
    float *tmp = input_left;
    if (active_coef != 1.0f) {
      __bang_mul_scalar(act_space, input_left, active_coef, number);
      tmp = act_space;
    }
    __bang_active_sigmoid((float *)act_space, (float *)tmp, number);
    __bang_mul((float *)input_left, (float *)input_left, (float *)act_space, number);
  }
}

/*************** functions for steps of each loop ***************/

template <typename T>
__mlu_func__ void gather_bias(T *bias_sram,
                              T *bias_nram,
                              uint32_t *gather_offset,
                              int expert_start,
                              int expert_end,
                              int expert_deal_start,
                              int tokens_deal_first,
                              int tokens_deal,
                              int inner_size,
                              bool is_gated) {
#if __BANG_ARCH__ > 500
  if (is_gated) {
    __gather_async((T *)bias_nram, (T *)bias_sram, gather_offset, inner_size * sizeof(T), SRAM2NRAM,
                   inner_size * sizeof(T), tokens_deal);
    __gather_async((T *)bias_nram + tokens_deal * inner_size, (T *)bias_sram + inner_size,
                   gather_offset, inner_size * sizeof(T), SRAM2NRAM, inner_size * sizeof(T),
                   tokens_deal);
  } else {
    __gather_async((T *)bias_nram, (T *)bias_sram, gather_offset, inner_size * sizeof(T), SRAM2NRAM,
                   inner_size * sizeof(T), tokens_deal);
  }
#else
  for (int i = 0; i < tokens_deal; i++) {
    if (is_gated) {
      __memcpy_async((T *)bias_nram + i * inner_size, (int8_t *)bias_sram + gather_offset[i],
                     inner_size * sizeof(T), SRAM2NRAM);
      __memcpy_async((T *)bias_nram + (tokens_deal * inner_size + i * inner_size),
                     (int8_t *)bias_sram + inner_size * sizeof(T) + gather_offset[i],
                     inner_size * sizeof(T), SRAM2NRAM);
    } else {
      __memcpy_async((T *)bias_nram + i * inner_size, (int8_t *)bias_sram + gather_offset[i],
                     inner_size * sizeof(T), SRAM2NRAM);
    }
  }
#endif
}

template <typename T>
__mlu_func__ void loadBiasInput(T *input,
                                T *left,
                                T *right,
                                T *bias_nram,
                                T *bias_sram,
                                uint32_t *gather_offset,
                                size_t input_offset,
                                int tokens_deal,
                                int inner_size,
                                uint32_t expert_start,
                                uint32_t expert_end,
                                int expert_deal_start,
                                uint32_t tokens_deal_first,
                                bool is_gated,
                                bool has_bias) {
  if (is_gated) {
    // if gated, stride io load input, left/right inner to input_left/right
    __memcpy_async((T *)left, (T *)input + input_offset, inner_size * sizeof(T), GDRAM2NRAM,
                   inner_size * sizeof(T), inner_size * 2 * sizeof(T), tokens_deal - 1);
    __memcpy_async((T *)right, (T *)input + input_offset + inner_size, inner_size * sizeof(T),
                   GDRAM2NRAM, inner_size * sizeof(T), inner_size * 2 * sizeof(T), tokens_deal - 1);
  } else {
    // if not gated, load input to input_left total
    __memcpy_async((T *)left, (T *)input + input_offset, tokens_deal * inner_size * sizeof(T),
                   GDRAM2NRAM);
  }
  if (has_bias) {
    __sync_compute();
    gather_bias((T *)bias_sram, (T *)bias_nram, gather_offset, expert_start, expert_end,
                expert_deal_start, tokens_deal_first, tokens_deal, inner_size, is_gated);
  }
}

template <typename T>
__mlu_func__ void computeAddActivation(T *bias_nram,
                                       T *left_dst,
                                       T *input_right,
                                       float *input_left,
                                       float *act_space,
                                       int tokens_deal,
                                       int inner_size,
                                       bool is_gated,
                                       bool has_bias,
                                       float active_coef,
                                       cnnlActivationMode_t act_type) {
  int number = tokens_deal * inner_size;
  if (has_bias) {
    add_bias((T *)left_dst, (T *)bias_nram, number);
    if (is_gated) {
      add_bias((T *)input_right, (T *)bias_nram + tokens_deal * inner_size, number);
    }
  }

  // cast half/bfloat16 to float to acvication, if float, left_dst is same as input_left
  if (std::is_same<T, half>::value) {
    __bang_half2float((float *)input_left, (half *)left_dst, number);
  }
#if __BANG_ARCH__ > 500
  if (std::is_same<T, bfloat16_t>::value) {
    __bang_bfloat162float((float *)input_left, (bfloat16_t *)left_dst, number);
  }
#endif

  // activation
  do_activation(input_left, act_space, number, active_coef, act_type);

  // if half/bfloat16, cast float to T to mul
  if (std::is_same<T, half>::value) {
    __bang_float2half((half *)input_left, (float *)input_left, number);
  }
#if __BANG_ARCH__ > 500
  if (std::is_same<T, bfloat16_t>::value) {
    __bang_float2bfloat16((bfloat16_t *)input_left, (float *)input_left, number);
  }
#endif

  if (is_gated) {
    mul_left_right((T *)input_left, (T *)input_right, tokens_deal * inner_size);
  }
}

template <typename T>
__mlu_func__ void storeOutput(T *output,
                              T *output_nram,
                              size_t output_offset,
                              int output_stride,
                              int tokens_deal,
                              int inner_size) {
  __memcpy_async((T *)output + output_offset, (T *)output_nram, inner_size * sizeof(T), NRAM2GDRAM,
                 output_stride * sizeof(T), inner_size * sizeof(T), tokens_deal - 1);
}

template <typename T>
__mlu_global__ void MLUAddBiasActivationKernel(T *output,
                                               const T *input,
                                               const T *bias,
                                               const int *cusum_token_count,
                                               int num_expert,
                                               int total_tokens,
                                               int inner_size,
                                               int output_stride,
                                               bool is_gated,
                                               cnnlActivationMode_t act_type,
                                               int start_expert_id,
                                               int expert_size,
                                               float active_coef) {
  // if bias and token_count is nullptr, not add bias, only activation and gated mul.
  bool has_bias = (bias != nullptr);

  // 1. distrubute nram space
  /* if not gated
      ---------------------------- ----------------------------------
     |       nram_token_count     |          bias_ping/pong          |
     |  num_expert * sizeof(int)  |  2 * x * inner_size * sizeof(T)  |
      ---------------------------- ----------------------------------
      --------------------------------------
     |           input_ping/pong            |
     |  2 * x * inner_size * sizeof(float)  |
      --------------------------------------
      ------------------------------------------------- -------------------
     |                   act_space                     |   gather_offset   |
     |  gelu: 0; silu: x * inner_size * sizeof(float)  |  x * sizeof(int)  |
      ------------------------------------------------- -------------------
  */

  /* if gated
      ---------------------------- ----------------------------------
     |       nram_token_count     |          bias_ping/pong          |
     |  num_expert * sizeof(int)  |  2 * x * real_inner * sizeof(T)  |
      ---------------------------- ----------------------------------
      -------------------------------------- ----------------------------------
     |        input_left_ping/pong          |       input_right_ping/pong      |
     |  2 * x * inner_size * sizeof(float)  |  2 * x * inner_size * sizeof(T)  |
      -------------------------------------- ----------------------------------
      ------------------------------------------------- -------------------
     |                   act_space                     |   gather_offset   |
     |  gelu: 0; silu: x * inner_size * sizeof(float)  |  x * sizeof(int)  |
      ------------------------------------------------- -------------------
 */

  // distribute sram
  int8_t *count_sram = (int8_t *)sram_buffer;
  int8_t *bias_sram = (int8_t *)count_sram + num_expert * sizeof(int);

  // distrubute nram
  int real_inner = (is_gated) ? inner_size * 2 : inner_size;
  int bias_nram_size = (has_bias) ? real_inner * sizeof(T) : 0;
  int act_space_size = (act_type == CNNL_ACTIVATION_GELU) ? 0 : inner_size * sizeof(float);
  int gated_ext_size = is_gated ? sizeof(T) : 0;
  int max_token_deal = (USE_NRAM_SIZE - (num_expert + 1) * sizeof(int)) /
                       (2 * inner_size * (sizeof(float) + gated_ext_size) + act_space_size +
                        2 * bias_nram_size + sizeof(int));

  int8_t *nram_count = (int8_t *)nram_buffer;
  int8_t *bias_nram = (int8_t *)nram_count + (num_expert + 1) * sizeof(int);
  int8_t *input_left = (int8_t *)bias_nram + 2 * max_token_deal * bias_nram_size;
  int8_t *input_right =
      (int8_t *)input_left + 2 * ((is_gated) ? max_token_deal * inner_size * sizeof(float) : 0);
  int8_t *act_space = (int8_t *)input_right +
                      2 * max_token_deal * inner_size * (is_gated ? sizeof(T) : sizeof(float));
  int8_t *gather_offset = (int8_t *)act_space + max_token_deal * act_space_size;

  // 2. cusum_token_count load to nram, because need to reuse in load bias.
  if (has_bias) {
    __memcpy((int *)nram_count, (int *)cusum_token_count, (num_expert + 1) * sizeof(int),
             GDRAM2NRAM);
    if (taskIdX == 0) {
      __bang_sub((int *)bias_nram, (int *)nram_count + 1, (int *)nram_count, num_expert);
      __sync();
      __memcpy((int *)count_sram, (int *)bias_nram, num_expert * sizeof(int), NRAM2SRAM);
    }
    __sync_cluster();
  }

  // 3. sram loop to compute
  // compute once load bias to sram due to sram_limit
  int max_expert_deal = (USE_SRAM_SIZE - num_expert * sizeof(int)) / (real_inner * sizeof(T));
  int real_expert = cusum_token_count == nullptr ? num_expert : expert_size;
  int sram_loop_rem = real_expert % max_expert_deal;
  int sram_loop = real_expert / max_expert_deal + (int)(sram_loop_rem != 0);
  if (!has_bias) {
    max_expert_deal = real_expert;
    sram_loop = 1;
    sram_loop_rem = 0;
  }
  for (int deal_loop = 0; deal_loop < sram_loop; deal_loop++) {
    // load current bias, compute each core deal number
    int expert_deal =
        (deal_loop == (sram_loop - 1) && sram_loop_rem != 0) ? sram_loop_rem : max_expert_deal;
    int expert_deal_start = deal_loop * max_expert_deal + start_expert_id;
    int expert_deal_end = expert_deal_start + expert_deal - 1;
    __sync_all();
    if (has_bias && __is_mpu()) {
      __memcpy((T *)bias_sram, (T *)bias + deal_loop * max_expert_deal * real_inner,
               expert_deal * real_inner * sizeof(T), GDRAM2SRAM);
    }
    __sync_all();

    // get tokens info of each core
    int tokens_total_cur = total_tokens;
    if (has_bias) {
      tokens_total_cur =
          __load_nram((int *)nram_count + expert_deal_end + 1) -
          (start_expert_id == 0 ? 0 : __load_nram((int *)nram_count + expert_deal_start));
    } else if (cusum_token_count != nullptr) {
      tokens_total_cur = __load_gdram(cusum_token_count + expert_deal_end + 1) -
                         __load_gdram(cusum_token_count + expert_deal_start);
    }

    if (sram_loop != 1) {
      tokens_total_cur = ((int *)nram_count)[expert_deal_end + 1] -
                         ((deal_loop == 0) ? 0 : ((int *)nram_count)[expert_deal_start]);
      if (deal_loop == 0 && start_expert_id != 0) {
        tokens_total_cur -= __load_nram((int *)nram_count + expert_deal_start);
      }
    }

    int tokens_core_rem = tokens_total_cur % taskDim;
    int tokens_cur_core = tokens_total_cur / taskDim + (taskId < tokens_core_rem);
    if (tokens_cur_core == 0) {
      continue;
    }

    // if ep, input start in current token, have a real start in total network
    int real_start =
        cusum_token_count != nullptr ? __load_gdram((int *)cusum_token_count + start_expert_id) : 0;
    int tokens_core_start = tokens_cur_core * taskId +
                            (taskId < tokens_core_rem ? 0 : tokens_core_rem) +
                            ((deal_loop == 0) ? 0 : ((int *)nram_count)[expert_deal_start]);
    if (deal_loop != 0) {
      tokens_core_start -= real_start;
    }

    uint32_t expert_start = 0;
    uint32_t expert_end = 0;
    uint32_t tokens_deal_first = 0;

    // 4. nram loop compute
    int nram_loop_rem = tokens_cur_core % max_token_deal;
    int nram_loop = tokens_cur_core / max_token_deal + (int)(nram_loop_rem != 0);
    int tokens_load = max_token_deal;
    int tokens_compute = max_token_deal;
    int tokens_store = max_token_deal;

    for (int loop = 0; loop < nram_loop + 2; loop++) {
      int inner_io_offset = (loop % 2) * max_token_deal * inner_size;
      int inner_com_offset = ((loop + 1) % 2) * max_token_deal * inner_size;
      int real_io_offset = (loop % 2) * max_token_deal * real_inner;
      int real_com_offset = ((loop + 1) % 2) * max_token_deal * real_inner;
      if (nram_loop_rem != 0) {
        if (loop > 1 && (loop - 2) == (nram_loop - 1)) {
          tokens_store = nram_loop_rem;
        }
        if (loop > 0 && (loop - 1) == (nram_loop - 1)) {
          tokens_compute = nram_loop_rem;
        }
        if (loop == (nram_loop - 1)) {
          tokens_load = nram_loop_rem;
        }
      }
      int tokens_cur_start = tokens_core_start + loop * max_token_deal;
      int tokens_cur_end = tokens_cur_start + tokens_load - 1;
      // get current load info
      if (loop < nram_loop && has_bias) {
        get_expert_info((int *)nram_count, (int *)count_sram, (uint32_t *)gather_offset, real_inner,
                        tokens_cur_start + real_start, tokens_cur_end + real_start, tokens_load,
                        expert_deal_start, expert_deal_end, expert_start, expert_end,
                        tokens_deal_first, sizeof(T));
      }

      // store
      if (loop > 1) {
        size_t output_offset = (tokens_core_start + (loop - 2) * max_token_deal) * output_stride;
        storeOutput((T *)output, (T *)((float *)input_left + inner_io_offset), output_offset,
                    output_stride, tokens_store, inner_size);
      }

      // compute
      if (loop > 0 && loop <= nram_loop) {
        T *left_dst = (T *)((float *)input_left + inner_com_offset) +
                      ((std::is_same<T, float>::value) ? 0 : tokens_compute * inner_size);
        computeAddActivation((T *)bias_nram + real_com_offset, (T *)left_dst,
                             (T *)input_right + inner_com_offset,
                             (float *)input_left + inner_com_offset, (float *)act_space,
                             tokens_compute, inner_size, is_gated, has_bias, active_coef, act_type);
      }

      // load
      if (loop < nram_loop) {
        T *left_dst = (T *)((float *)input_left + inner_io_offset) +
                      ((std::is_same<T, float>::value) ? 0 : tokens_load * inner_size);
        size_t input_offset = tokens_cur_start * real_inner;
        loadBiasInput((T *)input, (T *)left_dst, (T *)input_right + inner_io_offset,
                      (T *)bias_nram + real_io_offset, (T *)bias_sram, (uint32_t *)gather_offset,
                      input_offset, tokens_load, inner_size, expert_start, expert_end,
                      expert_deal_start, tokens_deal_first, is_gated, has_bias);
      }
      __sync();
    }
  }
}

}  // namespace kernels

KernelStatus invokeGroupAddBiasActivationKernel(cnrtQueue_t queue,
                                                void *output,
                                                const void *input,
                                                const void *bias,
                                                const int *cusum_token_count,
                                                int num_expert,
                                                int total_tokens,
                                                int inner_size,
                                                int output_stride,
                                                cnnlDataType_t dtype,
                                                bool is_gated,
                                                cnnlActivationMode_t act_type,
                                                int start_expert_id,
                                                int expert_size,
                                                float active_coef) {
  if (bias != NULL && cusum_token_count == NULL) {
    std::cerr << "[invokeGroupAddBiasActivationKernel]: "
              << "when have bias, cusum_token_count can not be nullptr.";
    return KernelStatus::KERNEL_STATUS_FAILED;
  }
  if (act_type != CNNL_ACTIVATION_GELU && act_type != CNNL_ACTIVATION_SWISH) {
    std::cerr << "[invokeGroupAddBiasActivationKernel]: "
              << "activation mode only supports gelu and swish now.";
    return KernelStatus::KERNEL_STATUS_FAILED;
  }
  CNdev dev;
  cnCtxGetDevice(&dev);
  int cluster_num;
  int core_num;
  CNRT_CHECK(cnrtDeviceGetAttribute(&cluster_num, cnrtAttrClusterCount, dev));
  CNRT_CHECK(cnrtDeviceGetAttribute(&core_num, cnrtAttrMcorePerCluster, dev));
  cnrtDim3_t dim{.x = (uint32_t)core_num, .y = (uint32_t)cluster_num, .z = 1};

  if (dtype == CNNL_DTYPE_FLOAT) {
    kernels::MLUAddBiasActivationKernel<<<dim, cnrtFuncTypeUnion1, queue>>>(
        (float *)output, (const float *)input, (const float *)bias, cusum_token_count, num_expert,
        total_tokens, inner_size, output_stride, is_gated, act_type, start_expert_id, expert_size,
        active_coef);
  } else if (dtype == CNNL_DTYPE_HALF) {
    kernels::MLUAddBiasActivationKernel<<<dim, cnrtFuncTypeUnion1, queue>>>(
        (half *)output, (const half *)(input), (const half *)bias, cusum_token_count, num_expert,
        total_tokens, inner_size, output_stride, is_gated, act_type, start_expert_id, expert_size,
        active_coef);
  } else if (dtype == CNNL_DTYPE_BFLOAT16) {
    if (!isBf16Supported()) {
      std::cerr << "[invokeGroupAddBiasActivationKernel]: MLU300 devices do not support bfloat16."
                << std::endl;
      return KernelStatus::KERNEL_STATUS_FAILED;
    }
    kernels::MLUAddBiasActivationKernel<<<dim, cnrtFuncTypeUnion1, queue>>>(
        (bfloat16_t *)output, (const bfloat16_t *)input, (const bfloat16_t *)bias,
        cusum_token_count, num_expert, total_tokens, inner_size, output_stride, is_gated, act_type,
        start_expert_id, expert_size, active_coef);
  } else {
    std::cerr << "[invokeGroupAddBiasActivationKernel]: add_bias_activation data_type not support, "
              << "only support float/half/bfloat16." << std::endl;
    return KernelStatus::KERNEL_STATUS_FAILED;
  }
  return KernelStatus::KERNEL_STATUS_SUCCESS;
}

}  // namespace tmo
