engine.cc

Go to the documentation of this file.

 
#include <rabit/base.h>
#include <dmlc/thread_local.h>
 
#include <memory>
#include "rabit/internal/engine.h"
#include "allreduce_base.h"
 
namespace rabit {
namespace engine {
// singleton sync manager
#ifndef RABIT_USE_BASE
#ifndef RABIT_USE_MOCK
using Manager = AllreduceBase;
#else
typedef AllreduceMock Manager;
#endif  // RABIT_USE_MOCK
#else
typedef AllreduceBase Manager;
#endif  // RABIT_USE_BASE
 
struct ThreadLocalEntry {
  std::unique_ptr<Manager> engine;
  bool initialized{false};
  ThreadLocalEntry() = default;
};
 
// define the threadlocal store.
using EngineThreadLocal = dmlc::ThreadLocalStore<ThreadLocalEntry>;
 
bool Init(int argc, char *argv[]) {
  ThreadLocalEntry* e = EngineThreadLocal::Get();
  if (e->engine.get() == nullptr) {
    e->initialized = true;
    e->engine.reset(new Manager());
    return e->engine->Init(argc, argv);
  } else {
    return true;
  }
}
 
bool Finalize() {
  ThreadLocalEntry* e = EngineThreadLocal::Get();
  if (e->engine.get() != nullptr) {
    if (e->engine->Shutdown()) {
      e->engine.reset(nullptr);
      e->initialized = false;
      return true;
    } else {
      return false;
    }
  } else {
    return true;
  }
}
 
IEngine *GetEngine() {
  // un-initialized default manager.
  static AllreduceBase default_manager;
  ThreadLocalEntry* e = EngineThreadLocal::Get();
  IEngine* ptr = e->engine.get();
  if (ptr == nullptr) {
    utils::Check(!e->initialized, "the rabit has not been initialized");
    return &default_manager;
  } else {
    return ptr;
  }
}
 
// perform in-place allgather, on sendrecvbuf
void Allgather(void *sendrecvbuf_, size_t total_size,
                   size_t slice_begin,
                   size_t slice_end,
                   size_t size_prev_slice) {
  GetEngine()->Allgather(sendrecvbuf_, total_size, slice_begin,
    slice_end, size_prev_slice);
}
 
 
// perform in-place allreduce, on sendrecvbuf
void Allreduce_(void *sendrecvbuf,  // NOLINT
                size_t type_nbytes,
                size_t count,
                IEngine::ReduceFunction red,
                mpi::DataType,
                mpi::OpType ,
                IEngine::PreprocFunction prepare_fun,
                void *prepare_arg) {
  GetEngine()->Allreduce(sendrecvbuf, type_nbytes, count, red, prepare_fun, prepare_arg);
}
}  // namespace engine
}  // namespace rabit