parallel_tree_learner.h

#ifndef LIGHTGBM_TREELEARNER_PARALLEL_TREE_LEARNER_H_
#define LIGHTGBM_TREELEARNER_PARALLEL_TREE_LEARNER_H_
 
#include "serial_tree_learner.h"
#include "gpu_tree_learner.h"
#include <LightGBM/network.h>
 
#include <LightGBM/utils/array_args.h>
 
#include <cstring>
#include <vector>
#include <memory>
 
namespace LightGBM {
 
template <typename TREELEARNER_T>
class FeatureParallelTreeLearner: public TREELEARNER_T {
public:
  explicit FeatureParallelTreeLearner(const Config* config);
  ~FeatureParallelTreeLearner();
  void Init(const Dataset* train_data, bool is_constant_hessian) override;
 
protected:
  void BeforeTrain() override;
  void FindBestSplitsFromHistograms(const std::vector<int8_t>& is_feature_used, bool use_subtract) override;
private:
  int rank_;
  int num_machines_;
  std::vector<char> input_buffer_;
  std::vector<char> output_buffer_;
};
 
template <typename TREELEARNER_T>
class DataParallelTreeLearner: public TREELEARNER_T {
public:
  explicit DataParallelTreeLearner(const Config* config);
  ~DataParallelTreeLearner();
  void Init(const Dataset* train_data, bool is_constant_hessian) override;
  void ResetConfig(const Config* config) override;
 
protected:
  void BeforeTrain() override;
  void FindBestSplits() override;
  void FindBestSplitsFromHistograms(const std::vector<int8_t>& is_feature_used, bool use_subtract) override;
  void Split(Tree* tree, int best_Leaf, int* left_leaf, int* right_leaf) override;
 
  inline data_size_t GetGlobalDataCountInLeaf(int leaf_idx) const override {
    if (leaf_idx >= 0) {
      return global_data_count_in_leaf_[leaf_idx];
    } else {
      return 0;
    }
  }
 
private:
  int rank_;
  int num_machines_;
  std::vector<char> input_buffer_;
  std::vector<char> output_buffer_;
  std::vector<bool> is_feature_aggregated_;
  std::vector<comm_size_t> block_start_;
  std::vector<comm_size_t> block_len_;
  std::vector<comm_size_t> buffer_write_start_pos_;
  std::vector<comm_size_t> buffer_read_start_pos_;
  comm_size_t reduce_scatter_size_;
  std::vector<data_size_t> global_data_count_in_leaf_;
};
 
template <typename TREELEARNER_T>
class VotingParallelTreeLearner: public TREELEARNER_T {
public:
  explicit VotingParallelTreeLearner(const Config* config);
  ~VotingParallelTreeLearner() { }
  void Init(const Dataset* train_data, bool is_constant_hessian) override;
  void ResetConfig(const Config* config) override;
 
protected:
  void BeforeTrain() override;
  bool BeforeFindBestSplit(const Tree* tree, int left_leaf, int right_leaf) override;
  void FindBestSplits() override;
  void FindBestSplitsFromHistograms(const std::vector<int8_t>& is_feature_used, bool use_subtract) override;
  void Split(Tree* tree, int best_Leaf, int* left_leaf, int* right_leaf) override;
 
  inline data_size_t GetGlobalDataCountInLeaf(int leaf_idx) const override {
    if (leaf_idx >= 0) {
      return global_data_count_in_leaf_[leaf_idx];
    } else {
      return 0;
    }
  }
  void GlobalVoting(int leaf_idx, const std::vector<LightSplitInfo>& splits,
    std::vector<int>* out);
  void CopyLocalHistogram(const std::vector<int>& smaller_top_features,
    const std::vector<int>& larger_top_features);
 
private:
  Config local_config_;
  int top_k_;
  int rank_;
  int num_machines_;
  std::vector<char> input_buffer_;
  std::vector<char> output_buffer_;
  std::vector<bool> smaller_is_feature_aggregated_;
  std::vector<bool> larger_is_feature_aggregated_;
  std::vector<comm_size_t> block_start_;
  std::vector<comm_size_t> block_len_;
  std::vector<comm_size_t> smaller_buffer_read_start_pos_;
  std::vector<comm_size_t> larger_buffer_read_start_pos_;
  comm_size_t reduce_scatter_size_;
  std::vector<data_size_t> global_data_count_in_leaf_;
  std::unique_ptr<LeafSplits> smaller_leaf_splits_global_;
  std::unique_ptr<LeafSplits> larger_leaf_splits_global_;
  std::unique_ptr<FeatureHistogram[]> smaller_leaf_histogram_array_global_;
  std::unique_ptr<FeatureHistogram[]> larger_leaf_histogram_array_global_;
 
  std::vector<HistogramBinEntry> smaller_leaf_histogram_data_;
  std::vector<HistogramBinEntry> larger_leaf_histogram_data_;
  std::vector<FeatureMetainfo> feature_metas_;
};
 
// To-do: reduce the communication cost by using bitset to communicate.
inline void SyncUpGlobalBestSplit(char* input_buffer_, char* output_buffer_, SplitInfo* smaller_best_split, SplitInfo* larger_best_split, int max_cat_threshold) {
  // sync global best info
  int size = SplitInfo::Size(max_cat_threshold);
  smaller_best_split->CopyTo(input_buffer_);
  larger_best_split->CopyTo(input_buffer_ + size);
  Network::Allreduce(input_buffer_, size * 2, size, output_buffer_,
                     [] (const char* src, char* dst, int size, comm_size_t len) {
    comm_size_t used_size = 0;
    LightSplitInfo p1, p2;
    while (used_size < len) {
      p1.CopyFrom(src);
      p2.CopyFrom(dst);
      if (p1 > p2) {
        std::memcpy(dst, src, size);
      }
      src += size;
      dst += size;
      used_size += size;
    }
  });
  // copy back
  smaller_best_split->CopyFrom(output_buffer_);
  larger_best_split->CopyFrom(output_buffer_ + size);
}
 
}  // namespace LightGBM
#endif   // LightGBM_TREELEARNER_PARALLEL_TREE_LEARNER_H_