/*
 * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/common_shape.h"

namespace phi {
template <typename T>
inline T xabs(const T x) {
  return x < static_cast<T>(0.0) ? -x : x;
}

template <typename T, typename ScaleT>
void per_channel_scale(
    ScaleT* scale, const T* input, size_t m, size_t n, float bound) {
  for (size_t i = 0; i < n; ++i) {
    float max = static_cast<float>(input[i]);
    for (size_t j = 0; j < m; ++j) {
      max = static_cast<float>(xabs(input[j * n + i])) > max
                ? static_cast<float>(xabs(input[j * n + i]))
                : max;
    }
    scale[i] = static_cast<ScaleT>(max / bound);
  }
}

template <typename T, typename ScaleT>
void group_wise_scale(ScaleT* scale,
                      const T* input,
                      size_t m,
                      size_t n,
                      float bound,
                      size_t group_size) {
  for (size_t i = 0; i < n; ++i) {
    for (size_t j = 0; j < m; j += group_size) {
      float max = static_cast<float>(0.f);
      for (size_t k = 0; k < group_size && j + k < m; ++k) {
        max = static_cast<float>(xabs(input[(j + k) * n + i])) > max
                  ? static_cast<float>(xabs(input[(j + k) * n + i]))
                  : max;
      }
      scale[static_cast<int>(j / group_size) * n + i] =
          static_cast<ScaleT>(max / bound);
    }
  }
}

template <typename T, int quant_bit = 8, typename ScaleT>
void per_channel_quant(int8_t* output,
                       const T* input,
                       const ScaleT* scale,
                       size_t num_rows,
                       size_t num_cols) {
  size_t bytes_per_out_col = num_cols * quant_bit / 8;
  for (size_t ii = 0; ii < num_rows; ++ii) {
    int8_t* current_quantized_weight_row = output + ii * bytes_per_out_col;
    const T* current_weight_row = input + ii * num_cols;
    for (size_t jj = 0; jj < bytes_per_out_col; ++jj) {
      if (quant_bit == 8) {
        const float col_scale = static_cast<float>(scale[jj]);
        const float weight_elt = static_cast<float>(current_weight_row[jj]);
        const float scaled_weight = round(weight_elt / col_scale);
        const int8_t clipped_weight = static_cast<int8_t>(
            std::max(-127.f, std::min(127.f, scaled_weight)));
        current_quantized_weight_row[jj] = clipped_weight;
      } else if (quant_bit == 4) {
        // We will pack two int4 elements per iteration of the inner loop.
        int8_t packed_int4s = 0;
        for (int packed_idx = 0; packed_idx < 2; ++packed_idx) {
          const size_t input_idx = 2 * jj + packed_idx;
          if (input_idx < num_cols) {
            const float col_scale = static_cast<float>(scale[input_idx]);
            const float weight_elt =
                static_cast<float>(current_weight_row[input_idx]);
            const float scaled_weight = round(weight_elt / col_scale);
            int int_weight = static_cast<int>(scaled_weight);
#ifdef PADDLE_WITH_HIP
            const int8_t clipped_weight =
                std::max(-7, std::min(7, int_weight)) + 8;
#else
            const int8_t clipped_weight = std::max(-7, std::min(7, int_weight));
#endif

            // Kill the sign extension bits (hence 0x0F mask) then shift to
            // upper bits if packing the second int4 and or the bits into the
            // final result.
            packed_int4s |= ((clipped_weight & 0x0F) << (4 * packed_idx));
          }
        }
        current_quantized_weight_row[jj] = packed_int4s;
      } else {
        common::errors::Unimplemented("Unsupported quantization bits: %d",
                                      quant_bit);
      }
    }
  }
}

template <typename T, int quant_bit = 8, typename ScaleT>
void group_wise_quant(int8_t* output,
                      const T* input,
                      const ScaleT* scale,
                      size_t num_rows,
                      size_t num_cols,
                      const int group_size) {
  size_t bytes_per_out_col = num_cols * quant_bit / 8;
  for (size_t ii = 0; ii < num_rows; ++ii) {
    int8_t* current_quantized_weight_row = output + ii * bytes_per_out_col;
    const T* current_weight_row = input + ii * num_cols;
    for (size_t jj = 0; jj < bytes_per_out_col; ++jj) {
      if (quant_bit == 8) {
        size_t scale_cur_offset = jj + (ii / group_size) * num_cols;
        const float col_scale = static_cast<float>(scale[scale_cur_offset]);
        const float weight_elt = static_cast<float>(current_weight_row[jj]);
        const float scaled_weight = round(weight_elt / col_scale);
        const int8_t clipped_weight = static_cast<int8_t>(
            std::max(-127.f, std::min(127.f, scaled_weight)));
        current_quantized_weight_row[jj] = clipped_weight;
      } else if (quant_bit == 4) {
        // We will pack two int4 elements per iteration of the inner loop.
        int8_t packed_int4s = 0;
        for (int packed_idx = 0; packed_idx < 2; ++packed_idx) {
          const size_t input_idx = 2 * jj + packed_idx;
          if (input_idx < num_cols) {
            size_t scale_cur_offset = input_idx + (ii / group_size) * num_cols;
            const float col_scale = static_cast<float>(scale[scale_cur_offset]);
            const float weight_elt =
                static_cast<float>(current_weight_row[input_idx]);
            const float scaled_weight = round(weight_elt / col_scale);
            int int_weight = static_cast<int>(scaled_weight);
#ifdef PADDLE_WITH_HIP
            const int8_t clipped_weight =
                std::max(-7, std::min(7, int_weight)) + 8;
#else
            const int8_t clipped_weight = std::max(-7, std::min(7, int_weight));
#endif

            // Kill the sign extension bits (hence 0x0F mask) then shift to
            // upper bits if packing the second int4 and or the bits into the
            // final result.
            packed_int4s |= ((clipped_weight & 0x0F) << (4 * packed_idx));
          }
        }
        current_quantized_weight_row[jj] = packed_int4s;
      } else {
        common::errors::Unimplemented("Unsupported quantization bits: %d",
                                      quant_bit);
      }
    }
  }
}

template <int quant_bit = 8>
void add_bias_and_interleave_inplace(int8_t* tensor_ptr, size_t num_elts) {
  const size_t num_bytes = num_elts * quant_bit / 8;

  for (size_t ii = 0; ii < num_bytes; ++ii) {
    if (quant_bit == 8) {
      tensor_ptr[ii] =
          static_cast<int8_t>(static_cast<int>(tensor_ptr[ii]) + 128);
    } else {
      int8_t transformed_packed_int4s = 0;
      int8_t transformed_first_elt =
          (int8_t(tensor_ptr[ii] << 4) >> 4) +
          8;  // The double shift here is to ensure sign extension
      int8_t transformed_second_elt = (tensor_ptr[ii] >> 4) + 8;

      if (!(transformed_first_elt >= 0 && transformed_first_elt <= 15)) {
        common::errors::InvalidArgument(
            "Illegal result for int4 transform (first elt)");
      }
      if (!(transformed_second_elt >= 0 && transformed_second_elt <= 15)) {
        common::errors::InvalidArgument(
            "Illegal result for int4 transform (second elt)");
      }
      // We don't need to mask in these ops since everything should be in the
      // range 0-15
      transformed_packed_int4s |= transformed_first_elt;
      transformed_packed_int4s |= (transformed_second_elt << 4);
      tensor_ptr[ii] = transformed_packed_int4s;
    }
  }
  if (quant_bit == 8) {
    for (size_t base = 0; base < num_elts; base += 4) {
      std::swap(tensor_ptr[base + 1], tensor_ptr[base + 2]);
    }
  } else {
    const size_t num_registers = num_bytes / 4;

    uint32_t* register_ptr = reinterpret_cast<uint32_t*>(tensor_ptr);
    for (size_t ii = 0; ii < num_registers; ++ii) {
      const uint32_t current_register = register_ptr[ii];
      uint32_t transformed_register = 0;

      for (int dest_idx = 0; dest_idx < 8; ++dest_idx) {
        const int src_idx =
            dest_idx < 4 ? 2 * dest_idx : 2 * (dest_idx - 4) + 1;
        const int src_shift = 4 * src_idx;
        const int dest_shift = 4 * dest_idx;

        const uint32_t src_bits = (current_register >> src_shift) & 0xF;
        transformed_register |= (src_bits << dest_shift);
      }
      register_ptr[ii] = transformed_register;
    }
  }
}

template <int quant_bit>
void permute_B_rows_for_mixed_gemm(int8_t* permuted_quantized_tensor,
                                   const int8_t* quantized_tensor,
                                   const std::vector<size_t>& shape) {
  // We only want to run this step for weight only quant.
  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];

  const int BITS_PER_ELT = quant_bit;
  const int K = 16 / BITS_PER_ELT;
  const int ELTS_PER_REG = 32 / BITS_PER_ELT;

  const uint32_t* input_byte_ptr =
      reinterpret_cast<const uint32_t*>(quantized_tensor);
  uint32_t* output_byte_ptr =
      reinterpret_cast<uint32_t*>(permuted_quantized_tensor);

  int B_ROWS_PER_MMA = 8 * K;
  const int elts_in_int32 = 32 / BITS_PER_ELT;

  const int num_vec_cols = num_cols / elts_in_int32;

  // The code is written as below so it works for both int8 and packed int4.
  for (size_t base_row = 0; base_row < num_rows; base_row += B_ROWS_PER_MMA) {
    for (int tile_row = 0; tile_row < B_ROWS_PER_MMA; ++tile_row) {
      for (int write_col = 0; write_col < num_vec_cols; ++write_col) {
        const int write_row = base_row + tile_row;
        const int tile_read_row = 8 * (((tile_row % ELTS_PER_REG) / 2)) +
                                  tile_row % 2 + 2 * (tile_row / ELTS_PER_REG);
        const int read_row = base_row + tile_read_row;
        const int read_col = write_col;

        const int64_t read_offset =
            static_cast<int64_t>(read_row) * num_vec_cols + read_col;
        const int64_t write_offset =
            static_cast<int64_t>(write_row) * num_vec_cols + write_col;
        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
      }
    }
  }
}

template <int quant_bit>
void subbyte_transpose_impl(int8_t* transposed_quantized_tensor,
                            const int8_t* quantized_tensor,
                            const std::vector<size_t>& shape) {
  const int bits_per_elt = quant_bit;

  // FT_CHECK_WITH_INFO(shape.size() == 2 || shape.size() == 3, "Shape must be
  // 2-D or 3-D");
  // const size_t num_experts = 1;
  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];

  const size_t col_bytes = num_cols * bits_per_elt / 8;
  const size_t col_bytes_trans = num_rows * bits_per_elt / 8;
  // const size_t num_bytes = size_t(num_experts) * num_rows * col_bytes;

  const uint8_t* input_byte_ptr =
      reinterpret_cast<const uint8_t*>(quantized_tensor);
  uint8_t* output_byte_ptr =
      reinterpret_cast<uint8_t*>(transposed_quantized_tensor);

  static constexpr int ELTS_PER_BYTE = 8 / quant_bit;

  static constexpr int M_TILE_L1 = 64;
  static constexpr int N_TILE_L1 = M_TILE_L1 / ELTS_PER_BYTE;
  uint8_t cache_buf[M_TILE_L1][N_TILE_L1];

  static constexpr int VECTOR_WIDTH = std::min(32, N_TILE_L1);

  // const int num_m_tiles = (num_rows + M_TILE_L1 - 1) / M_TILE_L1;
  // const int num_n_tiles = (col_bytes + N_TILE_L1 - 1) / N_TILE_L1;

  for (size_t row_tile_start = 0; row_tile_start < num_rows;
       row_tile_start += M_TILE_L1) {
    for (size_t col_tile_start_byte = 0; col_tile_start_byte < col_bytes;
         col_tile_start_byte += N_TILE_L1) {
      const int row_limit = std::min(row_tile_start + M_TILE_L1, num_rows);
      const int col_limit =
          std::min(col_tile_start_byte + N_TILE_L1, col_bytes);

      for (int ii = 0; ii < M_TILE_L1; ++ii) {
        const int row = row_tile_start + ii;

        for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
          const int col = col_tile_start_byte + jj;

          const size_t logical_src_offset = row * col_bytes + col;

          if (row < row_limit && col < col_limit) {
            for (int v = 0; v < VECTOR_WIDTH; ++v) {
              cache_buf[ii][jj + v] = input_byte_ptr[logical_src_offset + v];
            }
          }
        }
      }

      if (quant_bit == 8) {
        for (int ii = 0; ii < M_TILE_L1; ++ii) {
          for (int jj = ii + 1; jj < N_TILE_L1; ++jj) {
            std::swap(cache_buf[ii][jj], cache_buf[jj][ii]);
          }
        }
      } else if (quant_bit == 4) {
        for (int ii = 0; ii < M_TILE_L1; ++ii) {
          // Using M_TILE_L1 here is deliberate since we assume that the cache
          // tile is square in the number of elements (not necessarily the
          // number of bytes).
          for (int jj = ii + 1; jj < M_TILE_L1; ++jj) {
            const int ii_byte = ii / ELTS_PER_BYTE;
            const int ii_bit_offset = ii % ELTS_PER_BYTE;

            const int jj_byte = jj / ELTS_PER_BYTE;
            const int jj_bit_offset = jj % ELTS_PER_BYTE;

            uint8_t src_elt =
                0xF & (cache_buf[ii][jj_byte] >> (4 * jj_bit_offset));
            uint8_t tgt_elt =
                0xF & (cache_buf[jj][ii_byte] >> (4 * ii_bit_offset));

            cache_buf[ii][jj_byte] &= (0xF0 >> (4 * jj_bit_offset));
            cache_buf[jj][ii_byte] &= (0xF0 >> (4 * ii_bit_offset));

            cache_buf[ii][jj_byte] |= (tgt_elt << (4 * jj_bit_offset));
            cache_buf[jj][ii_byte] |= (src_elt << (4 * ii_bit_offset));
          }
        }
      } else {
        common::errors::Unimplemented("Unsupported quantization bits: %d",
                                      quant_bit);
      }

      const size_t row_tile_start_trans = col_tile_start_byte * ELTS_PER_BYTE;
      const size_t col_tile_start_byte_trans = row_tile_start / ELTS_PER_BYTE;

      const int row_limit_trans =
          std::min(row_tile_start_trans + M_TILE_L1, num_cols);
      const int col_limit_trans =
          std::min(col_tile_start_byte_trans + N_TILE_L1, col_bytes_trans);

      for (int ii = 0; ii < M_TILE_L1; ++ii) {
        const int row = row_tile_start_trans + ii;
        for (int jj = 0; jj < N_TILE_L1; jj += VECTOR_WIDTH) {
          const int col = col_tile_start_byte_trans + jj;

          const size_t logical_tgt_offset = row * col_bytes_trans + col;

          if (row < row_limit_trans && col < col_limit_trans) {
            for (int v = 0; v < VECTOR_WIDTH; ++v) {
              output_byte_ptr[logical_tgt_offset + v] = cache_buf[ii][jj + v];
            }
          }
        }
      }
    }
  }
}

template <int quant_bit>
void interleave_column_major_tensor(int8_t* interleaved_quantized_tensor,
                                    const int8_t* quantized_tensor,
                                    const std::vector<size_t>& shape) {
  // We only want to run this step for weight only quant.
  const size_t num_rows = shape.size() == 2 ? shape[0] : shape[1];
  const size_t num_cols = shape.size() == 2 ? shape[1] : shape[2];

  const size_t BITS_PER_ELT = quant_bit;
  const size_t elts_in_int32 = 32 / BITS_PER_ELT;

  const size_t rows_per_tile = 64;

  const uint32_t* input_byte_ptr =
      reinterpret_cast<const uint32_t*>(quantized_tensor);
  uint32_t* output_byte_ptr =
      reinterpret_cast<uint32_t*>(interleaved_quantized_tensor);

  const size_t num_vec_rows = num_rows / elts_in_int32;
  const size_t vec_rows_per_tile = rows_per_tile / elts_in_int32;
  const size_t interleave = 128 * 8 / quant_bit / rows_per_tile;
  for (size_t read_col = 0; read_col < num_cols; ++read_col) {
    const size_t write_col = read_col / interleave;
    for (size_t base_vec_row = 0; base_vec_row < num_vec_rows;
         base_vec_row += vec_rows_per_tile) {
      for (size_t vec_read_row = base_vec_row;
           vec_read_row <
           std::min(num_vec_rows, base_vec_row + vec_rows_per_tile);
           ++vec_read_row) {
        const size_t vec_write_row =
            interleave * base_vec_row +
            vec_rows_per_tile * (read_col % interleave) +
            vec_read_row % vec_rows_per_tile;
        const size_t read_offset =
            size_t(read_col) * num_vec_rows + vec_read_row;
        const size_t write_offset =
            size_t(write_col) * num_vec_rows * interleave + vec_write_row;
        output_byte_ptr[write_offset] = input_byte_ptr[read_offset];
      }
    }
  }
}
}  // namespace phi