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