| #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) |
|
|
