#include <torch/library.h>

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

TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("exact_solve_extract(Tensor! planes, Tensor X, Tensor primes, Tensor mu, Tensor pow32) -> ()");
  ops.def("exact_solve_matmul(Tensor! Cout, Tensor A_planes, Tensor B_planes, Tensor primes, Tensor inv, Tensor mu) -> ()");
  ops.def("exact_solve_mod_reduce(Tensor! out, Tensor C, int p, Tensor pw) -> ()");
  ops.def("exact_solve_residual(Tensor! out, Tensor r, Tensor t, int p) -> ()");
  ops.def("exact_solve_batched_modinv(Tensor! Bout, Tensor A, Tensor primes, Tensor! work, Tensor! singular) -> ()");
  ops.def("exact_solve_i64_matmul(Tensor! out, Tensor A, Tensor X) -> ()");
  ops.def("exact_solve_padic_residual(Tensor! out, Tensor r, Tensor axlow, int Pinv64) -> ()");
  ops.def("exact_solve_gpu_recon(Tensor! out, Tensor digits, Tensor Pl, Tensor Pm, Tensor Rr, Tensor Pmhalf, Tensor dl) -> ()");

#if defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
  ops.impl("exact_solve_extract", torch::kCUDA, &exact_solve_extract);
  ops.impl("exact_solve_matmul", torch::kCUDA, &exact_solve_matmul);
  ops.impl("exact_solve_mod_reduce", torch::kCUDA, &exact_solve_mod_reduce);
  ops.impl("exact_solve_residual", torch::kCUDA, &exact_solve_residual);
  ops.impl("exact_solve_batched_modinv", torch::kCUDA, &exact_solve_batched_modinv);
  ops.impl("exact_solve_i64_matmul", torch::kCUDA, &exact_solve_i64_matmul);
  ops.impl("exact_solve_padic_residual", torch::kCUDA, &exact_solve_padic_residual);
  ops.impl("exact_solve_gpu_recon", torch::kCUDA, &exact_solve_gpu_recon);
#endif
}

REGISTER_EXTENSION(TORCH_EXTENSION_NAME)