#pragma once
#include <cstddef>
#include <cstdint>
#include <memory>
#include <vector>
#include "vt/device.h"
#include "vt/dtype.h"
namespace vt {
enum class ReduceOp : uint8_t { kSum = 0, kMax, kMin, kProd };
class Communicator {
public:
virtual ~Communicator() = default;
virtual int rank() const = 0;
virtual int world_size() const = 0;
virtual void AllReduce(Queue& q, void* data, size_t count, DType dtype,
ReduceOp op) = 0;
virtual void AllGather(Queue& q, const void* sendbuf, void* recvbuf,
size_t count, DType dtype) = 0;
virtual void Send(Queue& q, const void* data, size_t count, DType dtype,
int peer) = 0;
virtual void Recv(Queue& q, void* data, size_t count, DType dtype,
int peer) = 0;
};
using CommAllReduceFn = void (*)(Communicator&, Queue&, void*, size_t, DType,
ReduceOp);
using CommAllGatherFn = void (*)(Communicator&, Queue&, const void*, void*,
size_t, DType);
using CommSendFn = void (*)(Communicator&, Queue&, const void*, size_t, DType,
int);
using CommRecvFn = void (*)(Communicator&, Queue&, void*, size_t, DType, int);
class CpuCommGroup {
public:
explicit CpuCommGroup(int world_size);
~CpuCommGroup();
CpuCommGroup(const CpuCommGroup&) = delete;
CpuCommGroup& operator=(const CpuCommGroup&) = delete;
int world_size() const;
Communicator& Rank(int r);
struct Shared;
private:
std::shared_ptr<Shared> shared_;
std::vector<std::unique_ptr<Communicator>> ranks_;
};
}