| 1 | // Provides torch::Tensor for ops.h (previously included transitively via |
| 2 | // cache.h, which is no longer included here after cache ops moved to |
| 3 | // _C_stable_libtorch). |
| 4 | #include <torch/all.h> |
| 5 | #include "ops.h" |
| 6 | #include "core/registration.h" |
| 7 | #include <torch/library.h> |
| 8 | #include <torch/version.h> |
| 9 | |
| 10 | // Note on op signatures: |
| 11 | // The X_meta signatures are for the meta functions corresponding to op X. |
| 12 | // They must be kept in sync with the signature for X. Generally, only |
| 13 | // functions that return Tensors require a meta function. |
| 14 | // |
| 15 | // See the following links for detailed docs on op registration and function |
| 16 | // schemas. |
| 17 | // https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.ptttacy8y1u9 |
| 18 | // https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/README.md#annotations |
| 19 | |
| 20 | #ifdef USE_ROCM |
| 21 | TORCH_LIBRARY_FRAGMENT(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) { |
| 22 | // Quick Reduce all-reduce kernels (ROCm-only; stays on legacy _C). |
| 23 | custom_ar.def( |
| 24 | "qr_all_reduce(int fa, Tensor inp, Tensor out, int quant_level, bool " |
| 25 | "cast_bf2half) -> ()"); |
| 26 | custom_ar.impl("qr_all_reduce", torch::kCUDA, &qr_all_reduce); |
| 27 | |
| 28 | custom_ar.def("init_custom_qr", &init_custom_qr); |
| 29 | custom_ar.def("qr_destroy", &qr_destroy); |
| 30 | custom_ar.def("qr_get_handle", &qr_get_handle); |
| 31 | |
| 32 | custom_ar.def("qr_open_handles(int _fa, Tensor[](b!) handles) -> ()"); |
| 33 | custom_ar.impl("qr_open_handles", torch::kCPU, &qr_open_handles); |
| 34 | |
| 35 | custom_ar.def("qr_max_size", &qr_max_size); |
| 36 | } |
| 37 | #endif |
| 38 | |
| 39 | REGISTER_EXTENSION(TORCH_EXTENSION_NAME) |