| | #include <torch/library.h> |
| |
|
| | #include "registration.h" |
| | #include "torch_binding.h" |
| |
|
| | TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { |
| | ops.def("__KERNEL_NAME_NORMALIZED__(Tensor! out, Tensor input) -> ()"); |
| | #if defined(CPU_KERNEL) |
| | ops.impl("__KERNEL_NAME_NORMALIZED__", torch::kCPU, &__KERNEL_NAME_NORMALIZED__); |
| | #elif defined(CUDA_KERNEL) || defined(ROCM_KERNEL) |
| | ops.impl("__KERNEL_NAME_NORMALIZED__", torch::kCUDA, &__KERNEL_NAME_NORMALIZED__); |
| | #elif defined(METAL_KERNEL) |
| | ops.impl("__KERNEL_NAME_NORMALIZED__", torch::kMPS, __KERNEL_NAME_NORMALIZED__); |
| | #elif defined(XPU_KERNEL) |
| | ops.impl("__KERNEL_NAME_NORMALIZED__", torch::kXPU, &__KERNEL_NAME_NORMALIZED__); |
| | #endif |
| | } |
| |
|
| | REGISTER_EXTENSION(TORCH_EXTENSION_NAME) |
| |
|
