/usr/local/lib64/python3.6/site-packages/torch/include/c10d
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/c10d/frontend.hpp (9252B)
#pragma once
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
namespace c10d {
#ifdef USE_C10D_GLOO
static const std::string GLOO_SOCKET_IFNAME_ENV = "GLOO_SOCKET_IFNAME";
#endif
inline std::vector split(
char separator,
const std::string& string) {
std::vector pieces;
std::stringstream ss(string);
std::string item;
while (std::getline(ss, item, separator)) {
pieces.push_back(std::move(item));
}
return pieces;
}
class Backend {
public:
// Maps to Backend.__new__ in Python.
static std::string get(const std::string&);
// TODO: How to support registering third_party backend?
static void registerBackend();
private:
// TODO: Should this be an enum list instead since this set doesn't
// change at all.
std::unordered_set registered_backends_;
};
class TORCH_PYTHON_API DistributedC10d : public torch::CustomClassHolder {
public:
static c10::intrusive_ptr get();
DistributedC10d() = default;
void initProcessGroup(
const std::string& backend,
const std::string& init_method,
const std::chrono::milliseconds& timeout,
int64_t world_size,
int64_t rank,
c10::intrusive_ptr store,
const std::string& group_name);
void destroyProcessGroup(c10::intrusive_ptr group);
int64_t getRank(const c10::intrusive_ptr& group) const;
int64_t getWorldSize(const c10::intrusive_ptr& group) const;
c10::intrusive_ptr isend(
at::Tensor tensor,
int64_t dst,
const c10::intrusive_ptr& group,
c10::optional& tag);
c10::intrusive_ptr irecv(
at::Tensor tensor,
int64_t src,
const c10::intrusive_ptr& group,
c10::optional& tag);
void send(
at::Tensor tensor,
int64_t dst,
const c10::intrusive_ptr& group,
c10::optional& tag);
int64_t recv(
at::Tensor tensor,
const c10::optional& src,
const c10::intrusive_ptr& group,
c10::optional& tag);
c10::intrusive_ptr broadcastMultiGPU(
std::vector& tensor_list,
int64_t src,
const c10::intrusive_ptr& group,
bool async_op = false,
int64_t src_tensor = 0);
c10::intrusive_ptr broadcast(
at::Tensor tensor,
int64_t src,
const c10::intrusive_ptr& group,
bool async_op = false);
c10::intrusive_ptr allReduceMultiGPU(
std::vector& tensor_list,
const c10::intrusive_ptr& group,
ReduceOp op = ReduceOp::SUM,
bool async_op = false);
c10::intrusive_ptr allReduce(
at::Tensor tensor,
const c10::intrusive_ptr& group,
ReduceOp op = ReduceOp::SUM,
bool async_op = false);
c10::intrusive_ptr allReduceCoalesced(
std::vector& tensors,
const c10::intrusive_ptr& group,
ReduceOp op = ReduceOp::SUM,
bool async_op = false);
c10::intrusive_ptr reduceMultiGPU(
std::vector& tensor_list,
int64_t dst,
const c10::intrusive_ptr& group,
ReduceOp op = ReduceOp::SUM,
bool async_op = false,
int64_t dst_tensor = 0);
c10::intrusive_ptr reduce(
at::Tensor tensor,
int64_t dst,
const c10::intrusive_ptr& group,
ReduceOp op = ReduceOp::SUM,
bool async_op = false);
c10::intrusive_ptr allGatherMultiGPU(
std::vector>& output_tensor_lists,
std::vector& input_tensor_list,
const c10::intrusive_ptr& group,
bool async_op = false);
c10::intrusive_ptr allGather(
std::vector& tensor_list,
at::Tensor tensor,
const c10::intrusive_ptr& group,
bool async_op = false);
c10::intrusive_ptr allGatherCoalesced(
std::vector>& output_tensor_lists,
std::vector& input_tensor_list,
const c10::intrusive_ptr& group,
bool async_op = false);
c10::intrusive_ptr gather(
at::Tensor tensor,
const c10::optional>& gather_list,
const c10::intrusive_ptr& group,
int64_t dst = 0,
bool async_op = false);
c10::intrusive_ptr scatter(
at::Tensor tensor,
std::vector& scatter_list,
const c10::intrusive_ptr& group,
int64_t src = 0,
bool async_op = false);
c10::intrusive_ptr reduceScatterMultiGPU(
std::vector& output_tensor_list,
std::vector>& input_tensor_lists,
const c10::intrusive_ptr& group,
ReduceOp op = ReduceOp::SUM,
bool async_op = false);
c10::intrusive_ptr reduceScatter(
at::Tensor output,
std::vector& input_tensor_list,
const c10::intrusive_ptr& group,
ReduceOp op = ReduceOp::SUM,
bool async_op = false);
c10::intrusive_ptr allToAllSingle(
at::Tensor output,
at::Tensor input,
std::vector& output_split_sizes,
std::vector& input_split_sizes,
const c10::intrusive_ptr& group,
bool async_op = false);
c10::intrusive_ptr allToAll(
std::vector& output_tensor_list,
std::vector& input_tensor_list,
const c10::intrusive_ptr& group,
bool async_op = false);
c10::intrusive_ptr barrier(
const c10::intrusive_ptr& group,
bool async_op = false);
c10::intrusive_ptr newGroup(
std::vector ranks,
std::chrono::milliseconds timeout,
Backend backend);
c10::intrusive_ptr worldProcessGroup();
c10::intrusive_ptr newProcessGroupHelper(
const int64_t world_size,
const int64_t rank,
const std::vector& group_ranks,
const std::string& backend_str,
const c10::intrusive_ptr& store,
c10::optional group_name,
int64_t timeout_milisesonds);
c10::intrusive_ptr getProcessGroupByName(
const std::string& name) const;
std::string getNameOfProcessGroup(
const c10::intrusive_ptr& pg) const;
void registerProcessGroupName(const c10::intrusive_ptr& process_group, const std::string& name);
private:
bool rankNotInGroup(const c10::intrusive_ptr& group) const;
int64_t getGroupRank(
const c10::intrusive_ptr& group,
const int64_t rank) const;
int64_t getGlobalRank(
const c10::intrusive_ptr& group,
const int64_t group_rank) const;
void checkDefaultPg() const;
int64_t getGroupSize(const c10::intrusive_ptr& group) const;
std::string getBackend(const c10::intrusive_ptr& group);
std::string backend_;
// TODO: Ask Alex what kind of equality we need. It determine whether we
// need to use ProcessGroup or ProcesGroup* as key.
std::unordered_map<
c10::intrusive_ptr,
std::pair>>
pg_map_;
// Note, this is different mapping relationship than original Python
// implementation.
std::unordered_map, std::string> pg_names_;
// Process group's global rank to local rank mapping
std::unordered_map<
c10::intrusive_ptr,
std::unordered_map>
pg_group_ranks_;
c10::intrusive_ptr default_pg_;
// Default value should be "env://"
std::string default_pg_init_method_;
int64_t group_count_;
};
// This class exists as a way to allow us to split NCCL-specific code into a
// different file. frontend_cuda.cpp will, if USE_C10D_NCCL is defined,
// override this NCCLProcessGroupProvider with one that will actually do
// something.
struct TORCH_API NCCLProcessGroupProvider {
virtual c10::intrusive_ptr get(
c10::intrusive_ptr /*prefix_store*/,
int64_t /*rank*/,
int64_t /*world_size*/,
std::chrono::milliseconds /*timeout*/) const {
AT_ERROR(
"Attempting to create NCCL-based process group while NCCL is either not enabled or built");
}
virtual ~NCCLProcessGroupProvider() = default;
};
TORCH_API void registerNCCLProcessGroupProvider(
NCCLProcessGroupProvider* provider);
TORCH_API void initCustomClassBindings();
} // namespace c10d