communicator-inl.h

 
#pragma once
#include <string>
#include <vector>
 
#include "communicator.h"
 
namespace xgboost {
namespace collective {
 
inline void Init(Json const& config) {
  Communicator::Init(config);
}
 
inline void Finalize() { Communicator::Finalize(); }
 
inline int GetRank() { return Communicator::Get()->GetRank(); }
 
inline int GetWorldSize() { return Communicator::Get()->GetWorldSize(); }
 
inline bool IsDistributed() { return Communicator::Get()->IsDistributed(); }
 
inline bool IsFederated() { return Communicator::Get()->IsFederated(); }
 
inline void Print(char const *message) { Communicator::Get()->Print(message); }
 
inline void Print(std::string const &message) { Communicator::Get()->Print(message); }
 
inline std::string GetProcessorName() { return Communicator::Get()->GetProcessorName(); }
 
inline void Broadcast(void *send_receive_buffer, size_t size, int root) {
  Communicator::Get()->Broadcast(send_receive_buffer, size, root);
}
 
inline void Broadcast(std::string *sendrecv_data, int root) {
  size_t size = sendrecv_data->length();
  Broadcast(&size, sizeof(size), root);
  if (sendrecv_data->length() != size) {
    sendrecv_data->resize(size);
  }
  if (size != 0) {
    Broadcast(&(*sendrecv_data)[0], size * sizeof(char), root);
  }
}
 
inline void Allgather(void *send_receive_buffer, std::size_t size) {
  Communicator::Get()->AllGather(send_receive_buffer, size);
}
 
inline void Allreduce(void *send_receive_buffer, size_t count, int data_type, int op) {
  Communicator::Get()->AllReduce(send_receive_buffer, count, static_cast<DataType>(data_type),
                                 static_cast<Operation>(op));
}
 
inline void Allreduce(void *send_receive_buffer, size_t count, DataType data_type, Operation op) {
  Communicator::Get()->AllReduce(send_receive_buffer, count, data_type, op);
}
 
template <Operation op>
inline void Allreduce(int8_t *send_receive_buffer, size_t count) {
  Communicator::Get()->AllReduce(send_receive_buffer, count, DataType::kInt8, op);
}
 
template <Operation op>
inline void Allreduce(uint8_t *send_receive_buffer, size_t count) {
  Communicator::Get()->AllReduce(send_receive_buffer, count, DataType::kUInt8, op);
}
 
template <Operation op>
inline void Allreduce(int32_t *send_receive_buffer, size_t count) {
  Communicator::Get()->AllReduce(send_receive_buffer, count, DataType::kInt32, op);
}
 
template <Operation op>
inline void Allreduce(uint32_t *send_receive_buffer, size_t count) {
  Communicator::Get()->AllReduce(send_receive_buffer, count, DataType::kUInt32, op);
}
 
template <Operation op>
inline void Allreduce(int64_t *send_receive_buffer, size_t count) {
  Communicator::Get()->AllReduce(send_receive_buffer, count, DataType::kInt64, op);
}
 
template <Operation op>
inline void Allreduce(uint64_t *send_receive_buffer, size_t count) {
  Communicator::Get()->AllReduce(send_receive_buffer, count, DataType::kUInt64, op);
}
 
// Specialization for size_t, which is implementation defined, so it might or might not
// be one of uint64_t/uint32_t/unsigned long long/unsigned long.
template <Operation op, typename T,
          typename = std::enable_if_t<std::is_same<size_t, T>{} && !std::is_same<uint64_t, T>{}> >
inline void Allreduce(T *send_receive_buffer, size_t count) {
  static_assert(sizeof(T) == sizeof(uint64_t));
  Communicator::Get()->AllReduce(send_receive_buffer, count, DataType::kUInt64, op);
}
 
template <Operation op>
inline void Allreduce(float *send_receive_buffer, size_t count) {
  Communicator::Get()->AllReduce(send_receive_buffer, count, DataType::kFloat, op);
}
 
template <Operation op>
inline void Allreduce(double *send_receive_buffer, size_t count) {
  Communicator::Get()->AllReduce(send_receive_buffer, count, DataType::kDouble, op);
}
 
template <typename T>
struct AllgatherVResult {
  std::vector<std::size_t> offsets;
  std::vector<std::size_t> sizes;
  std::vector<T> result;
};
 
template <typename T>
inline AllgatherVResult<T> AllgatherV(std::vector<T> const &inputs,
                                      std::vector<std::size_t> const &sizes) {
  auto num_inputs = sizes.size();
 
  // Gather the sizes across all workers.
  std::vector<std::size_t> all_sizes(num_inputs * GetWorldSize());
  std::copy_n(sizes.cbegin(), sizes.size(), all_sizes.begin() + num_inputs * GetRank());
  collective::Allgather(all_sizes.data(), all_sizes.size() * sizeof(std::size_t));
 
  // Calculate input offsets (std::exclusive_scan).
  std::vector<std::size_t> offsets(all_sizes.size());
  for (std::size_t i = 1; i < offsets.size(); i++) {
    offsets[i] = offsets[i - 1] + all_sizes[i - 1];
  }
 
  // Gather all the inputs.
  auto total_input_size = offsets.back() + all_sizes.back();
  std::vector<T> all_inputs(total_input_size);
  std::copy_n(inputs.cbegin(), inputs.size(), all_inputs.begin() + offsets[num_inputs * GetRank()]);
  // We cannot use allgather here, since each worker might have a different size.
  Allreduce<Operation::kMax>(all_inputs.data(), all_inputs.size());
 
  return {offsets, all_sizes, all_inputs};
}
 
}  // namespace collective
}  // namespace xgboost