#include "direct_w4a4_m1.h"

#include <torch/extension.h>

#include <optional>

namespace {

torch::Tensor direct_w4a4_m1_linear(
    const torch::Tensor& activation_qdata,
    const torch::Tensor& activation_block_scales,
    const torch::Tensor& activation_tensor_scale,
    const torch::Tensor& weight_qdata,
    const torch::Tensor& weight_block_scales,
    const torch::Tensor& weight_tensor_scale,
    const std::optional<torch::Tensor>& bias) {
  TORCH_CHECK(
      activation_qdata.is_cuda() &&
          activation_qdata.scalar_type() == at::kByte &&
          activation_qdata.dim() == 2 && activation_qdata.is_contiguous(),
      "activation qdata must be contiguous CUDA uint8 [padded_M,K/2]");
  TORCH_CHECK(
      activation_qdata.size(0) >= 1,
      "direct W4A4 M=1 requires at least one packed activation row");
  TORCH_CHECK(
      activation_block_scales.is_cuda() &&
          activation_block_scales.dim() == 2 &&
          activation_block_scales.is_contiguous() &&
          activation_block_scales.element_size() == 1,
      "activation block scales must be contiguous CUDA byte-sized [padded_M,padded_K/16]");
  TORCH_CHECK(
      activation_tensor_scale.is_cuda() &&
          activation_tensor_scale.scalar_type() == at::kFloat &&
          activation_tensor_scale.numel() == 1 &&
          activation_tensor_scale.is_contiguous(),
      "activation tensor scale must be one contiguous CUDA float32 value");
  TORCH_CHECK(
      weight_qdata.is_cuda() && weight_qdata.scalar_type() == at::kByte &&
          weight_qdata.dim() == 2 && weight_qdata.is_contiguous(),
      "weight qdata must be contiguous CUDA uint8 [N,K/2]");
  TORCH_CHECK(
      weight_block_scales.is_cuda() && weight_block_scales.dim() == 2 &&
          weight_block_scales.is_contiguous() &&
          weight_block_scales.element_size() == 1,
      "weight block scales must be contiguous CUDA byte-sized [padded_N,padded_K/16]");
  TORCH_CHECK(
      weight_tensor_scale.is_cuda() &&
          weight_tensor_scale.scalar_type() == at::kFloat &&
          weight_tensor_scale.numel() == 1 &&
          weight_tensor_scale.is_contiguous(),
      "weight tensor scale must be one contiguous CUDA float32 value");
  TORCH_CHECK(
      activation_qdata.device() == activation_block_scales.device() &&
          activation_qdata.device() == activation_tensor_scale.device() &&
          activation_qdata.device() == weight_qdata.device() &&
          activation_qdata.device() == weight_block_scales.device() &&
          activation_qdata.device() == weight_tensor_scale.device(),
      "all direct W4A4 tensors must use the same CUDA device");

  const int64_t in_features = activation_qdata.size(1) * 2;
  const int64_t out_features = weight_qdata.size(0);
  TORCH_CHECK(
      in_features > 0 && in_features % 32 == 0,
      "direct W4A4 M=1 requires K divisible by 32");
  TORCH_CHECK(
      weight_qdata.size(1) == activation_qdata.size(1),
      "activation and weight packed K dimensions differ");
  TORCH_CHECK(out_features > 0, "direct W4A4 M=1 requires positive N");
  TORCH_CHECK(
      activation_block_scales.size(0) >= 1 &&
          activation_block_scales.size(1) >= in_features / 16,
      "activation block-scale tensor is too small");
  TORCH_CHECK(
      weight_block_scales.size(0) >= out_features &&
          weight_block_scales.size(1) >= in_features / 16,
      "weight block-scale tensor is too small");

  if (bias.has_value()) {
    const auto& value = *bias;
    TORCH_CHECK(
        value.is_cuda() && value.scalar_type() == at::kBFloat16 &&
            value.dim() == 1 && value.is_contiguous() &&
            value.numel() == out_features &&
            value.device() == activation_qdata.device(),
        "bias must be contiguous CUDA bfloat16 [N]");
  }
  return direct_w4a4_m1_linear_cuda(
      activation_qdata,
      activation_block_scales,
      activation_tensor_scale,
      weight_qdata,
      weight_block_scales,
      weight_tensor_scale,
      bias);
}

}  // namespace

PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
  module.def(
      "linear",
      &direct_w4a4_m1_linear,
      "Direct packed-NVFP4 activation x packed-NVFP4 weight M=1 linear");
}