/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/
#include <climits>
#include "quant_to_paged_cache.mluh"

namespace tmo {
namespace kernels {

#define sizeof_(T) (uint32_t)sizeof(T)
#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b))

#define REM_FOR_STACK (32 * 1024)
__nram__ int8_t nram_buffer[__MLU_NRAM_SIZE__ * 1024 - REM_FOR_STACK];
#if __BANG_ARCH__ > 500
__nram__ const int nram_range_32[32] = {0,  1,  2,  3,  4,  5,  6,  7,  8,  9,  10,
                                        11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21,
                                        22, 23, 24, 25, 26, 27, 28, 29, 30, 31};
#endif

template <typename T>
__mlu_func__ void quantifyToInt8(int8_t *nram_output,
                                 float *nram_float,
                                 T *nram_input,
                                 float *nram_scale,
                                 float *nram_temp,
                                 int seq_len,
                                 int head_size) {
  // quantify
  __bang_transpose((T *)nram_input, (T *)nram_scale, seq_len, head_size);
  if (std::is_same<half, T>::value) {
    __bang_half2float(nram_float, (half *)nram_input, head_size * seq_len);
  }
  if (std::is_same<bfloat16_t, T>::value) {
#if __BANG_ARCH__ > 500
    __bang_bfloat162float(nram_float, (bfloat16_t *)nram_input, head_size * seq_len);
#endif
  }
  __bang_abs(nram_scale, nram_float, head_size * seq_len);
  __bang_maxpool(nram_scale, nram_scale, seq_len, head_size, 1, head_size, 1, 1, 1);
  __bang_mul_scalar(nram_scale, nram_scale, 1 / 127.f, seq_len);
  __bang_recip(nram_temp, nram_scale, seq_len);
  __bang_cycle_mul(nram_float, nram_float, nram_temp, head_size * seq_len, seq_len);
  __bang_float2int8_rn((int8_t *)nram_float, nram_float, head_size * seq_len, 0);
  __bang_transpose((int8_t *)nram_output, (int8_t *)nram_float, head_size, seq_len);
}

template <typename T>
__mlu_global__ void MLUQuantToPagedCacheKernel(T *key,
                                               T *value,
                                               int8_t *key_cache,
                                               int8_t *value_cache,
                                               float *key_cache_scale,
                                               float *value_cache_scale,
                                               int *slot_mapping,
                                               size_t key_stride0,
                                               size_t value_stride0,
                                               int tokens_num,
                                               int head_num,
                                               int block_size,
                                               int head_size,
                                               int tokens_block) {
/*******************************************************nram space***********************
 * nram:|     input    |   scale   | cache_offset | scale_offset | mask |  temp | index |
 * input size: tokens_block * head_num * head_size * sizeof(float)
 * scale size: equal to input size
 * cache_offset size: tokens_block * head_num * sizeof(float)
 * scale_offset size: equal to cache_offset size
 * mask size: CEIL_DIV(tokens_size * head_num, 8) * sizeof(int8_t)
 * temp size: CEIL_ALIGN(token_size * head_num, 8) * sizeof(int)
 * index size: head_num * sizeof(int)
 ****************************************************************************************/
#if __BANG_ARCH__ > 500
  int token_begin = taskId * tokens_block;
  if (token_begin >= tokens_num) return;
  int token_handle = std::min(tokens_block, tokens_num - token_begin);

  int seq_len = token_handle * head_num;
  int pad8_num = CEIL_DIV(seq_len, CHAR_BIT) * CHAR_BIT;
  int input_size = seq_len * head_size * sizeof_(float);
  int8_t *nram_input = nram_buffer;
  float *nram_scale = (float *)(nram_buffer + input_size);
  int *cache_offset = (int *)((int8_t *)nram_scale + input_size);
  int *scale_offset = cache_offset + pad8_num;
  int *nram_mask = scale_offset + pad8_num;
  int *nram_temp = nram_mask + pad8_num;
  int *head_index = nram_temp + pad8_num;

  // generate range: (0, 1, 2, ..., (head_num - 1))
  __memcpy(head_index, nram_range_32, std::min(head_num, 32) * sizeof_(int), NRAM2NRAM);
  int begin = 32;
  while (begin < head_num) {
    int count = std::min(begin, head_num - begin);
    __bang_add_scalar(head_index + begin, head_index, begin, count);
    begin += count;
  }

  // load slot(token_handle) -> expand(head_num, token_handle) ->transpose(token_handle, head_num)
  int token_size = token_handle * sizeof_(int);
  __memcpy(scale_offset, slot_mapping + token_begin, token_size, GDRAM2NRAM);
  __memcpy(nram_temp, scale_offset, token_size, NRAM2NRAM, token_size, 0, head_num - 1);
  __bang_transpose(scale_offset, nram_temp, head_num, token_handle);

  __bang_write_zero((float *)nram_temp, pad8_num);
  __bang_ge_bitindex((float *)nram_mask, (float *)scale_offset, (float *)nram_temp, pad8_num);

  // calculate cache/scale scatter offset
  __bang_div(cache_offset, scale_offset, (int)block_size, seq_len);
  __bang_rem(scale_offset, scale_offset, (int)block_size, seq_len);
  __bang_mul_scalar(cache_offset, cache_offset, head_num * block_size, seq_len);
  __bang_mul_scalar(head_index, head_index, block_size, head_num);
  __bang_cycle_add(cache_offset, cache_offset, head_index, seq_len, head_num);
  __bang_add(scale_offset, cache_offset, scale_offset, seq_len);
  __bang_mul_scalar(cache_offset, scale_offset, head_size, seq_len);
  __bang_mul_scalar(scale_offset, scale_offset, sizeof_(float), seq_len);

  int hidden_bytes = head_num * head_size * sizeof_(T);
  int8_t *nram_output = (int8_t *)(nram_input + seq_len * head_size * sizeof_(half));
  T *nram_input_origin = (T *)nram_input;
  if (!std::is_same<float, T>::value) {
    nram_input_origin = (T *)(nram_input + seq_len * head_size * (sizeof_(float) - sizeof_(T)));
  }
  if (key != nullptr && key_cache != nullptr && key_cache_scale != nullptr) {
    // (token_handle, head_num, head_size)
    __memcpy(nram_scale, key + token_begin * key_stride0, hidden_bytes, GDRAM2NRAM, hidden_bytes,
             key_stride0 * sizeof_(T), token_handle - 1);
    // quantify
    quantifyToInt8(nram_output, (float *)nram_input, nram_input_origin, nram_scale,
                   (float *)nram_temp, seq_len, head_size);
    // scatter to gdram
    __scatter(key_cache, (int8_t *)nram_output, (uint32_t *)cache_offset, nram_mask, head_size,
              NRAM2GDRAM, head_size, seq_len);
    __scatter(key_cache_scale, nram_scale, (uint32_t *)scale_offset, nram_mask, sizeof_(float),
              NRAM2GDRAM, sizeof_(float), seq_len);
  }

  if (value != nullptr && value_cache != nullptr && value_cache_scale != nullptr) {
    // (token_handle, head_num, head_size)
    __memcpy(nram_scale, value + token_begin * value_stride0, hidden_bytes, GDRAM2NRAM,
             hidden_bytes, value_stride0 * sizeof_(T), token_handle - 1);
    // quantify
    quantifyToInt8(nram_output, (float *)nram_input, nram_input_origin, nram_scale,
                   (float *)nram_temp, seq_len, head_size);
    // scatter to gdram
    __scatter(value_cache, (int8_t *)nram_output, (uint32_t *)cache_offset, nram_mask, head_size,
              NRAM2GDRAM, head_size, seq_len);
    __scatter(value_cache_scale, nram_scale, (uint32_t *)scale_offset, nram_mask, sizeof_(float),
              NRAM2GDRAM, sizeof_(float), seq_len);
  }
#endif
}
}  // namespace kernels

KernelStatus invokeQuantToPagedCache(cnrtQueue_t queue,
                                     cnnlDataType_t data_type,
                                     void *key,
                                     void *value,
                                     void *key_cache,
                                     void *value_cache,
                                     void *key_cache_scale,
                                     void *value_cache_scale,
                                     void *slot_mapping,
                                     size_t key_stride0,
                                     size_t value_stride0,
                                     int num_tokens,
                                     int num_heads,
                                     int block_num,
                                     int block_size,
                                     int head_size) {
  int dtype_size = 1;
  if (data_type == CNNL_DTYPE_HALF || data_type == CNNL_DTYPE_BFLOAT16) {
    dtype_size = 2;
  } else if (data_type == CNNL_DTYPE_FLOAT) {
    dtype_size = 4;
  } else {
    std::cerr << "invokeQuantToPagedCache: unsupport data type\n";
    return KernelStatus::KERNEL_STATUS_FAILED;
  }
  int64_t kv_range = block_num * block_size * num_heads * head_size * dtype_size;
  if (kv_range > UINT32_MAX) {
    std::cerr << "invokeQuantToPagedCache: The addressing range of kv_cache cannot exceed 4G."
              << std::endl;
    return KernelStatus::KERNEL_STATUS_FAILED;
  }
  if (is_arch300()) {
    std::cerr << "[invokeQuantToPagedCache]: kernel does not support MLU300 devices." << std::endl;
    return KernelStatus::KERNEL_STATUS_FAILED;
  }

  // nram_size_need: token_block * head_num * head_size * 2 +
  //                 token_block * head_num * 4 * sizeof(int) + head_num * sizeof(int)
  // nram uesd: 480KB
  int nram_size = 480 * 1024 - num_heads * sizeof(int);
  int hidden_bytes = num_heads * head_size * 2 * sizeof(float) + 4 * num_heads * sizeof(int);
  int seq_block = nram_size / hidden_bytes;
  if (seq_block <= 0) {
    std::cerr << "invokeQuantToPagedCache: "
              << "num_heads * head_size * dtype_size + 4 * num_heads * sizeof(int) "
              << "should be less than 480KB.\n";
    return KernelStatus::KERNEL_STATUS_FAILED;
  }
  if (seq_block > 16) {
    seq_block = seq_block / 16 * 16;
  }
  int cluster_num, core_dim;
  CNdev dev;
  cnCtxGetDevice(&dev);
  CNRT_CHECK(cnrtDeviceGetAttribute(&cluster_num, cnrtAttrClusterCount, dev));
  CNRT_CHECK(cnrtDeviceGetAttribute(&core_dim, cnrtAttrMcorePerCluster, dev));
  int core_num = core_dim * cluster_num;
  seq_block = std::min(seq_block, CEIL_DIV(num_tokens, core_num));
  uint32_t task_dim = CEIL_DIV(num_tokens, seq_block);
  cnrtDim3_t dim{1, task_dim, 1};

  if (data_type == CNNL_DTYPE_FLOAT) {
    kernels::MLUQuantToPagedCacheKernel<<<dim, cnrtFuncTypeBlock, queue>>>(
        (float *)key, (float *)value, (int8_t *)key_cache, (int8_t *)value_cache,
        (float *)key_cache_scale, (float *)value_cache_scale, (int *)slot_mapping, key_stride0,
        value_stride0, num_tokens, num_heads, block_size, head_size, seq_block);
  } else if (data_type == CNNL_DTYPE_HALF) {
    kernels::MLUQuantToPagedCacheKernel<<<dim, cnrtFuncTypeBlock, queue>>>(
        (half *)key, (half *)value, (int8_t *)key_cache, (int8_t *)value_cache,
        (float *)key_cache_scale, (float *)value_cache_scale, (int *)slot_mapping, key_stride0,
        value_stride0, num_tokens, num_heads, block_size, head_size, seq_block);
  } else {
    kernels::MLUQuantToPagedCacheKernel<<<dim, cnrtFuncTypeBlock, queue>>>(
        (bfloat16_t *)key, (bfloat16_t *)value, (int8_t *)key_cache, (int8_t *)value_cache,
        (float *)key_cache_scale, (float *)value_cache_scale, (int *)slot_mapping, key_stride0,
        value_stride0, num_tokens, num_heads, block_size, head_size, seq_block);
  }
  return KernelStatus::KERNEL_STATUS_SUCCESS;
}
}  // namespace tmo
