#include <torch/library.h>

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

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("exact_gemm_maxp() -> int");
  ops.def("exact_gemm_lmax() -> int");
  ops.def("exact_gemm_mm(Tensor! Cout, Tensor A, Tensor B, Tensor primes_cpu, Tensor inv_cpu, Tensor mu_cpu, Tensor pow32_cpu) -> ()");
  ops.def("exact_gemm_mm_timed(Tensor! Cout, Tensor! times, Tensor A, Tensor B, Tensor primes_cpu, Tensor inv_cpu, Tensor mu_cpu, Tensor pow32_cpu) -> ()");

  ops.impl("exact_gemm_maxp", &exact_gemm_maxp);
  ops.impl("exact_gemm_lmax", &exact_gemm_lmax);
#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
  ops.impl("exact_gemm_mm", torch::kCUDA, &exact_gemm_mm);
  ops.impl("exact_gemm_mm_timed", torch::kCUDA, &exact_gemm_mm_timed);
#endif
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)