| 1 | #include "ops.h" |
| 2 | #include "core/registration.h" |
| 3 | |
| 4 | #include <torch/csrc/stable/library.h> |
| 5 | |
| 6 | STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ag_rs) { |
| 7 | custom_ag_rs.def( |
| 8 | "custom_all_gather(int fa, Tensor inp, Tensor! out, int reg_buffer, " |
| 9 | "int reg_buffer_sz_bytes) -> ()"); |
| 10 | custom_ag_rs.def( |
| 11 | "mnnvl_lamport_all_gather(int fa, Tensor inp, Tensor! out, int " |
| 12 | "local_buffer, int multicast_buffer, int epoch_buffer, int " |
| 13 | "stage_sz_bytes) -> ()"); |
| 14 | custom_ag_rs.def( |
| 15 | "custom_reduce_scatter(int fa, Tensor inp, Tensor! out, int reg_buffer, " |
| 16 | "int reg_buffer_sz_bytes) -> ()"); |
| 17 | custom_ag_rs.def( |
| 18 | "mnnvl_lamport_reduce_scatter(int fa, Tensor inp, Tensor! out, int " |
| 19 | "local_buffer, int epoch_buffer, int stage_sz_bytes) -> ()"); |
| 20 | custom_ag_rs.def( |
| 21 | "mnnvl_multimem_reduce_scatter(int fa, Tensor inp, Tensor! out, int " |
| 22 | "local_buffer, int multicast_buffer, int stage_sz_bytes, int " |
| 23 | "block_limit) -> ()"); |
| 24 | } |
| 25 | |
| 26 | STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ag_rs) { |
| 27 | custom_ag_rs.impl("custom_all_gather", TORCH_BOX(&custom_all_gather)); |
| 28 | custom_ag_rs.impl("mnnvl_lamport_all_gather", |
| 29 | TORCH_BOX(&mnnvl_lamport_all_gather)); |
| 30 | custom_ag_rs.impl("custom_reduce_scatter", TORCH_BOX(&custom_reduce_scatter)); |
| 31 | custom_ag_rs.impl("mnnvl_lamport_reduce_scatter", |
| 32 | TORCH_BOX(&mnnvl_lamport_reduce_scatter)); |
| 33 | custom_ag_rs.impl("mnnvl_multimem_reduce_scatter", |
| 34 | TORCH_BOX(&mnnvl_multimem_reduce_scatter)); |
| 35 | } |