forked from jiuyuan/InfiniTensor
23 lines
450 B
C
23 lines
450 B
C
|
#pragma once
|
||
|
|
#include "object.h"
|
||
|
|
#include "ref.h"
|
||
|
|
|
||
|
|
namespace infini {
|
||
|
|
|
||
|
|
// base class
|
||
|
|
class CommunicatorObj : public Object {
|
||
|
|
protected:
|
||
|
|
int worldSize;
|
||
|
|
int rank;
|
||
|
|
|
||
|
|
public:
|
||
|
|
CommunicatorObj(int worldSize, int rank)
|
||
|
|
: worldSize(worldSize), rank(rank) {}
|
||
|
|
|
||
|
|
virtual ~CommunicatorObj() = default;
|
||
|
|
virtual int getWorldSize() const { return worldSize; }
|
||
|
|
virtual int getRank() const { return rank; }
|
||
|
|
};
|
||
|
|
|
||
|
|
} // namespace infini
|