45 lines
1.4 KiB
C++
45 lines
1.4 KiB
C++
#pragma once
|
|
|
|
#include <ATen/ATen.h>
|
|
#include <c10/macros/Macros.h>
|
|
|
|
namespace c10d::nccl_extension {
|
|
|
|
TORCH_API bool is_nccl_symmem_available();
|
|
|
|
TORCH_API void nccl_put(at::Tensor& tensor, const int64_t peer);
|
|
|
|
TORCH_API void nccl_get(at::Tensor& tensor, const int64_t peer);
|
|
|
|
TORCH_API void nccl_wait_for_signal(at::Tensor& sigpad, int64_t signal);
|
|
|
|
TORCH_API void nccl_put_with_signal(
|
|
at::Tensor& tensor,
|
|
int64_t signal,
|
|
int64_t peer);
|
|
|
|
// Simultaneously reduce N blocks of a 2-D input tensor from a shared symmetric
|
|
// memory buffer, routing each to a specific destination rank. Blocks are
|
|
// described by inclusive-prefix-sum offsets along `dim` (0 or 1); all blocks
|
|
// must have equal size.
|
|
TORCH_API void nccl_reduce_scatter_offset(
|
|
const at::Tensor& input,
|
|
at::TensorList out,
|
|
const std::string& group_name,
|
|
int64_t dim,
|
|
std::optional<at::IntArrayRef> offsets,
|
|
std::optional<at::IntArrayRef> dst_ranks,
|
|
const std::string& red_op);
|
|
|
|
// Permute-free all-to-all: scatter shards along `scatter_dim`, gather peer
|
|
// chunks along `gather_dim`. Supported: (scatter_dim=1, gather_dim=0) -> [rows,
|
|
// p*lc] to [p, rows, lc]; (scatter_dim=0, gather_dim=1) -> [p*lr, cols] to [lr,
|
|
// p, cols].
|
|
TORCH_API void nccl_all_to_all_permute(
|
|
const at::Tensor& input,
|
|
at::Tensor& out,
|
|
int64_t scatter_dim,
|
|
int64_t gather_dim,
|
|
const std::string& group_name);
|
|
} // namespace c10d::nccl_extension
|