engine.h

Go to the documentation of this file.

 
#ifndef RABIT_INTERNAL_ENGINE_H_
#define RABIT_INTERNAL_ENGINE_H_
#include <string>
#include "rabit/serializable.h"
 
namespace MPI {  // NOLINT
class Datatype;
}
 
namespace rabit {
namespace engine {
class IEngine {
 public:
  typedef void (PreprocFunction) (void *arg);  // NOLINT
  typedef void (ReduceFunction) (const void *src,  // NOLINT
                                 void *dst, int count,
                                 const MPI::Datatype &dtype);
  ~IEngine() = default;
  virtual void Allgather(void *sendrecvbuf,
                         size_t total_size,
                         size_t slice_begin,
                         size_t slice_end,
                         size_t size_prev_slice) = 0;
  virtual void Allreduce(void *sendrecvbuf_,
                         size_t type_nbytes,
                         size_t count,
                         ReduceFunction reducer,
                         PreprocFunction prepare_fun = nullptr,
                         void *prepare_arg = nullptr) = 0;
  virtual void Broadcast(void *sendrecvbuf_, size_t size, int root) = 0;
  virtual int LoadCheckPoint() = 0;
  virtual void CheckPoint() = 0;
  virtual int VersionNumber() const = 0;
  virtual int GetRingPrevRank() const = 0;
  virtual int GetRank() const = 0;
  virtual int GetWorldSize() const = 0;
  virtual bool IsDistributed() const = 0;
  virtual std::string GetHost() const = 0;
  virtual void TrackerPrint(const std::string &msg) = 0;
};
 
bool Init(int argc, char *argv[]);
bool Finalize();
IEngine *GetEngine();
 
namespace mpi {
enum OpType {
  kMax = 0,
  kMin = 1,
  kSum = 2,
  kBitwiseAND = 3,
  kBitwiseOR = 4,
  kBitwiseXOR = 5,
};
enum DataType {
  kChar = 0,
  kUChar = 1,
  kInt = 2,
  kUInt = 3,
  kLong = 4,
  kULong = 5,
  kFloat = 6,
  kDouble = 7,
  kLongLong = 8,
  kULongLong = 9
};
}  // namespace mpi
void Allgather(void* sendrecvbuf,
                   size_t total_size,
                   size_t slice_begin,
                   size_t slice_end,
                   size_t size_prev_slice);
void Allreduce_(void *sendrecvbuf,  // NOLINT
                size_t type_nbytes,
                size_t count,
                IEngine::ReduceFunction red,
                mpi::DataType dtype,
                mpi::OpType op,
                IEngine::PreprocFunction prepare_fun = nullptr,
                void *prepare_arg = nullptr);
}  // namespace engine
}  // namespace rabit
#endif  // RABIT_INTERNAL_ENGINE_H_