speed up multi-val bin subset for bagging (#2827)

* speed up multi-val bin subset for bagging

* remove the duplicated codes

* code refine

* some codes refactoring

* move `is_constant_hessian` into `TrainingShareStates`

* refine

* fix bug

* fix bug when num_groups_ < 0

* fix gpu

* fix gpu bagging

* fix gpu bug

* typo

* Update src/treelearner/serial_tree_learner.h

include/LightGBM/bin.h

@@ -303,7 +303,7 @@ class Bin {
   virtual void Push(int tid, data_size_t idx, uint32_t value) = 0;
 
 
-  virtual void CopySubset(const Bin* full_bin, const data_size_t* used_indices, data_size_t num_used_indices) = 0;
+  virtual void CopySubrow(const Bin* full_bin, const data_size_t* used_indices, data_size_t num_used_indices) = 0;
   /*!
   * \brief Get bin iterator of this bin for specific feature
   * \param min_bin min_bin of current used feature
@@ -453,23 +453,34 @@ class MultiValBin {
 
   virtual int32_t num_bin() const = 0;
 
-  virtual void ReSize(data_size_t num_data) = 0;
+  virtual double num_element_per_row() const = 0;
 
-  virtual void PushOneRow(int tid, data_size_t idx, const std::vector<uint32_t>& values) = 0;
 
-  virtual void CopySubset(const Bin* full_bin, const data_size_t* used_indices, data_size_t num_used_indices) = 0;
+  virtual void PushOneRow(int tid, data_size_t idx, const std::vector<uint32_t>& values) = 0;
 
-  virtual void ReSizeForSubFeature(int num_bin, int num_feature,
-                                   double estimate_element_per_row) = 0;
+  virtual void CopySubrow(const MultiValBin* full_bin,
+                          const data_size_t* used_indices,
+                          data_size_t num_used_indices) = 0;
 
-  virtual MultiValBin* CreateLike(int num_bin, int num_feature,
+  virtual MultiValBin* CreateLike(data_size_t num_data, int num_bin,
+                                  int num_feature,
                                   double estimate_element_per_row) const = 0;
 
-  virtual void CopySubFeature(const MultiValBin* full_bin,
-                              const std::vector<int>& used_feature_index,
-                              const std::vector<uint32_t>& lower,
-                              const std::vector<uint32_t>& upper,
-                              const std::vector<uint32_t>& delta) = 0;
+
+  virtual void CopySubcol(const MultiValBin* full_bin,
+                          const std::vector<int>& used_feature_index,
+                          const std::vector<uint32_t>& lower,
+                          const std::vector<uint32_t>& upper,
+                          const std::vector<uint32_t>& delta) = 0;
+
+  virtual void ReSize(data_size_t num_data, int num_bin, int num_feature,
+                      double estimate_element_per_row) = 0;
+
+  virtual void CopySubrowAndSubcol(
+      const MultiValBin* full_bin, const data_size_t* used_indices,
+      data_size_t num_used_indices, const std::vector<int>& used_feature_index,
+      const std::vector<uint32_t>& lower, const std::vector<uint32_t>& upper,
+      const std::vector<uint32_t>& delta) = 0;
 
   virtual void ConstructHistogram(
     const data_size_t* data_indices, data_size_t start, data_size_t end,

include/LightGBM/dataset.h

@@ -276,16 +276,22 @@ class Parser {
   static Parser* CreateParser(const char* filename, bool header, int num_features, int label_idx);
 };
 
-struct TrainingTempState {
-  std::vector<hist_t, Common::AlignmentAllocator<hist_t, kAlignedSize>>
-      hist_buf;
+struct TrainingShareStates {
+  bool is_colwise = true;
+  bool is_use_subcol = false;
+  bool is_use_subrow = false;
+  bool is_subrow_copied = false;
+  bool is_constant_hessian = true;
+  const data_size_t* bagging_use_indices;
+  data_size_t bagging_indices_cnt;
   int num_bin_aligned;
-  bool use_subfeature;
   std::unique_ptr<MultiValBin> multi_val_bin;
-  std::unique_ptr<MultiValBin> multi_val_bin_subfeature;
+  std::unique_ptr<MultiValBin> multi_val_bin_subset;
   std::vector<uint32_t> hist_move_src;
   std::vector<uint32_t> hist_move_dest;
   std::vector<uint32_t> hist_move_size;
+  std::vector<hist_t, Common::AlignmentAllocator<hist_t, kAlignedSize>>
+      hist_buf;
 
   void SetMultiValBin(MultiValBin* bin) {
     if (bin == nullptr) {
@@ -302,14 +308,14 @@ struct TrainingTempState {
   }
 
   hist_t* TempBuf() {
-    if (!use_subfeature) {
+    if (!is_use_subcol) {
       return nullptr;
     }
     return hist_buf.data() + hist_buf.size() - num_bin_aligned * 2;
   }
 
   void HistMove(const hist_t* src, hist_t* dest) {
-    if (!use_subfeature) {
+    if (!is_use_subcol) {
       return;
     }
 #pragma omp parallel for schedule(static)
@@ -436,16 +442,16 @@ class Dataset {
   }
   void ReSize(data_size_t num_data);
 
-  void CopySubset(const Dataset* fullset, const data_size_t* used_indices, data_size_t num_used_indices, bool need_meta_data);
+  void CopySubrow(const Dataset* fullset, const data_size_t* used_indices, data_size_t num_used_indices, bool need_meta_data);
 
   MultiValBin* GetMultiBinFromSparseFeatures() const;
 
   MultiValBin* GetMultiBinFromAllFeatures() const;
 
-  TrainingTempState* TestMultiThreadingMethod(
-    score_t* gradients, score_t* hessians,
-    const std::vector<int8_t>& is_feature_used, bool is_constant_hessian,
-    bool force_colwise, bool force_rowwise, bool* is_hist_col_wise) const;
+  TrainingShareStates* GetShareStates(
+      score_t* gradients, score_t* hessians,
+      const std::vector<int8_t>& is_feature_used, bool is_constant_hessian,
+      bool force_colwise, bool force_rowwise) const;
 
   LIGHTGBM_EXPORT void FinishLoad();
 
@@ -473,23 +479,21 @@ class Dataset {
   LIGHTGBM_EXPORT void CreateValid(const Dataset* dataset);
 
   void InitTrain(const std::vector<int8_t>& is_feature_used,
-                 bool is_colwise,
-                 TrainingTempState* temp_state) const;
+                 TrainingShareStates* share_state) const;
 
   void ConstructHistograms(const std::vector<int8_t>& is_feature_used,
                            const data_size_t* data_indices,
                            data_size_t num_data, const score_t* gradients,
                            const score_t* hessians, score_t* ordered_gradients,
-                           score_t* ordered_hessians, bool is_constant_hessian,
-                           bool is_colwise, TrainingTempState* temp_state,
+                           score_t* ordered_hessians,
+                           TrainingShareStates* share_state,
                            hist_t* histogram_data) const;
 
   void ConstructHistogramsMultiVal(const data_size_t* data_indices,
                                    data_size_t num_data,
                                    const score_t* gradients,
                                    const score_t* hessians,
-                                   bool is_constant_hessian,
-                                   TrainingTempState* temp_state,
+                                   TrainingShareStates* share_state,
                                    hist_t* histogram_data) const;
 
   void FixHistogram(int feature_idx, double sum_gradient, double sum_hessian, hist_t* data) const;

include/LightGBM/feature_group.h

@@ -174,12 +174,12 @@ class FeatureGroup {
     }
   }
 
-  inline void CopySubset(const FeatureGroup* full_feature, const data_size_t* used_indices, data_size_t num_used_indices) {
+  inline void CopySubrow(const FeatureGroup* full_feature, const data_size_t* used_indices, data_size_t num_used_indices) {
     if (!is_multi_val_) {
-      bin_data_->CopySubset(full_feature->bin_data_.get(), used_indices, num_used_indices);
+      bin_data_->CopySubrow(full_feature->bin_data_.get(), used_indices, num_used_indices);
     } else {
       for (int i = 0; i < num_feature_; ++i) {
-        multi_bin_data_[i]->CopySubset(full_feature->multi_bin_data_[i].get(), used_indices, num_used_indices);
+        multi_bin_data_[i]->CopySubrow(full_feature->multi_bin_data_[i].get(), used_indices, num_used_indices);
       }
     }
   }

include/LightGBM/tree_learner.h

@@ -37,7 +37,10 @@ class TreeLearner {
   */
   virtual void Init(const Dataset* train_data, bool is_constant_hessian) = 0;
 
-  virtual void ResetTrainingData(const Dataset* train_data) = 0;
+  virtual void ResetIsConstantHessian(bool is_constant_hessian) = 0;
+
+  virtual void ResetTrainingData(const Dataset* train_data,
+                                 bool is_constant_hessian) = 0;
 
   /*!
   * \brief Reset tree configs
@@ -52,7 +55,7 @@ class TreeLearner {
   * \param is_constant_hessian True if all hessians share the same value
   * \return A trained tree
   */
-  virtual Tree* Train(const score_t* gradients, const score_t* hessians, bool is_constant_hessian,
+  virtual Tree* Train(const score_t* gradients, const score_t* hessians,
                       const Json& forced_split_json) = 0;
 
   /*!
@@ -65,13 +68,13 @@ class TreeLearner {
 
   /*!
   * \brief Set bagging data
+  * \param subset subset of bagging
   * \param used_indices Used data indices
   * \param num_data Number of used data
   */
-  virtual void SetBaggingData(const data_size_t* used_indices,
-    data_size_t num_data) = 0;
-
-  virtual bool IsHistColWise() const = 0;
+  virtual void SetBaggingData(const Dataset* subset,
+                              const data_size_t* used_indices,
+                              data_size_t num_data) = 0;
 
   /*!
   * \brief Using last trained tree to predict score then adding to out_score;

src/boosting/gbdt.cpp

@@ -231,13 +231,14 @@ void GBDT::Bagging(int iter) {
     Log::Debug("Re-bagging, using %d data to train", bag_data_cnt_);
     // set bagging data to tree learner
     if (!is_use_subset_) {
-      tree_learner_->SetBaggingData(bag_data_indices_.data(), bag_data_cnt_);
+      tree_learner_->SetBaggingData(nullptr, bag_data_indices_.data(), bag_data_cnt_);
     } else {
       // get subset
       tmp_subset_->ReSize(bag_data_cnt_);
-      tmp_subset_->CopySubset(train_data_, bag_data_indices_.data(),
+      tmp_subset_->CopySubrow(train_data_, bag_data_indices_.data(),
                               bag_data_cnt_, false);
-      tree_learner_->ResetTrainingData(tmp_subset_.get());
+      tree_learner_->SetBaggingData(tmp_subset_.get(), bag_data_indices_.data(),
+                                    bag_data_cnt_);
     }
   }
 }
@@ -365,7 +366,7 @@ bool GBDT::TrainOneIter(const score_t* gradients, const score_t* hessians) {
         grad = gradients_.data() + offset;
         hess = hessians_.data() + offset;
       }
-      new_tree.reset(tree_learner_->Train(grad, hess, is_constant_hessian_, forced_splits_json_));
+      new_tree.reset(tree_learner_->Train(grad, hess, forced_splits_json_));
     }
 
     if (new_tree->num_leaves() > 1) {
@@ -693,8 +694,10 @@ void GBDT::ResetTrainingData(const Dataset* train_data, const ObjectiveFunction*
     feature_names_ = train_data_->feature_names();
     feature_infos_ = train_data_->feature_infos();
 
-    tree_learner_->ResetTrainingData(train_data);
+    tree_learner_->ResetTrainingData(train_data, is_constant_hessian_);
     ResetBaggingConfig(config_.get(), true);
+  } else {
+    tree_learner_->ResetIsConstantHessian(is_constant_hessian_);
   }
 }
 
@@ -750,7 +753,7 @@ void GBDT::ResetBaggingConfig(const Config* config, bool is_change_dataset) {
         (static_cast<double>(bag_data_cnt_) / num_data_) / config->bagging_freq;
     is_use_subset_ = false;
     const int group_threshold_usesubset = 100;
-    if (tree_learner_->IsHistColWise() && average_bag_rate <= 0.5
+    if (average_bag_rate <= 0.5
         && (train_data_->num_feature_groups() < group_threshold_usesubset)) {
       if (tmp_subset_ == nullptr || is_change_dataset) {
         tmp_subset_.reset(new Dataset(bag_data_cnt_));

src/boosting/rf.hpp

@@ -125,8 +125,7 @@ class RF : public GBDT {
           hess = tmp_hess_.data();
         }
 
-        new_tree.reset(tree_learner_->Train(grad, hess, is_constant_hessian_,
-          forced_splits_json_));
+        new_tree.reset(tree_learner_->Train(grad, hess, forced_splits_json_));
       }
 
       if (new_tree->num_leaves() > 1) {

src/c_api.cpp

@@ -1059,7 +1059,7 @@ int LGBM_DatasetGetSubset(
   }
   auto ret = std::unique_ptr<Dataset>(new Dataset(num_used_row_indices));
   ret->CopyFeatureMapperFrom(full_dataset);
-  ret->CopySubset(full_dataset, used_row_indices, num_used_row_indices, true);
+  ret->CopySubrow(full_dataset, used_row_indices, num_used_row_indices, true);
   *out = ret.release();
   API_END();
 }