communicator.h

 
#pragma once
#include <xgboost/json.h>
#include <xgboost/logging.h>
 
#include <memory>
#include <string>
 
namespace xgboost {
namespace collective {
 
enum class DataType {
  kInt8 = 0,
  kUInt8 = 1,
  kInt32 = 2,
  kUInt32 = 3,
  kInt64 = 4,
  kUInt64 = 5,
  kFloat = 6,
  kDouble = 7
};
 
inline std::size_t GetTypeSize(DataType data_type) {
  std::size_t size{0};
  switch (data_type) {
    case DataType::kInt8:
      size = sizeof(std::int8_t);
      break;
    case DataType::kUInt8:
      size = sizeof(std::uint8_t);
      break;
    case DataType::kInt32:
      size = sizeof(std::int32_t);
      break;
    case DataType::kUInt32:
      size = sizeof(std::uint32_t);
      break;
    case DataType::kInt64:
      size = sizeof(std::int64_t);
      break;
    case DataType::kUInt64:
      size = sizeof(std::uint64_t);
      break;
    case DataType::kFloat:
      size = sizeof(float);
      break;
    case DataType::kDouble:
      size = sizeof(double);
      break;
    default:
      LOG(FATAL) << "Unknown data type.";
  }
  return size;
}
 
enum class Operation {
  kMax = 0,
  kMin = 1,
  kSum = 2,
  kBitwiseAND = 3,
  kBitwiseOR = 4,
  kBitwiseXOR = 5
};
 
class DeviceCommunicator;
 
enum class CommunicatorType { kUnknown, kRabit, kFederated, kInMemory, kInMemoryNccl };
 
inline int CompareStringsCaseInsensitive(const char *s1, const char *s2) {
#ifdef _MSC_VER
  return _stricmp(s1, s2);
#else   // _MSC_VER
  return strcasecmp(s1, s2);
#endif  // _MSC_VER
}
 
class Communicator {
 public:
  static void Init(Json const &config);
 
  static void Finalize();
 
  static Communicator *Get() { return communicator_.get(); }
 
#if defined(XGBOOST_USE_CUDA)
  static DeviceCommunicator *GetDevice(int device_ordinal);
#endif
 
  virtual ~Communicator() = default;
 
  int GetWorldSize() const { return world_size_; }
 
  int GetRank() const { return rank_; }
 
  virtual bool IsDistributed() const = 0;
 
  virtual bool IsFederated() const = 0;
 
  virtual void AllGather(void *send_receive_buffer, std::size_t size) = 0;
 
  virtual void AllReduce(void *send_receive_buffer, std::size_t count, DataType data_type,
                         Operation op) = 0;
 
  virtual void Broadcast(void *send_receive_buffer, std::size_t size, int root) = 0;
 
  virtual std::string GetProcessorName() = 0;
 
  virtual void Print(std::string const &message) = 0;
 
  static CommunicatorType GetTypeFromEnv() {
    auto *env = std::getenv("XGBOOST_COMMUNICATOR");
    if (env != nullptr) {
      return StringToType(env);
    } else {
      return CommunicatorType::kUnknown;
    }
  }
 
  static CommunicatorType GetTypeFromConfig(Json const &config) {
    auto const &j_upper = config["XGBOOST_COMMUNICATOR"];
    if (IsA<String const>(j_upper)) {
      return StringToType(get<String const>(j_upper).c_str());
    }
    auto const &j_lower = config["xgboost_communicator"];
    if (IsA<String const>(j_lower)) {
      return StringToType(get<String const>(j_lower).c_str());
    }
    return CommunicatorType::kUnknown;
  }
 
 protected:
  Communicator(int world_size, int rank) : world_size_(world_size), rank_(rank) {
    if (world_size < 1) {
      LOG(FATAL) << "World size " << world_size << " is less than 1.";
    }
    if (rank < 0) {
      LOG(FATAL) << "Rank " << rank << " is less than 0.";
    }
    if (rank >= world_size) {
      LOG(FATAL) << "Rank " << rank << " is greater than world_size - 1: " << world_size - 1 << ".";
    }
  }
 
  virtual void Shutdown() = 0;
 
 private:
  static CommunicatorType StringToType(char const *str) {
    CommunicatorType result = CommunicatorType::kUnknown;
    if (!CompareStringsCaseInsensitive("rabit", str)) {
      result = CommunicatorType::kRabit;
    } else if (!CompareStringsCaseInsensitive("federated", str)) {
      result = CommunicatorType::kFederated;
    } else if (!CompareStringsCaseInsensitive("in-memory", str)) {
      result = CommunicatorType::kInMemory;
    } else if (!CompareStringsCaseInsensitive("in-memory-nccl", str)) {
      result = CommunicatorType::kInMemoryNccl;
    } else {
      LOG(FATAL) << "Unknown communicator type " << str;
    }
    return result;
  }
 
  static thread_local std::unique_ptr<Communicator> communicator_;
  static thread_local CommunicatorType type_;
#if defined(XGBOOST_USE_CUDA)
  static thread_local std::unique_ptr<DeviceCommunicator> device_communicator_;
#endif
 
  int const world_size_;
  int const rank_;
};
 
}  // namespace collective
}  // namespace xgboost