#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, 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) -> ()");

#if defined(CPU_KERNEL)
  ops.impl("quantize_act",      torch::kCPU, &quantize_act);
  ops.impl("bitnet_gemv_fused", torch::kCPU, &bitnet_gemv_fused);
  ops.impl("bitnet_gemm",       torch::kCPU, &bitnet_gemm);
#endif
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)