#include "smallm_gemv.h"

#include <torch/extension.h>

#include <optional>

namespace {

torch::Tensor smallm_nvfp4_linear(
    const torch::Tensor& input,
    const torch::Tensor& packed_weight,
    const torch::Tensor& weight_block_scales,
    const torch::Tensor& weight_tensor_scale,
    const std::optional<torch::Tensor>& bias) {
  TORCH_CHECK(input.is_cuda(), "small-M GEMV requires a CUDA input");
  TORCH_CHECK(
      input.scalar_type() == at::kBFloat16,
      "small-M GEMV input must be bfloat16");
  TORCH_CHECK(
      input.dim() >= 1 && input.is_contiguous(),
      "small-M GEMV input must be contiguous");
  TORCH_CHECK(
      packed_weight.is_cuda() && packed_weight.scalar_type() == at::kByte &&
          packed_weight.dim() == 2 && packed_weight.is_contiguous(),
      "packed weight 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(
      input.device() == packed_weight.device() &&
          input.device() == weight_block_scales.device() &&
          input.device() == weight_tensor_scale.device(),
      "all small-M GEMV tensors must use the same CUDA device");

  const int64_t out_features = packed_weight.size(0);
  const int64_t in_features = packed_weight.size(1) * 2;
  TORCH_CHECK(
      input.size(-1) == in_features,
      "small-M GEMV expected input width ",
      in_features,
      " but got ",
      input.size(-1));
  TORCH_CHECK(
      in_features > 0 && in_features % 32 == 0,
      "small-M GEMV requires K divisible by 32");
  TORCH_CHECK(
      out_features > 0,
      "small-M GEMV requires positive N");
  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() == input.device(),
        "bias must be contiguous CUDA bfloat16 [N]");
  }
  return smallm_nvfp4_linear_cuda(
      input,
      packed_weight,
      weight_block_scales,
      weight_tensor_scale,
      bias);
}

}  // namespace

PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
  module.def(
      "linear",
      &smallm_nvfp4_linear,
      "Fused BF16-activation x packed-NVFP4-weight small-M GEMV");
}