// Copyright © 2025 Apple Inc.

#include "mlx/backend/common/quantized.h"
#include "mlx/backend/cuda/device.h"
#include "mlx/backend/cuda/kernel_utils.cuh"
#include "mlx/backend/cuda/quantized/mxfp8_quantize.cuh"
#include "mlx/backend/cuda/quantized/nvfp4_quantize.cuh"
#include "mlx/backend/cuda/quantized/quantized.h"
#include "mlx/backend/cuda/vector_types.cuh"
#include "mlx/dtype_utils.h"

#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#include <cutlass/float8.h>
#include <cutlass/numeric_conversion.h>

constexpr float F8E4M3_MAX = 448.0f;
constexpr float F4E2M1_MAX = 6.0f;

namespace mlx::core {
namespace cu {

template <int bits>
struct Dequantize {
  __device__ float operator()(uint8_t x) {
    if constexpr (bits == 8) {
      return float(*(cutlass::float_e4m3_t*)(&x));
    } else {
      return float(*(cutlass::float_e2m1_t*)(&x));
    }
  }
};

template <typename T>
__device__ __forceinline__ void absmax_x2(T& out, const T& x1, const T& x2) {
  if constexpr (
      (std::is_same<T, __nv_bfloat162>::value) ||
      (std::is_same<T, __half2>::value)) {
    T a = x1;
    T b = x2;
    out = __hmax2(__habs2(a), __habs2(b));
  } else if constexpr (std::is_same<T, float2>::value) {
    float2 a = x1;
    float2 b = x2;
    out.x = fmaxf(fabsf(a.x), fabsf(b.x));
    out.y = fmaxf(fabsf(a.y), fabsf(b.y));
  }
}

namespace cg = cooperative_groups;

template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
__global__ void fp_quantize_dequantize(
    T* w,
    T* out,
    size_t size,
    float* global_scale = nullptr) {
  const bool use_global_scale = global_scale != nullptr;
  const float scale_enc =
      use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f;
  const float inv_scale_enc = use_global_scale ? 1.0f / scale_enc : 1.0f;

  using Tx2 = Vector2_t<T>;
  uint32_t rbits = 0; // reserved bits for future use
  auto block_size = cg::this_thread_block().dim_threads();
  auto block_idx = cg::this_thread_block().group_index();
  auto idx_in_block = cg::this_thread_block().thread_index();
  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;

  size_t thread_idx = tidx + grid_dim_x * size_t(tidy);
  size_t base_idx = thread_idx * group_size;

  if (base_idx >= size) {
    return;
  }

  auto w_tile = load_vector<group_size, T>(w, thread_idx);
  float scale_dec_b = 0.0f;

  Tx2 amax_2x = Tx2{0.0f, 0.0f};

#pragma unroll
  for (int i = 0; i < group_size; i += 2) {
    auto pair = Tx2{w_tile[i], w_tile[i + 1]};
    absmax_x2<Tx2>(amax_2x, amax_2x, pair);
  }

  scale_dec_b = static_cast<float>(
      max(fabsf(static_cast<float>(amax_2x.x)),
          fabsf(static_cast<float>(amax_2x.y))));

  scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX;
  scale_dec_b *= scale_enc;
  // Convert to mx scale or nv scale
  using ScaleType = std::conditional_t<
      use_mx_scale,
      cutlass::float_ue8m0_t,
      cutlass::float_e4m3_t>;
  auto s = ScaleType(scale_dec_b);
  float scale_enc_b = scale_enc / float(s);
  float scale_dec = float(s) * inv_scale_enc;
  AlignedVector<T, group_size> w_hat;

#pragma unroll
  for (int i = 0; i < group_size / 8; i++) {
    auto& w = *reinterpret_cast<cutlass::Array<T, 8>*>(&w_tile[i * 8]);
    cutlass::NumericArrayConverter<float, T, 8> fp32_t;
    auto scaled = fp32_t(w) * scale_enc_b;
    cutlass::Array<float, 8> dq;
    if constexpr (bits == 8) {
      cutlass::NumericArrayConverter<cutlass::float_e4m3_t, float, 8> fp8_fp32;
      auto quant = fp8_fp32(scaled);
      cutlass::NumericArrayConverter<float, cutlass::float_e4m3_t, 8> fp32_fp8;
      dq = fp32_fp8(quant);
    } else {
      cutlass::NumericArrayConverter<cutlass::float_e2m1_t, float, 8> fp4_fp32;
      auto quant = fp4_fp32(scaled);
      cutlass::NumericArrayConverter<float, cutlass::float_e2m1_t, 8> fp32_fp4;
      dq = fp32_fp4(quant);
    }
    cutlass::NumericArrayConverter<T, float, 8> t_fp32;
    *reinterpret_cast<cutlass::Array<T, 8>*>(&w_hat[i * 8]) =
        t_fp32(dq * scale_dec);
  }
  store_vector<group_size>(out, thread_idx, w_hat);
}

template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
__global__ void fp_quantize_rowwise(
    T* w,
    uint8_t* out,
    uint8_t* scales,
    size_t size,
    float* global_scale = nullptr) {
  // NVFP4 conversion:
  // Global encode scale: (448 × 6) / *global_scale
  // Per-block decode scale: S_dec_b = (block_amax / 6) × S_enc → stored as FP8
  // E4M3 Per-block encode scale: S_enc_b = S_enc / S_dec_b
  const bool use_global_scale = global_scale != nullptr;
  const float scale_enc =
      use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f;

  using Tx2 = Vector2_t<T>;
  using Tx4 = Vector4_t<T>;
  uint32_t rbits = 0; // reserved bits for future use
  auto block_size = cg::this_thread_block().dim_threads();
  auto block_idx = cg::this_thread_block().group_index();
  auto idx_in_block = cg::this_thread_block().thread_index();
  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
  auto tidy = block_idx.y * block_size.y + idx_in_block.y;
  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;

  size_t thread_idx = tidx + grid_dim_x * size_t(tidy);
  size_t base_idx = thread_idx * group_size;

  if (base_idx >= size) {
    return;
  }

  auto w_tile = load_vector<group_size, T>(w, thread_idx);
  float scale_dec_b = 0.0f;

  Tx2 amax_2x = Tx2{0.0f, 0.0f};

#pragma unroll
  for (int i = 0; i < group_size; i += 2) {
    auto pair = Tx2{w_tile[i], w_tile[i + 1]};
    absmax_x2<Tx2>(amax_2x, amax_2x, pair);
  }

  scale_dec_b = static_cast<float>(
      max(fabsf(static_cast<float>(amax_2x.x)),
          fabsf(static_cast<float>(amax_2x.y))));

  scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX;
  scale_dec_b *= scale_enc;
  // Convert to mx scale or nv scale
  using ScaleType = std::conditional_t<
      use_mx_scale,
      cutlass::float_ue8m0_t,
      cutlass::float_e4m3_t>;
  auto s = ScaleType(scale_dec_b);
  uint8_t q_scale = s.storage;
  float scale_enc_b = scale_enc / float(s);

  scales[thread_idx] = q_scale;
  constexpr int elem_per_byte = bits == 8 ? 1 : 2;
  AlignedVector<uint8_t, group_size / elem_per_byte> quantized;

#pragma unroll
  for (int i = 0; i < group_size / 4; i++) {
    Tx4 w_Tx4 = *reinterpret_cast<Tx4*>(&w_tile[i * 4]);
    if constexpr (bits == 8) {
      uint32_t quantized_val =
          scale_cvt_Tx4_to_fp8x4<T, USE_SR>(w_Tx4, scale_enc_b, rbits);
      *reinterpret_cast<uint32_t*>(&quantized[i * 4]) = quantized_val;
    } else {
      uint16_t quantized_val =
          scale_cvt_Tx4_to_fp4x4<T, USE_SR>(w_Tx4, scale_enc_b, rbits);
      *reinterpret_cast<uint16_t*>(&quantized[i * 2]) = quantized_val;
    }
  }
  store_vector<group_size / elem_per_byte>(out, thread_idx, quantized);
}

template <typename T, int group_size, int bits, bool use_mx_scale, bool USE_SR>
__global__ void fp_quantize_columnwise(
    T* w,
    uint8_t* out,
    uint8_t* scales,
    size_t size,
    int M,
    int K,
    float* global_scale = nullptr) {
  // Input: [M, K] with strides [1, M] (M-major)
  // Quantized output: [M, K/elem_per_byte] row-major (K-major)
  // Scales: [M, K/group_size] row-major (K-major)
  // Quantize along K (last dimension, groups of group_size elements)
  const bool use_global_scale = global_scale != nullptr;
  const float scale_enc =
      use_global_scale ? (F8E4M3_MAX * F4E2M1_MAX) / *global_scale : 1.0f;

  using Tx2 = Vector2_t<T>;
  using Tx4 = Vector4_t<T>;
  uint32_t rbits = 0;

  auto block_idx = cg::this_thread_block().group_index();
  auto idx_in_block = cg::this_thread_block().thread_index();

  constexpr int BLOCK_X = 16;
  constexpr int BLOCK_Y = 32;
  constexpr int elem_per_byte = (bits == 8) ? 1 : 2;
  constexpr int bytes_per_group = group_size / elem_per_byte;

  constexpr int rows_per_block = BLOCK_X;
  constexpr int cols_per_block = BLOCK_Y * group_size;
  constexpr int local_cols = cols_per_block / elem_per_byte;
  constexpr int bytes_per_block = rows_per_block * local_cols;

  constexpr int SMEM_PAD = 4;
  constexpr int padded_local_cols = local_cols + SMEM_PAD;

  auto tidx = idx_in_block.x;
  auto tidy = idx_in_block.y;

  int num_col_blocks = (K + cols_per_block - 1) / cols_per_block;
  auto bidx = block_idx.x % num_col_blocks;
  auto bidy = block_idx.x / num_col_blocks;

  T thread_data[group_size];

  __shared__ uint8_t quantized_smem[rows_per_block * padded_local_cols];
  __shared__ uint8_t scales_smem[BLOCK_X][BLOCK_Y + SMEM_PAD];

  int row_base = bidy * rows_per_block + tidx;
  int col_base = bidx * cols_per_block + tidy * group_size;

  bool valid = (row_base < M) && (col_base + group_size <= K);
  if (valid) {
#pragma unroll
    for (int i = 0; i < group_size; i++) {
      auto index = row_base + (col_base + i) * M;
      thread_data[i] = w[index];
    }

    // Compute scale
    Tx2 amax_2x = Tx2{0.0f, 0.0f};
#pragma unroll
    for (int r = 0; r < group_size; r += 2) {
      auto pair = Tx2{thread_data[r], thread_data[r + 1]};
      absmax_x2<Tx2>(amax_2x, amax_2x, pair);
    }
    float scale_dec_b =
        max(fabsf(static_cast<float>(amax_2x.x)),
            fabsf(static_cast<float>(amax_2x.y)));
    scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX;
    scale_dec_b *= scale_enc;
    // Convert to mx scale or nv scale
    using ScaleType = std::conditional_t<
        use_mx_scale,
        cutlass::float_ue8m0_t,
        cutlass::float_e4m3_t>;
    auto s = ScaleType(scale_dec_b);
    float scale_enc_b = scale_enc / float(s);
    scales_smem[tidx][tidy] = s.storage;

    int shared_idx = tidx * padded_local_cols + tidy * bytes_per_group;

#pragma unroll
    for (int j = 0; j < group_size / 4; j++) {
      Tx4 w_Tx4 = *reinterpret_cast<Tx4*>(&thread_data[j * 4]);
      if constexpr (bits == 8) {
        uint32_t quantized_val =
            scale_cvt_Tx4_to_fp8x4<T, USE_SR>(w_Tx4, scale_enc_b, rbits);
        *reinterpret_cast<uint32_t*>(&quantized_smem[shared_idx + j * 4]) =
            quantized_val;
      } else {
        uint16_t quantized_val =
            scale_cvt_Tx4_to_fp4x4<T, USE_SR>(w_Tx4, scale_enc_b, rbits);
        *reinterpret_cast<uint16_t*>(&quantized_smem[shared_idx + j * 2]) =
            quantized_val;
      }
    }
  }
  __syncthreads();

  int output_cols = K / elem_per_byte;
  int num_groups_per_row = K / group_size;
  int linear_tid = tidx + tidy * BLOCK_X;
  // Write back quantized values
#pragma unroll
  for (int i = linear_tid; i < bytes_per_block; i += BLOCK_X * BLOCK_Y) {
    int local_row = i / local_cols;
    int local_col = i % local_cols;

    int global_row = bidy * rows_per_block + local_row;
    int global_col = bidx * local_cols + local_col;

    if (global_row < M && global_col < output_cols) {
      int physical_idx = local_row * padded_local_cols + local_col;
      out[global_row * output_cols + global_col] = quantized_smem[physical_idx];
    }
  }
  // Write back scales
  constexpr int num_scales = BLOCK_X * BLOCK_Y;
#pragma unroll
  for (int i = linear_tid; i < num_scales; i += BLOCK_X * BLOCK_Y) {
    int local_row = i / BLOCK_Y;
    int local_col = i % BLOCK_Y;

    int global_row = bidy * BLOCK_X + local_row;
    int global_col = bidx * BLOCK_Y + local_col;

    if (global_row < M && global_col < num_groups_per_row) {
      scales[global_row * num_groups_per_row + global_col] =
          scales_smem[local_row][local_col];
    }
  }
}

template <typename T, int group_size, int bits, bool use_mx_scale>
__global__ void fp_dequantize(
    const uint8_t* w,
    const uint8_t* scales,
    T* out,
    size_t size,
    float* global_scale = nullptr) {
  auto block_size = cg::this_thread_block().dim_threads();
  auto block_idx = cg::this_thread_block().group_index();
  auto idx_in_block = cg::this_thread_block().thread_index();

  auto tidx = block_idx.x * block_size.x + idx_in_block.x;
  auto tidy = block_idx.y * block_size.y + idx_in_block.y;

  auto grid_dim_x = cg::this_grid().dim_blocks().x * block_size.x;

  constexpr int pack_factor = bits == 8 ? 1 : 2;
  const bool use_global_scale = global_scale != nullptr;
  const float inv_scale_enc = use_mx_scale
      ? 1.0f
      : (use_global_scale ? (*global_scale) / (F8E4M3_MAX * F4E2M1_MAX) : 1.0f);
  size_t offset = tidx + grid_dim_x * size_t(tidy);
  size_t oindex = offset * pack_factor;

  if (oindex >= size) {
    return;
  }

  size_t gindex = oindex / group_size;
  using ScaleType = std::conditional_t<
      use_mx_scale,
      cutlass::float_ue8m0_t,
      cutlass::float_e4m3_t>;
  auto scale = float(((ScaleType*)(scales))[gindex]) * inv_scale_enc;

  out += oindex;

  uint32_t val = w[offset];
#pragma clang loop unroll(full)
  for (int i = 0; i < pack_factor; i++) {
    uint8_t d;
    if (bits == 4) {
      d = (val >> (bits * i)) & 0x0f;
    } else if (bits == 8) {
      d = val;
    }
    out[i] = static_cast<T>(scale * Dequantize<bits>{}(d));
  }
}

inline std::tuple<dim3, dim3>
get_columnwise_quantize_launch_args(size_t size, int group_size, int M, int K) {
  constexpr int BLOCK_X = 16;
  constexpr int BLOCK_Y = 32;
  int rows_per_block = BLOCK_X;
  int cols_per_block = BLOCK_Y * group_size;

  dim3 grid;
  grid.x =
      cuda::ceil_div(K, cols_per_block) * cuda::ceil_div(M, rows_per_block);
  grid.y = 1;
  grid.z = 1;

  dim3 block(BLOCK_X, BLOCK_Y);

  return std::make_tuple(grid, block);
}

} // namespace cu

void fp_quantize_dequantize(
    const array& w,
    array& what,
    int group_size,
    int bits,
    const std::optional<array>& global_scale /* = std::nullopt */,
    cu::CommandEncoder& enc,
    const Stream& s) {
  enc.set_input_array(w);
  if (global_scale.has_value()) {
    enc.set_input_array(global_scale.value());
  }
  enc.set_output_array(what);
  dispatch_float_types(w.dtype(), "fp_quantize_dequantize", [&](auto type_tag) {
    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
    if constexpr (!std::is_same_v<T, double>) {
      auto kernel = cu::fp_quantize_dequantize<T, 32, 4, true, false>;
      if (bits == 8) {
        kernel = cu::fp_quantize_dequantize<T, 32, 8, true, false>;
      } else if (group_size == 16) {
        kernel = cu::fp_quantize_dequantize<T, 16, 4, false, false>;
      }
      bool large = w.size() > UINT_MAX;
      auto [num_blocks, block_dims] =
          get_launch_args(w.size(), w.shape(), w.strides(), large, group_size);

      enc.add_kernel_node(
          kernel,
          num_blocks,
          block_dims,
          gpu_ptr<T>(w),
          gpu_ptr<T>(what),
          w.size(),
          global_scale.has_value() ? gpu_ptr<float>(global_scale.value())
                                   : nullptr);
    }
  });
}

void fp_quantize(
    const array& w,
    array& wq,
    array& scales,
    int group_size,
    int bits,
    const std::optional<array>& global_scale /* = std::nullopt */,
    cu::CommandEncoder& enc,
    const Stream& s) {
  enc.set_input_array(w);
  if (global_scale.has_value()) {
    enc.set_input_array(global_scale.value());
  }
  enc.set_output_array(wq);
  enc.set_output_array(scales);
  if (w.strides().back() != 1) {
    dispatch_float_types(w.dtype(), "fp_quantize_columnwise", [&](auto type_tag) {
      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
      if constexpr (!std::is_same_v<T, double>) {
        auto M = w.shape(-2);
        auto K = w.shape(-1);
        auto kernel = cu::fp_quantize_columnwise<T, 32, 4, true, false>;
        if (bits == 8) {
          kernel = cu::fp_quantize_columnwise<T, 32, 8, true, false>;
        } else if (group_size == 16) {
          kernel = cu::fp_quantize_columnwise<T, 16, 4, false, false>;
        }
        auto [num_blocks, block_dims] =
            cu::get_columnwise_quantize_launch_args(w.size(), group_size, M, K);
        enc.add_kernel_node(
            kernel,
            num_blocks,
            block_dims,
            gpu_ptr<T>(w),
            gpu_ptr<uint8_t>(wq),
            gpu_ptr<uint8_t>(scales),
            w.size(),
            M,
            K,
            global_scale.has_value() ? gpu_ptr<float>(global_scale.value())
                                     : nullptr);
      } else {
        throw std::runtime_error(
            "[Quantize::eval_gpu] Can not quantize input with type float64.");
      }
    });
  } else {
    dispatch_float_types(w.dtype(), "fp_quantize_rowwise", [&](auto type_tag) {
      using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
      if constexpr (!std::is_same_v<T, double>) {
        auto kernel = cu::fp_quantize_rowwise<T, 32, 4, true, false>;
        if (bits == 8) {
          kernel = cu::fp_quantize_rowwise<T, 32, 8, true, false>;
        } else if (group_size == 16) {
          kernel = cu::fp_quantize_rowwise<T, 16, 4, false, false>;
        }
        bool large = w.size() > UINT_MAX;
        auto [num_blocks, block_dims] = get_launch_args(
            w.size(), w.shape(), w.strides(), large, group_size);

        enc.add_kernel_node(
            kernel,
            num_blocks,
            block_dims,
            gpu_ptr<T>(w),
            gpu_ptr<uint8_t>(wq),
            gpu_ptr<uint8_t>(scales),
            w.size(),
            global_scale.has_value() ? gpu_ptr<float>(global_scale.value())
                                     : nullptr);
      } else {
        throw std::runtime_error(
            "[Quantize::eval_gpu] Can not quantize input with type float64.");
      }
    });
  }
}

void fp_dequantize(
    const array& wq,
    const array& scales,
    array& w,
    int group_size,
    int bits,
    const std::optional<array>& global_scale /* = std::nullopt */,
    cu::CommandEncoder& enc,
    const Stream& s) {
  constexpr int uint8_per_uint32 = 4;
  int packs_per_int = 8 / bits;

  size_t size = w.size() / packs_per_int;
  bool large = size > UINT_MAX;
  auto grid_shape = w.shape();
  grid_shape.back() *= uint8_per_uint32;

  enc.set_input_array(wq);
  enc.set_input_array(scales);
  if (global_scale.has_value()) {
    enc.set_input_array(global_scale.value());
  }
  enc.set_output_array(w);
  dispatch_float_types(w.dtype(), "fp_dequantize", [&](auto type_tag) {
    using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
    if constexpr (!std::is_same_v<T, double>) {
      auto kernel = cu::fp_dequantize<T, 32, 4, true>;
      if (bits == 8) {
        kernel = cu::fp_dequantize<T, 32, 8, true>;
      } else if (group_size == 16) {
        kernel = cu::fp_dequantize<T, 16, 4, false>;
      }
      auto [num_blocks, block_dims] =
          get_launch_args(size, grid_shape, w.strides(), large);
      enc.add_kernel_node(
          kernel,
          num_blocks,
          block_dims,
          gpu_ptr<uint8_t>(wq),
          gpu_ptr<uint8_t>(scales),
          gpu_ptr<T>(w),
          w.size(),
          global_scale.has_value() ? gpu_ptr<float>(global_scale.value())
                                   : nullptr);
    } else {
      throw std::runtime_error(
          "[Quantize::eval_gpu] Can not dequantize to output with type float64.");
    }
  });
}

} // namespace mlx::core
