/*************************************************************************
 * 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 <cassert>
#include <iostream>
#include "offline_quant_to_linear_cache.mluh"
// clang-format off
#include <mlu.h>
// clang-format on

namespace tmo {

namespace kernels {

#define NRAM_BUFFER_SIZE (480 * 1024)
__nram__ int8_t nram_buffer[NRAM_BUFFER_SIZE];

#define sizeof_(T) (uint32_t)sizeof(T)

template <typename T>
__mlu_func__ void quantify(int8_t *nram_output,
                           float *nram_input_float,
                           T *nram_input,
                           float *nram_scale,
                           int input_num,
                           int scale_num) {
  if (std::is_same<half, T>::value) {
    __bang_half2float(nram_input_float, (half *)nram_input, input_num);
  } else if (std::is_same<bfloat16_t, T>::value) {
#if __BANG_ARCH__ > 500
    __bang_bfloat162float(nram_input_float, (bfloat16_t *)nram_input, input_num);
#endif
  }
  __bang_cycle_mul(nram_input_float, nram_input_float, nram_scale, input_num, scale_num);
  __bang_float2int8_rn(nram_output, (float *)nram_input_float, input_num, 0);
}

template <typename T>
__mlu_func__ void quantPerHead(int8_t *output_gdram,
                               int8_t *output_nram,
                               const T *input_gdram,
                               T *input_nram,
                               const float *scale_gdram,
                               float *scale_nram,
                               float *input_nram_float,
                               T *trans_nram,
                               int seq,
                               int head_num,
                               int head_size,
                               size_t in_hstr_bytes,     // context head_num stide bytes
                               size_t in_sstr_bytes,     // context seq stide bytes
                               size_t scale_hstr_bytes,  // scale head_num stide bytes
                               size_t out_hstr_bytes,    // cache head_num stride bytes
                               size_t out_sstr_bytes     // cache seq stride bytes
) {
  constexpr int dtype_size = sizeof_(T);
  // nram_input: (head_num, seq, head_size)
  int io1_size = head_size * dtype_size;
  __memcpy(trans_nram, input_gdram, io1_size, GDRAM2NRAM, seq * io1_size, head_num - 1, io1_size,
           seq - 1, in_hstr_bytes, head_num - 1, in_sstr_bytes, seq - 1);

  // nram_scale:(head_num, seq);
  int io2_size = seq * sizeof_(float);
  __memcpy(scale_nram, scale_gdram, io2_size, GDRAM2NRAM, io2_size, scale_hstr_bytes, head_num - 1);
  __bang_recip(scale_nram, scale_nram, head_num * seq);

  if (std::is_same<T, half>::value || std::is_same<T, bfloat16_t>::value) {
    __bang_transpose((half *)input_nram, (half *)trans_nram, head_num * seq, head_size);
  } else {
    __bang_transpose(input_nram_float, (float *)trans_nram, head_num * seq, head_size);
  }
  quantify<T>(output_nram, input_nram_float, input_nram, scale_nram, head_size * head_num * seq,
              head_num * seq);
  __bang_transpose((int8_t *)trans_nram, output_nram, head_size, head_num * seq);

  __memcpy(output_gdram, trans_nram, head_size, NRAM2GDRAM, out_hstr_bytes, head_num - 1,
           out_sstr_bytes, seq - 1, seq * head_size, head_num - 1, head_size, seq - 1);
}

template <typename T>
__mlu_global__ void MLUOfflineQuantToLinearCacheKernelPerHead(int8_t *key_cache,
                                                              int8_t *value_cache,
                                                              const float *key_cache_scale,
                                                              const float *value_cache_scale,
                                                              const int *cache_bs_offsets,
                                                              const int *cache_seq_offsets,
                                                              const T *key,
                                                              const T *value,
                                                              const int *context_seq_offsets,
                                                              const int *context_lens,
                                                              const int batch,
                                                              const int head_num,
                                                              const int head_size,
                                                              const int max_context_len,
                                                              const int cache_mem_len,
                                                              const size_t context_bs_stride,
                                                              const size_t context_head_stride,
                                                              const size_t context_seq_stride,
                                                              const size_t cache_bs_stride,
                                                              const size_t cache_head_stride,
                                                              const size_t cache_seq_stride,
                                                              const size_t cache_scale_head_stride,
                                                              const bool packed,
                                                              const int seq_block) {
  bool handle_key = (key != nullptr && key_cache != nullptr && key_cache_scale != nullptr);
  bool handle_value = (value != nullptr && value_cache != nullptr && value_cache_scale != nullptr);
  if ((!handle_key) && (!handle_value)) {
    return;
  }

  constexpr int dtype_size = sizeof_(T);
  size_t in_hstr_bytes = context_head_stride * dtype_size;
  size_t in_sstr_bytes = context_seq_stride * dtype_size;
  size_t out_hstr_bytes = cache_head_stride * sizeof_(int8_t);
  size_t out_sstr_bytes = cache_seq_stride * sizeof_(int8_t);
  size_t scale_hstr_bytes = cache_scale_head_stride * sizeof_(float);

  /* ***************************nram space ****************************
   * |      scale    |      input/output     |     trans     |
   * scale size:[head_num, seq_block], float
   * input size:[head_num, seq_block, head_size], float
   * trans size:, [head_size, head_num, seq_block], T
   */
  float *scale_nram = (float *)nram_buffer;
  float *input_nram_float = nullptr;
  T *trans_nram = nullptr, *input_nram = nullptr;
  input_nram_float = scale_nram + head_num * seq_block;
  trans_nram = (T *)(input_nram_float + head_num * seq_block * head_size);

  if (std::is_same<T, half>::value || std::is_same<T, bfloat16_t>::value) {
    // need cast from input_nram to input_nram_float
    input_nram = (T *)input_nram_float + seq_block * head_num * head_size;
  } else {
    input_nram = (T *)input_nram_float;
  }
  int8_t *output_nram = (int8_t *)input_nram_float;  // output and input share nram space

  for (int bs_idx = taskIdY; bs_idx < batch; bs_idx += taskDimY) {
    int context_len = __load_gdram(context_lens + bs_idx);
    int seq_len = packed ? (__load_gdram(context_lens + bs_idx + 1) - context_len) : context_len;
    int task_seq_begin = taskIdZ * seq_block;
    if (task_seq_begin >= seq_len) continue;
    int seq = std::min(seq_len - task_seq_begin, seq_block);

    // context offset
    size_t context_offset = 0;
    if (packed) {
      context_offset = (context_len + task_seq_begin) * context_seq_stride;
    } else {
      int seq_offset = (context_seq_offsets == nullptr) ? 0 : context_seq_offsets[bs_idx];
      context_offset =
          (bs_idx * context_bs_stride + (seq_offset + task_seq_begin) * context_seq_stride);
    }

    // cache offset
    int cache_seq_offset = (cache_seq_offsets == nullptr ? 0 : cache_seq_offsets[bs_idx]);
    int cache_bs_offset = cache_bs_offsets == nullptr ? bs_idx : cache_bs_offsets[bs_idx];
    if (cache_seq_offset < 0 || cache_bs_offset < 0) {
      continue;
    }
    cache_seq_offset += task_seq_begin;
    size_t cache_offset = (cache_bs_offset * cache_bs_stride + cache_seq_offset * cache_seq_stride);
    // per_head, nram input[head_num, seq, head_size], nram scale[head_num, seq]
    if (handle_key) {
      quantPerHead(key_cache + cache_offset, output_nram, key + context_offset, input_nram,
                   key_cache_scale + cache_seq_offset, scale_nram, input_nram_float, trans_nram,
                   seq, head_num, head_size, in_hstr_bytes, in_sstr_bytes, scale_hstr_bytes,
                   out_hstr_bytes, out_sstr_bytes);
    }

    if (handle_value) {
      quantPerHead(value_cache + cache_offset, output_nram, value + context_offset, input_nram,
                   value_cache_scale + cache_seq_offset, scale_nram, input_nram_float, trans_nram,
                   seq, head_num, head_size, in_hstr_bytes, in_sstr_bytes, scale_hstr_bytes,
                   out_hstr_bytes, out_sstr_bytes);
    }
  }
}

template <typename T>
__mlu_global__ void MLUOfflineQuantToLinearCacheKernelPerChannel(
    int8_t *key_cache,
    int8_t *value_cache,
    const float *key_cache_scale,
    const float *value_cache_scale,
    const int *cache_bs_offsets,
    const int *cache_seq_offsets,
    const T *key,
    const T *value,
    const int *context_seq_offsets,
    const int *context_lens,
    const int batch,
    const int head_num,
    const int head_size,
    const int max_context_len,
    const int cache_mem_len,
    const size_t context_bs_stride,
    const size_t context_head_stride,
    const size_t context_seq_stride,
    const size_t cache_bs_stride,
    const size_t cache_head_stride,
    const size_t cache_seq_stride,
    const size_t cache_scale_head_stride,
    const bool packed,
    const int seq_block) {
  bool handle_key = (key != nullptr && key_cache != nullptr && key_cache_scale != nullptr);
  bool handle_value = (value != nullptr && value_cache != nullptr && value_cache_scale != nullptr);
  if ((!handle_key) && (!handle_value)) {
    return;
  }

  constexpr int dtype_size = sizeof_(T);
  size_t in_hstr_bytes = context_head_stride * dtype_size;
  size_t in_sstr_bytes = context_seq_stride * dtype_size;
  size_t out_hstr_bytes = cache_head_stride * sizeof_(int8_t);
  size_t out_sstr_bytes = cache_seq_stride * sizeof_(int8_t);
  size_t scale_hstr_bytes = cache_scale_head_stride * sizeof_(float);

  /* *********************************nram space **************************************
   * per_chennel: |scale[head_num, head_size]| input[seq_block, head_num, head_size]|
   */
  float *scale_nram = (float *)nram_buffer;
  float *input_nram_float = scale_nram + head_num * head_size;

  T *input_nram = (T *)input_nram_float;
  if (std::is_same<T, half>::value || std::is_same<T, bfloat16_t>::value) {
    // need cast from input_nram to input_nram_float
    input_nram = (T *)input_nram_float + seq_block * head_num * head_size;
  }
  int8_t *output_nram = (int8_t *)input_nram_float;  // output and input share nram space
  int size1 = head_size * sizeof_(float);
  int size2 = head_size * dtype_size;
  int scale_num = head_num * head_size;

  if (handle_key) {
    // load offline scale nram_scale:(head_num, head_size);
    __memcpy(scale_nram, key_cache_scale, size1, GDRAM2NRAM, size1, scale_hstr_bytes, head_num - 1);
    __bang_recip(scale_nram, scale_nram, scale_num);

    for (int bs_idx = taskIdY; bs_idx < batch; bs_idx += taskDimY) {
      int context_len = __load_gdram(context_lens + bs_idx);
      int seq_len = packed ? (__load_gdram(context_lens + bs_idx + 1) - context_len) : context_len;
      int task_seq_begin = taskIdZ * seq_block;
      if (task_seq_begin >= seq_len) continue;
      int seq = std::min(seq_len - task_seq_begin, seq_block);

      // context offset
      size_t context_offset = 0;
      if (packed) {
        context_offset = (context_len + task_seq_begin) * context_seq_stride;
      } else {
        int seq_offset = (context_seq_offsets == nullptr) ? 0 : context_seq_offsets[bs_idx];
        context_offset =
            (bs_idx * context_bs_stride + (seq_offset + task_seq_begin) * context_seq_stride);
      }

      // cache offset
      int cache_seq_offset = (cache_seq_offsets == nullptr ? 0 : cache_seq_offsets[bs_idx]);
      int cache_bs_offset = cache_bs_offsets == nullptr ? bs_idx : cache_bs_offsets[bs_idx];
      if (cache_seq_offset < 0 || cache_bs_offset < 0) {
        continue;
      }
      cache_seq_offset += task_seq_begin;
      size_t cache_offset =
          (cache_bs_offset * cache_bs_stride + cache_seq_offset * cache_seq_stride);

      __memcpy(input_nram, key + context_offset, size2, GDRAM2NRAM, size2, head_num - 1,
               head_num * size2, seq - 1, in_hstr_bytes, head_num - 1, in_sstr_bytes, seq - 1);

      quantify<T>((int8_t *)output_nram, input_nram_float, input_nram, scale_nram, seq * scale_num,
                  scale_num);

      __memcpy(key_cache + cache_offset, output_nram, head_size, NRAM2GDRAM, out_hstr_bytes,
               head_num - 1, out_sstr_bytes, seq - 1, head_size, head_num - 1, scale_num, seq - 1);
    }
  }

  if (handle_value) {
    // load offline scale nram_scale:(head_num, head_size);
    __memcpy(scale_nram, value_cache_scale, size1, GDRAM2NRAM, size1, scale_hstr_bytes,
             head_num - 1);
    __bang_recip(scale_nram, scale_nram, scale_num);

    for (int bs_idx = taskIdY; bs_idx < batch; bs_idx += taskDimY) {
      int context_len = __load_gdram(context_lens + bs_idx);
      int seq_len = packed ? (__load_gdram(context_lens + bs_idx + 1) - context_len) : context_len;
      int task_seq_begin = taskIdZ * seq_block;
      if (task_seq_begin >= seq_len) continue;
      int seq = std::min(seq_len - task_seq_begin, seq_block);

      // context offset
      size_t context_offset = 0;
      if (packed) {
        context_offset = (context_len + task_seq_begin) * context_seq_stride;
      } else {
        int seq_offset = (context_seq_offsets == nullptr) ? 0 : context_seq_offsets[bs_idx];
        context_offset =
            (bs_idx * context_bs_stride + (seq_offset + task_seq_begin) * context_seq_stride);
      }

      // cache offset
      int cache_seq_offset = (cache_seq_offsets == nullptr ? 0 : cache_seq_offsets[bs_idx]);
      int cache_bs_offset = cache_bs_offsets == nullptr ? bs_idx : cache_bs_offsets[bs_idx];
      if (cache_seq_offset < 0 || cache_bs_offset < 0) {
        continue;
      }
      cache_seq_offset += task_seq_begin;
      size_t cache_offset =
          (cache_bs_offset * cache_bs_stride + cache_seq_offset * cache_seq_stride);

      __memcpy(input_nram, value + context_offset, size2, GDRAM2NRAM, size2, head_num - 1,
               head_num * size2, seq - 1, in_hstr_bytes, head_num - 1, in_sstr_bytes, seq - 1);

      quantify<T>((int8_t *)output_nram, input_nram_float, input_nram, scale_nram, seq * scale_num,
                  scale_num);

      __memcpy(value_cache + cache_offset, output_nram, head_size, NRAM2GDRAM, out_hstr_bytes,
               head_num - 1, out_sstr_bytes, seq - 1, head_size, head_num - 1, scale_num, seq - 1);
    }
  }
}
}  // namespace kernels

#define LAUNCH_OFFLINE_QUANT_KERNEL(Dtype, Name)                                                   \
  kernels::MLUOfflineQuantToLinearCacheKernel##Name<Dtype><<<dim, cnrtFuncTypeBlock, queue>>>(     \
      (int8_t *)key_cache, (int8_t *)value_cache, (float *)key_cache_scale,                        \
      (float *)value_cache_scale, (int *)cache_bs_offsets, (int *)cache_seq_offsets, (Dtype *)key, \
      (Dtype *)value, (int *)context_seq_offsets, (int *)context_lens, batch, head_num, head_size, \
      max_context_len, cache_mem_len, context_bs_stride, context_head_stride, context_seq_stride,  \
      cache_bs_stride, cache_head_stride, cache_seq_stride, cache_scale_head_stride, packed,       \
      seq_block);

KernelStatus invokeOfflineQuantToLinearCache(cnrtQueue_t queue,
                                             void *key_cache,
                                             void *value_cache,
                                             const void *key_cache_scale,
                                             const void *value_cache_scale,
                                             const void *cache_bs_offsets,
                                             const void *cache_seq_offsets,
                                             const void *key,
                                             const void *value,
                                             const void *context_seq_offsets,
                                             const void *context_lens,
                                             const cnnlDataType_t dtype,
                                             const int batch,
                                             const int head_num,
                                             const int head_size,
                                             const int max_context_len,
                                             const int cache_mem_len,
                                             const size_t context_bs_stride,
                                             const size_t context_head_stride,
                                             const size_t context_seq_stride,
                                             const size_t cache_bs_stride,
                                             const size_t cache_head_stride,
                                             const size_t cache_seq_stride,
                                             const size_t cache_scale_head_stride,
                                             const bool packed,
                                             const int quant_mode) {
  constexpr int nram_size = 480 * 1024;
  int dtype_size = dtype == CNNL_DTYPE_FLOAT ? sizeof(float) : sizeof(half);
  int seq_block = 0;
  if (quant_mode == 0) {
    seq_block = nram_size / (head_num * head_size * sizeof(float)) - 1;
    if (seq_block <= 0) {
      std::cerr << __func__ << "," << __LINE__
                << " :head_num * head_size * sizeof(float) should be less than 240KB when "
                   "quant_mode is 0."
                << std::endl;
    }
  } else {
    seq_block = nram_size /
                (head_num * sizeof(float) + head_num * head_size * (sizeof(float) + dtype_size));
    if (seq_block <= 0) {
      std::cerr << __func__ << "," << __LINE__
                << " :head_num * sizeof(float) + head_num * head_size * (sizeof(float) + "
                   "context_dtype_size)) "
                << " should be less than 480KB when quant_mode is not 0." << std::endl;
    }
  }
  seq_block = std::min(seq_block, max_context_len);
  if (seq_block > 16 && seq_block < max_context_len) {
    seq_block = seq_block / 16 * 16;
  }
  int seq_seg = (max_context_len + seq_block - 1) / seq_block;

  CNdev dev;
  int cluster_dim, core_dim;
  cnCtxGetDevice(&dev);
  CNRT_CHECK(cnrtDeviceGetAttribute(&cluster_dim, cnrtAttrClusterCount, dev));
  CNRT_CHECK(cnrtDeviceGetAttribute(&core_dim, cnrtAttrMcorePerCluster, dev));

  uint32_t core_num = cluster_dim * core_dim;
  uint32_t task_y_dim = std::min((uint32_t)batch, core_num);
  cnrtDim3_t dim{1, task_y_dim, (uint32_t)seq_seg};

  if (dtype == CNNL_DTYPE_HALF) {
    if (quant_mode == 0) {
      LAUNCH_OFFLINE_QUANT_KERNEL(half, PerChannel);
    } else {
      LAUNCH_OFFLINE_QUANT_KERNEL(half, PerHead);
    }
  } else if (dtype == CNNL_DTYPE_BFLOAT16) {
    if (quant_mode == 0) {
      LAUNCH_OFFLINE_QUANT_KERNEL(bfloat16_t, PerChannel);
    } else {
      LAUNCH_OFFLINE_QUANT_KERNEL(bfloat16_t, PerHead);
    }
  } else {
    if (quant_mode == 0) {
      LAUNCH_OFFLINE_QUANT_KERNEL(float, PerChannel);
    } else {
      LAUNCH_OFFLINE_QUANT_KERNEL(float, PerHead);
    }
  }
  return KernelStatus::KERNEL_STATUS_SUCCESS;
}
}  // namespace tmo
