network.h

#ifndef LIGHTGBM_NETWORK_H_
#define LIGHTGBM_NETWORK_H_
 
#include <LightGBM/utils/log.h>
 
#include <LightGBM/meta.h>
#include <LightGBM/config.h>
 
#include <functional>
#include <vector>
#include <memory>
 
namespace LightGBM {
 
class Linkers;
 
class BruckMap {
public:
  int k;
  std::vector<int> in_ranks;
  std::vector<int> out_ranks;
  BruckMap();
  explicit BruckMap(int n);
  static BruckMap Construct(int rank, int num_machines);
};
 
enum RecursiveHalvingNodeType {
  Normal,  // normal node, 1 group only have 1 machine
  GroupLeader,  // leader of group when number of machines in this group is 2.
  Other  // non-leader machines in group
};
 
class RecursiveHalvingMap {
public:
  int k;
  RecursiveHalvingNodeType type;
  bool is_power_of_2;
  int neighbor;
  std::vector<int> ranks;
  std::vector<int> send_block_start;
  std::vector<int> send_block_len;
  std::vector<int> recv_block_start;
  std::vector<int> recv_block_len;
 
  RecursiveHalvingMap();
 
  RecursiveHalvingMap(int k, RecursiveHalvingNodeType _type, bool _is_power_of_2);
 
  static RecursiveHalvingMap Construct(int rank, int num_machines);
};
 
class Network {
public:
  static void Init(Config config);
  static void Init(int num_machines, int rank, ReduceScatterFunction reduce_scatter_ext_fun, AllgatherFunction allgather_ext_fun);
  static void Dispose();
  static inline int rank();
  static inline int num_machines();
 
  static void Allreduce(char* input, comm_size_t input_size, int type_size,
                        char* output, const ReduceFunction& reducer);
 
  static void AllreduceByAllGather(char* input, comm_size_t input_size, int type_size, char* output,
                                   const ReduceFunction& reducer);
 
  static void Allgather(char* input, comm_size_t send_size, char* output);
 
  static void Allgather(char* input, const comm_size_t* block_start, const comm_size_t* block_len, char* output, comm_size_t all_size);
 
  static void ReduceScatter(char* input, comm_size_t input_size, int type_size,
                            const comm_size_t* block_start, const comm_size_t* block_len, char* output, comm_size_t output_size,
                            const ReduceFunction& reducer);
 
  template<class T>
  static T GlobalSyncUpByMin(T& local) {
    T global = local;
    Allreduce(reinterpret_cast<char*>(&local),
              sizeof(local), sizeof(local),
              reinterpret_cast<char*>(&global),
              [] (const char* src, char* dst, int type_size, comm_size_t len) {
      comm_size_t used_size = 0;
      const T *p1;
      T *p2;
      while (used_size < len) {
        p1 = reinterpret_cast<const T *>(src);
        p2 = reinterpret_cast<T *>(dst);
        if (*p1 < *p2) {
          std::memcpy(dst, src, type_size);
        }
        src += type_size;
        dst += type_size;
        used_size += type_size;
      }
    });
    return global;
  }
 
  template<class T>
  static T GlobalSyncUpByMax(T& local) {
    T global = local;
    Allreduce(reinterpret_cast<char*>(&local),
              sizeof(local), sizeof(local),
              reinterpret_cast<char*>(&global),
              [] (const char* src, char* dst, int type_size, comm_size_t len) {
      comm_size_t used_size = 0;
      const T *p1;
      T *p2;
      while (used_size < len) {
        p1 = reinterpret_cast<const T *>(src);
        p2 = reinterpret_cast<T *>(dst);
        if (*p1 > *p2) {
          std::memcpy(dst, src, type_size);
        }
        src += type_size;
        dst += type_size;
        used_size += type_size;
      }
    });
    return global;
  }
 
  template<class T>
  static T GlobalSyncUpByMean(T& local) {
    T global = (T)0;
    Allreduce(reinterpret_cast<char*>(&local),
              sizeof(local), sizeof(local),
              reinterpret_cast<char*>(&global),
              [](const char* src, char* dst, int type_size, comm_size_t len) {
      comm_size_t used_size = 0;
      const T *p1;
      T *p2;
      while (used_size < len) {
        p1 = reinterpret_cast<const T *>(src);
        p2 = reinterpret_cast<T *>(dst);
        *p2 += *p1;
        src += type_size;
        dst += type_size;
        used_size += type_size;
      }
    });
    return static_cast<T>(global / num_machines_);
  }
 
  template<class T>
  static void GlobalSum(std::vector<T>& local) {
    std::vector<T> global(local.size(), 0);
    Allreduce(reinterpret_cast<char*>(local.data()),
              static_cast<comm_size_t>(sizeof(T) * local.size()), sizeof(T),
              reinterpret_cast<char*>(global.data()),
              [](const char* src, char* dst, int type_size, comm_size_t len) {
      comm_size_t used_size = 0;
      const T *p1;
      T *p2;
      while (used_size < len) {
        p1 = reinterpret_cast<const T *>(src);
        p2 = reinterpret_cast<T *>(dst);
        *p2 += *p1;
        src += type_size;
        dst += type_size;
        used_size += type_size;
      }
    });
    for (size_t i = 0; i < local.size(); ++i) {
      local[i] = global[i];
    }
  }
 
private:
  static void AllgatherBruck(char* input, const comm_size_t* block_start, const comm_size_t* block_len, char* output, comm_size_t all_size);
 
  static void AllgatherRecursiveDoubling(char* input, const comm_size_t* block_start, const comm_size_t* block_len, char* output, comm_size_t all_size);
 
  static void AllgatherRing(char* input, const comm_size_t* block_start, const comm_size_t* block_len, char* output, comm_size_t all_size);
 
  static void ReduceScatterRecursiveHalving(char* input, comm_size_t input_size, int type_size,
                                            const comm_size_t* block_start, const comm_size_t* block_len, char* output, comm_size_t output_size,
                                            const ReduceFunction& reducer);
 
  static void ReduceScatterRing(char* input, comm_size_t input_size, int type_size,
                                const comm_size_t* block_start, const comm_size_t* block_len, char* output, comm_size_t output_size,
                                const ReduceFunction& reducer);
 
  static THREAD_LOCAL int num_machines_;
  static THREAD_LOCAL int rank_;
  static THREAD_LOCAL std::unique_ptr<Linkers> linkers_;
  static THREAD_LOCAL BruckMap bruck_map_;
  static THREAD_LOCAL RecursiveHalvingMap recursive_halving_map_;
  static THREAD_LOCAL std::vector<comm_size_t> block_start_;
  static THREAD_LOCAL std::vector<comm_size_t> block_len_;
  static THREAD_LOCAL std::vector<char> buffer_;
  static THREAD_LOCAL comm_size_t buffer_size_;
  static THREAD_LOCAL ReduceScatterFunction reduce_scatter_ext_fun_;
  static THREAD_LOCAL AllgatherFunction allgather_ext_fun_;
};
 
inline int Network::rank() {
  return rank_;
}
 
inline int Network::num_machines() {
  return num_machines_;
}
 
}  // namespace LightGBM
 
#endif   // LightGBM_NETWORK_H_