// 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/fp_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>

namespace mlx::core {
namespace cu {

inline void create_2D_tensor_map(
    CUtensorMap* tensorMap,
    void* input_ptr,
    CUtensorMapDataType dtype,
    uint64_t rows,
    uint64_t cols,
    uint32_t tile_y,
    uint32_t tile_x,
    uint64_t stride_bytes,
    CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_NONE) {
  constexpr uint32_t rank = 2; // 2D
  uint64_t global_dim[rank] = {cols, rows};
  // For row-major layout
  uint64_t strides[rank - 1] = {stride_bytes};
  uint32_t tile_dim[rank] = {tile_x, tile_y};
  uint32_t elem_stride[rank] = {1, 1};

  CHECK_CUDA_ERROR(cuTensorMapEncodeTiled(
      tensorMap,
      dtype,
      rank,
      input_ptr,
      global_dim,
      strides,
      tile_dim,
      elem_stride,
      CU_TENSOR_MAP_INTERLEAVE_NONE,
      swizzle,
      CU_TENSOR_MAP_L2_PROMOTION_NONE,
      CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
}

inline std::tuple<dim3, dim3> get_columnwise_quantize_mxfp8_launch_args(
    size_t grid_dim_x_size, // rows
    size_t grid_dim_y_size, // cols
    size_t block_size_x, // ROWS_PER_BLOCK
    size_t block_size_y, // COL_PER_BLOCK
    int threads_per_block) {
  dim3 grid;
  grid.x = cuda::ceil_div(grid_dim_x_size, block_size_x);
  grid.y = cuda::ceil_div(grid_dim_y_size, block_size_y);
  grid.z = 1;

  dim3 block(threads_per_block, 1, 1);
  return std::make_tuple(grid, block);
}

inline CUtensorMapDataType get_tma_dtype(Dtype dtype) {
  switch (dtype) {
    case float16:
      return CU_TENSOR_MAP_DATA_TYPE_FLOAT16;
    case bfloat16:
      return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16;
    case float32:
      return CU_TENSOR_MAP_DATA_TYPE_FLOAT32;
    default:
      throw std::runtime_error(
          "[fp_quantize_columnwise_tma] Unsupported dtype for TMA");
  }
}

inline std::tuple<dim3, dim3> get_columnwise_quantize_fallback_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_rowwise(
    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);
  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_quantize_columnwise_fallback(
    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);
  dispatch_float_types(
      w.dtype(), "fp_quantize_columnwise_fallback", [&](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_fallback<T, 32, 4, true, false>;
          if (bits == 8) {
            kernel = cu::fp_quantize_columnwise_fallback<T, 32, 8, true, false>;
          } else if (group_size == 16) {
            kernel =
                cu::fp_quantize_columnwise_fallback<T, 16, 4, false, false>;
          }
          auto [num_blocks, block_dims] =
              cu::get_columnwise_quantize_fallback_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.");
        }
      });
}

void fp_quantize_columnwise_mxfp8(
    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);
  enc.set_output_array(wq);
  enc.set_output_array(scales);

  size_t rows = w.shape(-1);
  size_t cols = w.size() / rows;
  size_t stride_bytes = w.strides(-1) * w.itemsize();

  dispatch_float_types(
      w.dtype(), "fp_quantize_columnwise_mxfp8", [&](auto type_tag) {
        using T = cuda_type_t<MLX_GET_TYPE(type_tag)>;
        if constexpr (!std::is_same_v<T, double>) {
          constexpr int THREADS_PER_BLOCK = 64;
          constexpr int ROWS_PER_BLOCK = 64;
          constexpr int COLS_PER_BLOCK = 64;
          constexpr size_t TILE_M = 32;
          constexpr size_t TILE_K = COLS_PER_BLOCK;
          constexpr size_t STAGES = ROWS_PER_BLOCK / TILE_M;

          // For columnwise: grid.x = cols, grid.y = rows
          // scales_per_stage = TILE_K (one scale per column per stage)
          auto [grid, block] = cu::get_columnwise_quantize_mxfp8_launch_args(
              cols, rows, COLS_PER_BLOCK, ROWS_PER_BLOCK, THREADS_PER_BLOCK);

          CUtensorMap tensor_map_input;
          CUtensorMap tensor_map_output;

          cu::create_2D_tensor_map(
              &tensor_map_input,
              gpu_ptr<void>(w),
              cu::get_tma_dtype(w.dtype()),
              rows,
              cols,
              TILE_M,
              TILE_K,
              stride_bytes);

          cu::create_2D_tensor_map(
              &tensor_map_output,
              gpu_ptr<uint8_t>(wq),
              CU_TENSOR_MAP_DATA_TYPE_UINT8,
              cols,
              rows,
              TILE_K,
              TILE_M,
              rows);

          auto kernel = cu::fp_quantize_columnwise_mxfp8<
              T,
              false,
              THREADS_PER_BLOCK,
              COLS_PER_BLOCK,
              ROWS_PER_BLOCK>;
          enc.add_kernel_node(
              kernel,
              grid,
              block,
              tensor_map_input,
              tensor_map_output,
              gpu_ptr<uint8_t>(scales),
              rows,
              cols);
        } else {
          throw std::runtime_error(
              "[fp_quantize_columnwise_tma] Cannot quantize input with type float64.");
        }
      });
}

void fp_quantize_columnwise(
    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) {
  // Use TMA version for SM100+ with MXFP8 (bits=8, group_size=32)
  // NVFP4 todo
  const size_t rows = w.shape(-1);
  const size_t cols = w.size() / rows;
  const bool has_full_tma_tiles = ((rows % 128) == 0) && ((cols % 128) == 0);
  bool use_tma =
      (enc.device().compute_capability_major() >= 10 && bits == 8 &&
       group_size == 32 && has_full_tma_tiles);
  if (use_tma) {
    fp_quantize_columnwise_mxfp8(
        w, wq, scales, group_size, bits, global_scale, enc, s);
  } else {
    fp_quantize_columnwise_fallback(
        w, wq, scales, group_size, bits, global_scale, enc, s);
  }
}

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) {
  if (w.strides(-1) == 1) {
    fp_quantize_rowwise(w, wq, scales, group_size, bits, global_scale, enc, s);
  } else {
    fp_quantize_columnwise(
        w, wq, scales, group_size, bits, global_scale, enc, s);
  }
}

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
