#include <torch/library.h>

#include "registration.h"
#include "torch_binding.h"

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("quantize_act(Tensor! Aq, Tensor! scale, Tensor a) -> ()");
  ops.def("bitnet_gemv_fused(Tensor! out, Tensor a_bf16, Tensor w_packed, Tensor scale_wt) -> ()");
  ops.def("bitnet_gemm(Tensor! out, Tensor a_int8, Tensor w_packed, Tensor scale_act, Tensor scale_wt, Tensor? scratch) -> ()");
  ops.def("ternarize_pack(Tensor! w_packed, Tensor! gamma, Tensor W) -> ()");
  ops.def("dequant_w(Tensor! W_hat, Tensor w_packed, Tensor gamma) -> ()");
  ops.def("dequant_x(Tensor! X_hat, Tensor x_q, Tensor scale) -> ()");

#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
  ops.impl("quantize_act",      torch::kCUDA, &quantize_act);
  ops.impl("bitnet_gemv_fused", torch::kCUDA, &bitnet_gemv_fused);
  ops.impl("bitnet_gemm",       torch::kCUDA, &bitnet_gemm);
  ops.impl("ternarize_pack",    torch::kCUDA, &ternarize_pack);
  ops.impl("dequant_w",         torch::kCUDA, &dequant_w);
  ops.impl("dequant_x",         torch::kCUDA, &dequant_x);
#endif
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)