speed up sub-feature in row-wise parallelism (#2764)

* commit

* refactoring

* Update src/io/bin.cpp

* Apply suggestions from code review

* bug

* code clean

* remove warning

* commit

* update parameter

docs/Parameters.rst

@@ -202,8 +202,6 @@ Learning Control Parameters
 
       -  ``num_threads`` is large, e.g. ``>20``
 
-      -  you want to use small ``feature_fraction`` (e.g. ``0.5``) to speed up
-
       -  you want to reduce memory cost
 
    -  **Note**: when both ``force_col_wise`` and ``force_row_wise`` are ``false``, LightGBM will firstly try them both, and then use the faster one. To remove the overhead of testing set the faster one to ``true`` manually

include/LightGBM/bin.h

@@ -458,6 +458,18 @@ class MultiValBin {
 
   virtual void CopySubset(const Bin* full_bin, const data_size_t* used_indices, data_size_t num_used_indices) = 0;
 
+  virtual void ReSizeForSubFeature(int num_bin, int num_feature,
+                                   double estimate_element_per_row) = 0;
+
+  virtual MultiValBin* CreateLike(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 ConstructHistogram(
     const data_size_t* data_indices, data_size_t start, data_size_t end,
     const score_t* gradients, const score_t* hessians,
@@ -477,7 +489,13 @@ class MultiValBin {
 
   virtual bool IsSparse() = 0;
 
-  static MultiValBin* CreateMultiValBin(data_size_t num_data, int num_bin, int num_feature, double sparse_rate);
+  static MultiValBin* CreateMultiValBin(data_size_t num_data, int num_bin,
+                                        int num_feature, double sparse_rate);
+
+  static MultiValBin* CreateMultiValDenseBin(data_size_t num_data, int num_bin,
+                                             int num_feature);
+
+  static MultiValBin* CreateMultiValSparseBin(data_size_t num_data, int num_bin, double estimate_element_per_row);
 
   virtual MultiValBin* Clone() = 0;
 };

include/LightGBM/config.h

@@ -219,7 +219,6 @@ struct Config {
   // desc = enabling this is recommended when:
   // descl2 = the number of columns is large, or the total number of bins is large
   // descl2 = ``num_threads`` is large, e.g. ``>20``
-  // descl2 = you want to use small ``feature_fraction`` (e.g. ``0.5``) to speed up
   // descl2 = you want to reduce memory cost
   // desc = **Note**: when both ``force_col_wise`` and ``force_row_wise`` are ``false``, LightGBM will firstly try them both, and then use the faster one. To remove the overhead of testing set the faster one to ``true`` manually
   // desc = **Note**: this parameter cannot be used at the same time with ``force_row_wise``, choose only one of them

include/LightGBM/dataset.h

@@ -277,6 +277,53 @@ 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;
+  int num_bin_aligned;
+  bool use_subfeature;
+  std::unique_ptr<MultiValBin> multi_val_bin;
+  std::unique_ptr<MultiValBin> multi_val_bin_subfeature;
+  std::vector<uint32_t> hist_move_src;
+  std::vector<uint32_t> hist_move_dest;
+  std::vector<uint32_t> hist_move_size;
+
+  void SetMultiValBin(MultiValBin* bin) {
+    if (bin == nullptr) {
+      return;
+    }
+    multi_val_bin.reset(bin);
+    int num_threads = 1;
+#pragma omp parallel
+#pragma omp master
+    { num_threads = omp_get_num_threads(); }
+   num_bin_aligned =
+        (bin->num_bin() + kAlignedSize - 1) / kAlignedSize * kAlignedSize;
+    size_t new_size = static_cast<size_t>(num_bin_aligned) * 2 * num_threads;
+    if (new_size > hist_buf.size()) {
+      hist_buf.resize(static_cast<size_t>(num_bin_aligned) * 2 * num_threads);
+    }
+  }
+
+  hist_t* TempBuf() {
+    if (!use_subfeature) {
+      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) {
+      return;
+    }
+#pragma omp parallel for schedule(static)
+    for (int i = 0; i < static_cast<int>(hist_move_src.size()); ++i) {
+      std::copy_n(src + hist_move_src[i], hist_move_size[i],
+                  dest + hist_move_dest[i]);
+    }
+  }
+};
+
 /*! \brief The main class of data set,
 *          which are used to training or validation
 */
@@ -399,7 +446,9 @@ class Dataset {
 
   MultiValBin* GetMultiBinFromAllFeatures() const;
 
-  MultiValBin* TestMultiThreadingMethod(score_t* gradients, score_t* hessians, const std::vector<int8_t>& is_feature_used, bool is_constant_hessian,
+  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;
 
   LIGHTGBM_EXPORT void FinishLoad();
@@ -429,17 +478,24 @@ 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;
+
   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,
-                           const MultiValBin* multi_val_bin, bool is_colwise,
+                           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,
                            hist_t* histogram_data) const;
 
-  void ConstructHistogramsMultiVal(const MultiValBin* multi_val_bin, const data_size_t* data_indices, data_size_t num_data,
-                                   const score_t* gradients, const score_t* hessians,
+  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,
                                    hist_t* histogram_data) const;
 
   void FixHistogram(int feature_idx, double sum_gradient, double sum_hessian, hist_t* data) const;
@@ -654,7 +710,6 @@ class Dataset {
   bool use_missing_;
   bool zero_as_missing_;
   std::vector<int> feature_need_push_zeros_;
-  mutable std::vector<hist_t, Common::AlignmentAllocator<hist_t, kAlignedSize>> hist_buf_;
 };
 
 }  // namespace LightGBM

src/io/bin.cpp

@@ -666,21 +666,38 @@ namespace LightGBM {
   MultiValBin* MultiValBin::CreateMultiValBin(data_size_t num_data, int num_bin, int num_feature, double sparse_rate) {
     const double multi_val_bin_sparse_threshold = 0.25f;
     if (sparse_rate >= multi_val_bin_sparse_threshold) {
-      if (num_bin <= 256) {
-        return new MultiValSparseBin<uint8_t>(num_data, num_bin);
-      } else if (num_bin <= 65536) {
-        return new MultiValSparseBin<uint16_t>(num_data, num_bin);
-      } else {
-        return new MultiValSparseBin<uint32_t>(num_data, num_bin);
-      }
+      const double average_element_per_row = (1.0 - sparse_rate) * num_feature;
+      return CreateMultiValSparseBin(num_data, num_bin,
+                                     average_element_per_row);
     } else {
-      if (num_bin <= 256) {
-        return new MultiValDenseBin<uint8_t>(num_data, num_bin, num_feature);
-      } else if (num_bin <= 65536) {
-        return new MultiValDenseBin<uint16_t>(num_data, num_bin, num_feature);
-      } else {
-        return new MultiValDenseBin<uint32_t>(num_data, num_bin, num_feature);
-      }
+      return CreateMultiValDenseBin(num_data, num_bin, num_feature);
+    }
+  }
+
+  MultiValBin* MultiValBin::CreateMultiValDenseBin(data_size_t num_data,
+                                                   int num_bin,
+                                                   int num_feature) {
+    if (num_bin <= 256) {
+      return new MultiValDenseBin<uint8_t>(num_data, num_bin, num_feature);
+    } else if (num_bin <= 65536) {
+      return new MultiValDenseBin<uint16_t>(num_data, num_bin, num_feature);
+    } else {
+      return new MultiValDenseBin<uint32_t>(num_data, num_bin, num_feature);
+    }
+  }
+
+  MultiValBin* MultiValBin::CreateMultiValSparseBin(data_size_t num_data,
+                                                    int num_bin,
+                                                    double estimate_element_per_row) {
+    if (num_bin <= 256) {
+      return new MultiValSparseBin<uint8_t>(num_data, num_bin,
+                                            estimate_element_per_row);
+    } else if (num_bin <= 65536) {
+      return new MultiValSparseBin<uint16_t>(num_data, num_bin,
+                                             estimate_element_per_row);
+    } else {
+      return new MultiValSparseBin<uint32_t>(num_data, num_bin,
+                                             estimate_element_per_row);
     }
   }
 

src/io/dataset.cpp

@@ -482,14 +482,21 @@ void Dataset::FinishLoad() {
   is_finish_load_ = true;
 }
 
-void PushDataToMultiValBin(int num_threads, data_size_t num_data, const std::vector<uint32_t> most_freq_bins,
-  const std::vector<uint32_t> offsets, std::vector<std::vector<std::unique_ptr<BinIterator>>>& iters, MultiValBin* ret) {
-  Common::FunctionTimer fun_time("Dataset::PushDataToMultiValBin", global_timer);
+
+void PushDataToMultiValBin(
+    int num_threads, data_size_t num_data,
+    const std::vector<uint32_t> most_freq_bins,
+    const std::vector<uint32_t> offsets,
+    std::vector<std::vector<std::unique_ptr<BinIterator>>>& iters,
+    MultiValBin* ret) {
+  Common::FunctionTimer fun_time("Dataset::PushDataToMultiValBin",
+                                 global_timer);
   const data_size_t min_block_size = 4096;
-  const int n_block = std::min(num_threads, (num_data + min_block_size - 1) / min_block_size);
+  const int n_block =
+      std::min(num_threads, (num_data + min_block_size - 1) / min_block_size);
   const data_size_t block_size = (num_data + n_block - 1) / n_block;
   if (ret->IsSparse()) {
-    #pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static)
     for (int tid = 0; tid < n_block; ++tid) {
       std::vector<uint32_t> cur_data;
       cur_data.reserve(most_freq_bins.size());
@@ -515,7 +522,7 @@ void PushDataToMultiValBin(int num_threads, data_size_t num_data, const std::vec
       }
     }
   } else {
-    #pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static)
     for (int tid = 0; tid < n_block; ++tid) {
       std::vector<uint32_t> cur_data(most_freq_bins.size(), 0);
       data_size_t start = tid * block_size;
@@ -543,7 +550,8 @@ void PushDataToMultiValBin(int num_threads, data_size_t num_data, const std::vec
 }
 
 MultiValBin* Dataset::GetMultiBinFromSparseFeatures() const {
-  Common::FunctionTimer fun_time("Dataset::GetMultiBinFromSparseFeatures", global_timer);
+  Common::FunctionTimer fun_time("Dataset::GetMultiBinFromSparseFeatures",
+                                 global_timer);
   int multi_group_id = -1;
   for (int i = 0; i < num_groups_; ++i) {
     if (feature_groups_[i]->is_multi_val_) {
@@ -560,39 +568,42 @@ MultiValBin* Dataset::GetMultiBinFromSparseFeatures() const {
   const auto& offsets = feature_groups_[multi_group_id]->bin_offsets_;
   const int num_feature = feature_groups_[multi_group_id]->num_feature_;
   int num_threads = 1;
-  #pragma omp parallel
-  #pragma omp master
-  {
-    num_threads = omp_get_num_threads();
-  }
+#pragma omp parallel
+#pragma omp master
+  { num_threads = omp_get_num_threads(); }
 
   std::vector<std::vector<std::unique_ptr<BinIterator>>> iters(num_threads);
   std::vector<uint32_t> most_freq_bins;
   double sum_sparse_rate = 0;
   for (int i = 0; i < num_feature; ++i) {
     for (int tid = 0; tid < num_threads; ++tid) {
-      iters[tid].emplace_back(feature_groups_[multi_group_id]->SubFeatureIterator(i));
+      iters[tid].emplace_back(
+          feature_groups_[multi_group_id]->SubFeatureIterator(i));
     }
-    most_freq_bins.push_back(feature_groups_[multi_group_id]->bin_mappers_[i]->GetMostFreqBin());
-    sum_sparse_rate += feature_groups_[multi_group_id]->bin_mappers_[i]->sparse_rate();
+    most_freq_bins.push_back(
+        feature_groups_[multi_group_id]->bin_mappers_[i]->GetMostFreqBin());
+    sum_sparse_rate +=
+        feature_groups_[multi_group_id]->bin_mappers_[i]->sparse_rate();
   }
   sum_sparse_rate /= num_feature;
-  Log::Debug("Dataset::GetMultiBinFromSparseFeatures: sparse rate %f", sum_sparse_rate);
+  Log::Debug("Dataset::GetMultiBinFromSparseFeatures: sparse rate %f",
+             sum_sparse_rate);
   std::unique_ptr<MultiValBin> ret;
-  ret.reset(MultiValBin::CreateMultiValBin(num_data_, offsets.back(), num_feature, sum_sparse_rate));
-  PushDataToMultiValBin(num_threads, num_data_, most_freq_bins, offsets, iters, ret.get());
+  ret.reset(MultiValBin::CreateMultiValBin(num_data_, offsets.back(),
+                                           num_feature, sum_sparse_rate));
+  PushDataToMultiValBin(num_threads, num_data_, most_freq_bins, offsets, iters,
+                        ret.get());
   ret->FinishLoad();
   return ret.release();
 }
 
 MultiValBin* Dataset::GetMultiBinFromAllFeatures() const {
-  Common::FunctionTimer fun_time("Dataset::GetMultiBinFromAllFeatures", global_timer);
+  Common::FunctionTimer fun_time("Dataset::GetMultiBinFromAllFeatures",
+                                 global_timer);
   int num_threads = 1;
-  #pragma omp parallel
-  #pragma omp master
-  {
-    num_threads = omp_get_num_threads();
-  }
+#pragma omp parallel
+#pragma omp master
+  { num_threads = omp_get_num_threads(); }
   double sum_dense_ratio = 0;
 
   std::unique_ptr<MultiValBin> ret;
@@ -613,7 +624,8 @@ MultiValBin* Dataset::GetMultiBinFromAllFeatures() const {
         }
         offsets.push_back(num_total_bin);
         for (int tid = 0; tid < num_threads; ++tid) {
-          iters[tid].emplace_back(feature_groups_[gid]->SubFeatureIterator(fid));
+          iters[tid].emplace_back(
+              feature_groups_[gid]->SubFeatureIterator(fid));
         }
       }
     } else {
@@ -630,86 +642,105 @@ MultiValBin* Dataset::GetMultiBinFromAllFeatures() const {
     }
   }
   sum_dense_ratio /= static_cast<double>(most_freq_bins.size());
-  Log::Debug("Dataset::GetMultiBinFromAllFeatures: sparse rate %f", 1.0 - sum_dense_ratio);
-  ret.reset(MultiValBin::CreateMultiValBin(num_data_, num_total_bin, static_cast<int>(most_freq_bins.size()), 1.0 - sum_dense_ratio));
-  PushDataToMultiValBin(num_threads, num_data_, most_freq_bins, offsets, iters, ret.get());
+  Log::Debug("Dataset::GetMultiBinFromAllFeatures: sparse rate %f",
+             1.0 - sum_dense_ratio);
+  ret.reset(MultiValBin::CreateMultiValBin(
+      num_data_, num_total_bin, static_cast<int>(most_freq_bins.size()),
+      1.0 - sum_dense_ratio));
+  PushDataToMultiValBin(num_threads, num_data_, most_freq_bins, offsets, iters,
+                        ret.get());
   ret->FinishLoad();
   return ret.release();
 }
 
-MultiValBin* Dataset::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 {
-  int num_threads = 1;
-  #pragma omp parallel
-  #pragma omp master
-  {
-    num_threads = omp_get_num_threads();
-  }
-  Common::FunctionTimer fun_timer("Dataset::TestMultiThreadingMethod", global_timer);
+TrainingTempState* Dataset::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 {
+  Common::FunctionTimer fun_timer("Dataset::TestMultiThreadingMethod",
+                                  global_timer);
   if (force_colwise && force_rowwise) {
-    Log::Fatal("Cannot set both `force_col_wise` and `force_row_wise` to `true` at the same time");
+    Log::Fatal(
+        "Cannot set both `force_col_wise` and `force_row_wise` to `true` at "
+        "the same time");
   }
   if (num_groups_ <= 0) {
     return nullptr;
   }
   if (force_colwise) {
     *is_hist_col_wise = true;
-    return GetMultiBinFromSparseFeatures();
+    TrainingTempState* temp_state = new TrainingTempState();
+    temp_state->SetMultiValBin(GetMultiBinFromSparseFeatures());
+    return temp_state;
   } else if (force_rowwise) {
     *is_hist_col_wise = false;
-    auto ret = GetMultiBinFromAllFeatures();
-    const int num_bin_aligned = (ret->num_bin() + kAlignedSize - 1) / kAlignedSize * kAlignedSize;
-    hist_buf_.resize(static_cast<size_t>(num_bin_aligned) * 2 * num_threads);
-    return ret;
+    TrainingTempState* temp_state = new TrainingTempState();
+    temp_state->SetMultiValBin(GetMultiBinFromAllFeatures());
+    return temp_state;
   } else {
     std::unique_ptr<MultiValBin> sparse_bin;
     std::unique_ptr<MultiValBin> all_bin;
-    std::chrono::duration<double, std::milli> col_wise_init_time, row_wise_init_time;
+    std::unique_ptr<TrainingTempState> colwise_state;
+    std::unique_ptr<TrainingTempState> rowwise_state;
+    colwise_state.reset(new TrainingTempState());
+    rowwise_state.reset(new TrainingTempState());
+
+    std::chrono::duration<double, std::milli> col_wise_init_time,
+        row_wise_init_time;
     auto start_time = std::chrono::steady_clock::now();
-    sparse_bin.reset(GetMultiBinFromSparseFeatures());
+    colwise_state->SetMultiValBin(GetMultiBinFromSparseFeatures());
     col_wise_init_time = std::chrono::steady_clock::now() - start_time;
+
     start_time = std::chrono::steady_clock::now();
-    all_bin.reset(GetMultiBinFromAllFeatures());
-    std::vector<hist_t, Common::AlignmentAllocator<hist_t, kAlignedSize>> hist_data(NumTotalBin() * 2);
-    const int num_bin_aligned = (all_bin->num_bin() + kAlignedSize - 1) / kAlignedSize * kAlignedSize;
-    hist_buf_.resize(static_cast<size_t>(num_bin_aligned) * 2 * num_threads);
+    rowwise_state->SetMultiValBin(GetMultiBinFromAllFeatures());
+    std::vector<hist_t, Common::AlignmentAllocator<hist_t, kAlignedSize>>
+        hist_data(NumTotalBin() * 2);
+
     row_wise_init_time = std::chrono::steady_clock::now() - start_time;
-    Log::Debug("init for col-wise cost %f seconds, init for row-wise cost %f seconds",
-               col_wise_init_time * 1e-3, row_wise_init_time * 1e-3);
+    Log::Debug(
+        "init for col-wise cost %f seconds, init for row-wise cost %f seconds",
+        col_wise_init_time * 1e-3, row_wise_init_time * 1e-3);
+    InitTrain(is_feature_used, true, colwise_state.get());
+    InitTrain(is_feature_used, false, rowwise_state.get());
     std::chrono::duration<double, std::milli> col_wise_time, row_wise_time;
     start_time = std::chrono::steady_clock::now();
-    ConstructHistograms(is_feature_used, nullptr, num_data_, gradients, hessians, gradients, hessians, is_constant_hessian, sparse_bin.get(), true, hist_data.data());
+    ConstructHistograms(is_feature_used, nullptr, num_data_, gradients,
+                        hessians, gradients, hessians, is_constant_hessian,
+                        true, colwise_state.get(), hist_data.data());
     col_wise_time = std::chrono::steady_clock::now() - start_time;
     start_time = std::chrono::steady_clock::now();
-    ConstructHistogramsMultiVal(all_bin.get(), nullptr, num_data_, gradients, hessians, is_constant_hessian, hist_data.data());
+    ConstructHistogramsMultiVal(nullptr, num_data_, gradients, hessians,
+                                is_constant_hessian, rowwise_state.get(),
+                                hist_data.data());
     row_wise_time = std::chrono::steady_clock::now() - start_time;
     Log::Debug("col-wise cost %f seconds, row-wise cost %f seconds",
                col_wise_time * 1e-3, row_wise_time * 1e-3);
     if (col_wise_time < row_wise_time) {
       *is_hist_col_wise = true;
-      hist_buf_.clear();
       auto overhead_cost = row_wise_init_time + row_wise_time + col_wise_time;
       Log::Warning(
-          "Auto-choosing col-wise multi-threading, the overhead of testing was %f "
+          "Auto-choosing col-wise multi-threading, the overhead of testing was "
+          "%f "
           "seconds.\nYou can set `force_col_wise=true` to remove the "
           "overhead.",
           overhead_cost * 1e-3);
-      return sparse_bin.release();
+      return colwise_state.release();
     } else {
       *is_hist_col_wise = false;
       auto overhead_cost = col_wise_init_time + row_wise_time + col_wise_time;
       Log::Warning(
-          "Auto-choosing row-wise multi-threading, the overhead of testing was %f "
+          "Auto-choosing row-wise multi-threading, the overhead of testing was "
+          "%f "
           "seconds.\nYou can set `force_row_wise=true` to remove the "
           "overhead.\nAnd if memory is not enough, you can set "
           "`force_col_wise=true`.",
           overhead_cost * 1e-3);
-      if (all_bin->IsSparse()) {
+      if (rowwise_state->multi_val_bin->IsSparse()) {
         Log::Debug("Using Sparse Multi-Val Bin");
       } else {
         Log::Debug("Using Dense Multi-Val Bin");
       }
-      return all_bin.release();
+      return rowwise_state.release();
     }
   }
 }
@@ -1099,51 +1130,195 @@ void Dataset::DumpTextFile(const char* text_filename) {
   fclose(file);
 }
 
-void Dataset::ConstructHistogramsMultiVal(const MultiValBin* multi_val_bin, const data_size_t* data_indices, data_size_t num_data,
-                                          const score_t* gradients, const score_t* hessians,
-                                          bool is_constant_hessian,
-                                          hist_t* hist_data) const {
-  Common::FunctionTimer fun_time("Dataset::ConstructHistogramsMultiVal", global_timer);
+
+void Dataset::InitTrain(const std::vector<int8_t>& is_feature_used,
+                        bool is_colwise, TrainingTempState* temp_state) const {
+  Common::FunctionTimer fun_time("Dataset::InitTrain", global_timer);
+  temp_state->use_subfeature = false;
+  if (temp_state->multi_val_bin == nullptr) {
+    return;
+  }
+  global_timer.Start("Dataset::InitTrain.Prep");
+  double sum_used_dense_ratio = 0.0;
+  double sum_dense_ratio = 0.0;
+  int num_used = 0;
+  int total = 0;
+  std::vector<int> used_feature_index;
+  for (int i = 0; i < num_groups_; ++i) {
+    int f_start = group_feature_start_[i];
+    if (feature_groups_[i]->is_multi_val_ ) {
+      for (int j = 0; j < feature_groups_[i]->num_feature_; ++j) {
+        const auto dense_rate =
+            1.0 - feature_groups_[i]->bin_mappers_[j]->sparse_rate();
+        if (is_feature_used[f_start + j]) {
+          ++num_used;
+          used_feature_index.push_back(total);
+          sum_used_dense_ratio += dense_rate;
+        }
+        sum_dense_ratio += dense_rate;
+        ++total;
+      }
+    } else if (!is_colwise) {
+      bool is_group_used = false;
+      double dense_rate = 0;
+      for (int j = 0; j < feature_groups_[i]->num_feature_; ++j) {
+        if (is_feature_used[f_start + j]) {
+          is_group_used = true;
+        }
+        dense_rate += 1.0 - feature_groups_[i]->bin_mappers_[j]->sparse_rate();
+      }
+      if (is_group_used) {
+        ++num_used;
+        used_feature_index.push_back(total);
+        sum_used_dense_ratio += dense_rate;
+      }
+      sum_dense_ratio += dense_rate;
+      ++total;
+    }
+  }
+  global_timer.Stop("Dataset::InitTrain.Prep");
+  const double k_subfeature_threshold = 0.6;
+  if (sum_used_dense_ratio >= sum_dense_ratio * k_subfeature_threshold) {
+    return;
+  }
+  temp_state->use_subfeature = true;
+  global_timer.Start("Dataset::InitTrain.Prep");
+  std::vector<uint32_t> upper_bound;
+  std::vector<uint32_t> lower_bound;
+  std::vector<uint32_t> delta;
+  temp_state->hist_move_src.clear();
+  temp_state->hist_move_dest.clear();
+  temp_state->hist_move_size.clear();
+
+  int num_total_bin = 1;
+  int new_num_total_bin = 1;
+
+  for (int i = 0; i < num_groups_; ++i) {
+    int f_start = group_feature_start_[i];
+    if (feature_groups_[i]->is_multi_val_) {
+      for (int j = 0; j < feature_groups_[i]->num_feature_; ++j) {
+        const auto& bin_mapper = feature_groups_[i]->bin_mappers_[j];
+        int cur_num_bin = bin_mapper->num_bin();
+        if (bin_mapper->GetMostFreqBin() == 0) {
+          cur_num_bin -= 1;
+        }
+        num_total_bin += cur_num_bin;
+        if (is_feature_used[f_start + j]) {
+          new_num_total_bin += cur_num_bin;
+
+          lower_bound.push_back(num_total_bin - cur_num_bin);
+          upper_bound.push_back(num_total_bin);
+          
+          temp_state->hist_move_src.push_back(
+              (new_num_total_bin - cur_num_bin) * 2);
+          temp_state->hist_move_dest.push_back((num_total_bin - cur_num_bin) *
+                                               2);
+          temp_state->hist_move_size.push_back(cur_num_bin * 2);
+          delta.push_back(num_total_bin - new_num_total_bin);
+        }
+      }
+    } else if (!is_colwise) {
+      bool is_group_used = false;
+      for (int j = 0; j < feature_groups_[i]->num_feature_; ++j) {
+        if (is_feature_used[f_start + j]) {
+          is_group_used = true;
+          break;
+        }
+      }
+      int cur_num_bin = feature_groups_[i]->bin_offsets_.back() - 1;
+      num_total_bin += cur_num_bin;
+      if (is_group_used) {
+        new_num_total_bin += cur_num_bin;
+
+        lower_bound.push_back(num_total_bin - cur_num_bin);
+        upper_bound.push_back(num_total_bin);
+
+        temp_state->hist_move_src.push_back((new_num_total_bin - cur_num_bin) *
+                                            2);
+        temp_state->hist_move_dest.push_back((num_total_bin - cur_num_bin) * 2);
+        temp_state->hist_move_size.push_back(cur_num_bin * 2);
+        delta.push_back(num_total_bin - new_num_total_bin);
+      }
+    }
+  }
+  // avoid out of range
+  lower_bound.push_back(num_total_bin);
+  upper_bound.push_back(num_total_bin);
+  global_timer.Stop("Dataset::InitTrain.Prep");
+  global_timer.Start("Dataset::InitTrain.Resize");
+  if (temp_state->multi_val_bin_subfeature == nullptr) {
+    temp_state->multi_val_bin_subfeature.reset(
+        temp_state->multi_val_bin->CreateLike(new_num_total_bin, num_used,
+                                              sum_used_dense_ratio));
+  } else {
+    temp_state->multi_val_bin_subfeature->ReSizeForSubFeature(
+        new_num_total_bin, num_used, sum_used_dense_ratio);
+  }
+  global_timer.Stop("Dataset::InitTrain.Resize");
+  global_timer.Start("Dataset::InitTrain.CopySubFeature");
+  temp_state->multi_val_bin_subfeature->CopySubFeature(
+      temp_state->multi_val_bin.get(), used_feature_index, lower_bound,
+      upper_bound, delta);
+  global_timer.Stop("Dataset::InitTrain.CopySubFeature");
+}
+
+void Dataset::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, hist_t* hist_data) const {
+  Common::FunctionTimer fun_time("Dataset::ConstructHistogramsMultiVal",
+                                 global_timer);
+  const auto multi_val_bin = temp_state->use_subfeature
+                                 ? temp_state->multi_val_bin_subfeature.get()
+                                 : temp_state->multi_val_bin.get();
   if (multi_val_bin == nullptr) {
     return;
   }
   int num_threads = 1;
-  #pragma omp parallel
-  #pragma omp master
-  {
-    num_threads = omp_get_num_threads();
-  }
+#pragma omp parallel
+#pragma omp master
+  { num_threads = omp_get_num_threads(); }
 
   global_timer.Start("Dataset::sparse_bin_histogram");
   const int num_bin = multi_val_bin->num_bin();
-  const int num_bin_aligned = (num_bin + kAlignedSize - 1) / kAlignedSize * kAlignedSize;
+  const int num_bin_aligned =
+      (num_bin + kAlignedSize - 1) / kAlignedSize * kAlignedSize;
   const int min_data_block_size = 1024;
-  const int n_data_block = std::min(num_threads, (num_data + min_data_block_size - 1) / min_data_block_size);
+  const int n_data_block = std::min(
+      num_threads, (num_data + min_data_block_size - 1) / min_data_block_size);
   const int data_block_size = (num_data + n_data_block - 1) / n_data_block;
 
-  const size_t buf_size = static_cast<size_t>(n_data_block - 1) * num_bin_aligned * 2;
-  if (hist_buf_.size() < buf_size) {
-    hist_buf_.resize(buf_size);
+  const size_t buf_size =
+      static_cast<size_t>(n_data_block - 1) * num_bin_aligned * 2;
+  if (temp_state->hist_buf.size() < buf_size) {
+    temp_state->hist_buf.resize(buf_size);
   }
-
-  #pragma omp parallel for schedule(static)
+  auto origin_hist_data = hist_data;
+  if (temp_state->use_subfeature) {
+    hist_data = temp_state->TempBuf();
+  }
+#pragma omp parallel for schedule(static)
   for (int tid = 0; tid < n_data_block; ++tid) {
     data_size_t start = tid * data_block_size;
     data_size_t end = std::min(start + data_block_size, num_data);
     auto data_ptr = hist_data;
     if (tid > 0) {
-      data_ptr = hist_buf_.data() + static_cast<size_t>(num_bin_aligned) * 2 * (tid - 1);
+      data_ptr = temp_state->hist_buf.data() +
+                 static_cast<size_t>(num_bin_aligned) * 2 * (tid - 1);
     }
     std::memset(reinterpret_cast<void*>(data_ptr), 0, num_bin * kHistEntrySize);
     if (data_indices != nullptr && num_data < num_data_) {
       if (!is_constant_hessian) {
-        multi_val_bin->ConstructHistogram(data_indices, start, end, gradients, hessians, data_ptr);
+        multi_val_bin->ConstructHistogram(data_indices, start, end, gradients,
+                                          hessians, data_ptr);
       } else {
-        multi_val_bin->ConstructHistogram(data_indices, start, end, gradients, data_ptr);
+        multi_val_bin->ConstructHistogram(data_indices, start, end, gradients,
+                                          data_ptr);
       }
     } else {
       if (!is_constant_hessian) {
-        multi_val_bin->ConstructHistogram(start, end, gradients, hessians, data_ptr);
+        multi_val_bin->ConstructHistogram(start, end, gradients, hessians,
+                                          data_ptr);
       } else {
         multi_val_bin->ConstructHistogram(start, end, gradients, data_ptr);
       }
@@ -1153,27 +1328,30 @@ void Dataset::ConstructHistogramsMultiVal(const MultiValBin* multi_val_bin, cons
 
   global_timer.Start("Dataset::sparse_bin_histogram_merge");
   const int min_bin_block_size = 512;
-  const int n_bin_block = std::min(num_threads, (num_bin + min_bin_block_size - 1) / min_bin_block_size);
+  const int n_bin_block = std::min(
+      num_threads, (num_bin + min_bin_block_size - 1) / min_bin_block_size);
   const int bin_block_size = (num_bin + n_bin_block - 1) / n_bin_block;
   if (!is_constant_hessian) {
-    #pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static)
     for (int t = 0; t < n_bin_block; ++t) {
       const int start = t * bin_block_size;
       const int end = std::min(start + bin_block_size, num_bin);
       for (int tid = 1; tid < n_data_block; ++tid) {
-        auto src_ptr = hist_buf_.data() + static_cast<size_t>(num_bin_aligned) * 2 * (tid - 1);
+        auto src_ptr = temp_state->hist_buf.data() +
+                       static_cast<size_t>(num_bin_aligned) * 2 * (tid - 1);
         for (int i = start * 2; i < end * 2; ++i) {
           hist_data[i] += src_ptr[i];
         }
       }
     }
   } else {
-    #pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static)
     for (int t = 0; t < n_bin_block; ++t) {
       const int start = t * bin_block_size;
       const int end = std::min(start + bin_block_size, num_bin);
       for (int tid = 1; tid < n_data_block; ++tid) {
-        auto src_ptr = hist_buf_.data() + static_cast<size_t>(num_bin_aligned) * 2 * (tid - 1);
+        auto src_ptr = temp_state->hist_buf.data() +
+                       static_cast<size_t>(num_bin_aligned) * 2 * (tid - 1);
         for (int i = start * 2; i < end * 2; ++i) {
           hist_data[i] += src_ptr[i];
         }
@@ -1184,21 +1362,25 @@ void Dataset::ConstructHistogramsMultiVal(const MultiValBin* multi_val_bin, cons
     }
   }
   global_timer.Stop("Dataset::sparse_bin_histogram_merge");
+  global_timer.Start("Dataset::sparse_bin_histogram_move");
+  temp_state->HistMove(hist_data, origin_hist_data);
+  global_timer.Stop("Dataset::sparse_bin_histogram_move");
 }
 
-void Dataset::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,
-                                  const MultiValBin* multi_val_bin, bool is_colwise,
-                                  hist_t* hist_data) const {
+void Dataset::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,
+    hist_t* hist_data) const {
   Common::FunctionTimer fun_timer("Dataset::ConstructHistograms", global_timer);
   if (num_data < 0 || hist_data == nullptr) {
     return;
   }
   if (!is_colwise) {
-    return ConstructHistogramsMultiVal(multi_val_bin, data_indices, num_data, gradients, hessians, is_constant_hessian, hist_data);
+    return ConstructHistogramsMultiVal(data_indices, num_data, gradients,
+                                       hessians, is_constant_hessian,
+                                       temp_state, hist_data);
   }
   global_timer.Start("Dataset::Get used group");
   std::vector<int> used_dense_group;
@@ -1230,13 +1412,13 @@ void Dataset::ConstructHistograms(const std::vector<int8_t>& is_feature_used,
     auto ptr_ordered_hess = hessians;
     if (data_indices != nullptr && num_data < num_data_) {
       if (!is_constant_hessian) {
-        #pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static)
         for (data_size_t i = 0; i < num_data; ++i) {
           ordered_gradients[i] = gradients[data_indices[i]];
           ordered_hessians[i] = hessians[data_indices[i]];
         }
       } else {
-        #pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static)
         for (data_size_t i = 0; i < num_data; ++i) {
           ordered_gradients[i] = gradients[data_indices[i]];
         }
@@ -1245,7 +1427,7 @@ void Dataset::ConstructHistograms(const std::vector<int8_t>& is_feature_used,
       ptr_ordered_hess = ordered_hessians;
       if (!is_constant_hessian) {
         OMP_INIT_EX();
-        #pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static)
         for (int gi = 0; gi < num_used_dense_group; ++gi) {
           OMP_LOOP_EX_BEGIN();
           int group = used_dense_group[gi];
@@ -1256,14 +1438,15 @@ void Dataset::ConstructHistograms(const std::vector<int8_t>& is_feature_used,
                       num_bin * kHistEntrySize);
           // construct histograms for smaller leaf
           feature_groups_[group]->bin_data_->ConstructHistogram(
-              data_indices, 0, num_data, ptr_ordered_grad, ptr_ordered_hess, data_ptr);
+              data_indices, 0, num_data, ptr_ordered_grad, ptr_ordered_hess,
+              data_ptr);
           OMP_LOOP_EX_END();
         }
         OMP_THROW_EX();
 
       } else {
         OMP_INIT_EX();
-        #pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static)
         for (int gi = 0; gi < num_used_dense_group; ++gi) {
           OMP_LOOP_EX_BEGIN();
           int group = used_dense_group[gi];
@@ -1286,7 +1469,7 @@ void Dataset::ConstructHistograms(const std::vector<int8_t>& is_feature_used,
     } else {
       if (!is_constant_hessian) {
         OMP_INIT_EX();
-        #pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static)
         for (int gi = 0; gi < num_used_dense_group; ++gi) {
           OMP_LOOP_EX_BEGIN();
           int group = used_dense_group[gi];
@@ -1303,7 +1486,7 @@ void Dataset::ConstructHistograms(const std::vector<int8_t>& is_feature_used,
         OMP_THROW_EX();
       } else {
         OMP_INIT_EX();
-        #pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static)
         for (int gi = 0; gi < num_used_dense_group; ++gi) {
           OMP_LOOP_EX_BEGIN();
           int group = used_dense_group[gi];
@@ -1327,8 +1510,9 @@ void Dataset::ConstructHistograms(const std::vector<int8_t>& is_feature_used,
   }
   global_timer.Stop("Dataset::dense_bin_histogram");
   if (multi_val_groud_id >= 0) {
-    ConstructHistogramsMultiVal(multi_val_bin, data_indices, num_data, gradients, hessians, is_constant_hessian,
-                                hist_data + group_bin_boundaries_[multi_val_groud_id] * 2);
+    ConstructHistogramsMultiVal(
+        data_indices, num_data, gradients, hessians, is_constant_hessian,
+        temp_state, hist_data + group_bin_boundaries_[multi_val_groud_id] * 2);
   }
 }
 

src/io/multi_val_dense_bin.hpp

@@ -5,6 +5,7 @@
 #ifndef LIGHTGBM_IO_MULTI_VAL_DENSE_BIN_HPP_
 #define LIGHTGBM_IO_MULTI_VAL_DENSE_BIN_HPP_
 
+#include <LightGBM/utils/openmp_wrapper.h>
 #include <LightGBM/bin.h>
 
 #include <cstdint>
@@ -34,9 +35,6 @@ class MultiValDenseBin : public MultiValBin {
 
   void PushOneRow(int , data_size_t idx, const std::vector<uint32_t>& values) override {
     auto start = RowPtr(idx);
-#ifdef DEBUG
-    CHECK(num_feature_ == static_cast<int>(values.size()));
-#endif  // DEBUG
     for (auto i = 0; i < num_feature_; ++i) {
       data_[start + i] = static_cast<VAL_T>(values[i]);
     }
@@ -128,10 +126,60 @@ class MultiValDenseBin : public MultiValBin {
 
   void CopySubset(const Bin* full_bin, const data_size_t* used_indices, data_size_t num_used_indices) override {
     auto other_bin = dynamic_cast<const MultiValDenseBin<VAL_T>*>(full_bin);
-    data_.clear();
+    data_.resize(num_feature_ * num_used_indices);
     for (data_size_t i = 0; i < num_used_indices; ++i) {
-      for (int64_t j = other_bin->RowPtr(used_indices[i]); j < other_bin->RowPtr(used_indices[i] + 1); ++j) {
-        data_.push_back(other_bin->data_[j]);
+      auto j_start = RowPtr(i);
+      auto other_j_start = other_bin->RowPtr(used_indices[i]);
+      for (int64_t j = other_j_start;
+           j < other_bin->RowPtr(used_indices[i] + 1); ++j) {
+        data_[j - other_j_start + j_start] = other_bin->data_[j];
+      }
+    }
+  }
+
+  MultiValBin* CreateLike(int num_bin, int num_feature, double) const override {
+    return new MultiValDenseBin<VAL_T>(num_data_, num_bin, num_feature);
+  }
+
+  void ReSizeForSubFeature(int num_bin, int num_feature, double) override {
+    num_bin_ = num_bin;
+    num_feature_ = num_feature;
+    size_t new_size = static_cast<size_t>(num_feature_) * num_data_;
+    if (data_.size() < new_size) {
+      data_.resize(new_size, 0);
+    }
+  }
+
+  void CopySubFeature(const MultiValBin* full_bin,
+                      const std::vector<int>& used_feature_index,
+                      const std::vector<uint32_t>&,
+                      const std::vector<uint32_t>&,
+                      const std::vector<uint32_t>& delta) override {
+    const auto other = reinterpret_cast<const MultiValDenseBin<VAL_T>*>(full_bin);
+    int num_threads = 1;
+#pragma omp parallel
+#pragma omp master
+    { num_threads = omp_get_num_threads(); }
+
+    const int min_block_size = 1024;
+    const int n_block = std::min(
+        num_threads, (num_data_ + min_block_size - 1) / min_block_size);
+    const data_size_t block_size = (num_data_ + n_block - 1) / n_block;
+#pragma omp parallel for schedule(static, 1)
+    for (int tid = 0; tid < n_block; ++tid) {
+      data_size_t start = tid * block_size;
+      data_size_t end = std::min(num_data_, start + block_size);
+      for (data_size_t i = start; i < end; ++i) {
+        const auto j_start = RowPtr(i);
+        const auto other_j_start = other->RowPtr(i);
+        for (int j = 0; j < num_feature_; ++j) {
+          if (other->data_[other_j_start + used_feature_index[j]] > 0) {
+            data_[j_start + j] = static_cast<VAL_T>(
+                other->data_[other_j_start + used_feature_index[j]] - delta[j]);
+          } else {
+            data_[j_start + j] = 0;
+          }
+        }
       }
     }
   }

src/io/multi_val_sparse_bin.hpp

 /*!
  * Copyright (c) 2020 Microsoft Corporation. All rights reserved.
- * Licensed under the MIT License. See LICENSE file in the project root for license information.
+ * Licensed under the MIT License. See LICENSE file in the project root for
+ * license information.
  */
 #ifndef LIGHTGBM_IO_MULTI_VAL_SPARSE_BIN_HPP_
 #define LIGHTGBM_IO_MULTI_VAL_SPARSE_BIN_HPP_
@@ -14,75 +15,94 @@
 
 namespace LightGBM {
 
-
 template <typename VAL_T>
 class MultiValSparseBin : public MultiValBin {
  public:
-  explicit MultiValSparseBin(data_size_t num_data, int num_bin)
-    : num_data_(num_data), num_bin_(num_bin) {
+  explicit MultiValSparseBin(data_size_t num_data, int num_bin,
+                             double estimate_element_per_row)
+      : num_data_(num_data),
+        num_bin_(num_bin),
+        estimate_element_per_row_(estimate_element_per_row) {
     row_ptr_.resize(num_data_ + 1, 0);
-    data_.reserve(num_data_);
+    data_size_t estimate_num_data =
+        static_cast<data_size_t>(num_data_ * estimate_element_per_row_ * 1.1);
     int num_threads = 1;
-    #pragma omp parallel
-    #pragma omp master
-    {
-      num_threads = omp_get_num_threads();
-    }
+#pragma omp parallel
+#pragma omp master
+    { num_threads = omp_get_num_threads(); }
     if (num_threads > 1) {
       t_data_.resize(num_threads - 1);
+      for (size_t i = 0; i < t_data_.size(); ++i) {
+        t_data_[i].resize(estimate_num_data / num_threads);
+      }
     }
+    t_size_.resize(num_threads, 0);
+    data_.resize(estimate_num_data / num_threads);
   }
 
-  ~MultiValSparseBin() {
-  }
+  ~MultiValSparseBin() {}
 
-  data_size_t num_data() const override {
-    return num_data_;
-  }
+  data_size_t num_data() const override { return num_data_; }
 
-  int num_bin() const override {
-    return num_bin_;
-  }
+  int num_bin() const override { return num_bin_; }
 
-  void PushOneRow(int tid, data_size_t idx, const std::vector<uint32_t> & values) override {
+  void PushOneRow(int tid, data_size_t idx,
+                  const std::vector<uint32_t>& values) override {
+    const int pre_alloc_size = 50;
     row_ptr_[idx + 1] = static_cast<data_size_t>(values.size());
     if (tid == 0) {
+      if (t_size_[tid] + row_ptr_[idx + 1] > static_cast<data_size_t>(data_.size())) {
+        data_.resize(t_size_[tid] + row_ptr_[idx + 1] * pre_alloc_size);
+      }
       for (auto val : values) {
-        data_.push_back(static_cast<VAL_T>(val));
+        data_[t_size_[tid]++] = static_cast<VAL_T>(val);
       }
     } else {
+      if (t_size_[tid] + row_ptr_[idx + 1] > static_cast<data_size_t>(t_data_[tid - 1].size())) {
+        t_data_[tid - 1].resize(t_size_[tid] +
+                                row_ptr_[idx + 1] * pre_alloc_size);
+      }
       for (auto val : values) {
-        t_data_[tid - 1].push_back(static_cast<VAL_T>(val));
+        t_data_[tid - 1][t_size_[tid]++] = static_cast<VAL_T>(val);
       }
     }
   }
 
-  void FinishLoad() override {
+  void MergeData(const data_size_t* sizes) {
+    Common::FunctionTimer fun_time("MultiValSparseBin::MergeData", global_timer);
     for (data_size_t i = 0; i < num_data_; ++i) {
       row_ptr_[i + 1] += row_ptr_[i];
     }
     if (t_data_.size() > 0) {
-      std::vector<size_t> offsets;
-      offsets.push_back(data_.size());
+      std::vector<data_size_t> offsets(1 + t_data_.size());
+      offsets[0] = sizes[0];
       for (size_t tid = 0; tid < t_data_.size() - 1; ++tid) {
-        offsets.push_back(offsets.back() + t_data_[tid].size());
+        offsets[tid + 1] = offsets[tid] + sizes[tid + 1];
       }
       data_.resize(row_ptr_[num_data_]);
-      #pragma omp parallel for schedule(static)
+#pragma omp parallel for schedule(static)
       for (int tid = 0; tid < static_cast<int>(t_data_.size()); ++tid) {
-        std::copy_n(t_data_[tid].data(), t_data_[tid].size(),
+        std::copy_n(t_data_[tid].data(), sizes[tid + 1],
                     data_.data() + offsets[tid]);
       }
+    } else {
+      data_.resize(row_ptr_[num_data_]);
     }
+  }
+
+  void FinishLoad() override {
+    MergeData(t_size_.data());
+    t_size_.clear();
     row_ptr_.shrink_to_fit();
     data_.shrink_to_fit();
     t_data_.clear();
     t_data_.shrink_to_fit();
+    // update estimate_element_per_row_ by all data
+    estimate_element_per_row_ =
+        static_cast<double>(row_ptr_[num_data_]) / num_data_;
   }
 
-  bool IsSparse() override {
-    return true;
-  }
+  bool IsSparse() override { return true; }
 
   void ReSize(data_size_t num_data) override {
     if (num_data_ != num_data) {
@@ -90,14 +110,16 @@ class MultiValSparseBin : public MultiValBin {
     }
   }
 
-  #define ACC_GH(hist, i, g, h) \
+#define ACC_GH(hist, i, g, h)               \
   const auto ti = static_cast<int>(i) << 1; \
-  hist[ti] += g; \
-  hist[ti + 1] += h; \
+  hist[ti] += g;                            \
+  hist[ti + 1] += h;
 
-  template<bool use_indices, bool use_prefetch, bool use_hessians>
-  void ConstructHistogramInner(const data_size_t* data_indices, data_size_t start, data_size_t end,
-    const score_t* gradients, const score_t* hessians, hist_t* out) const {
+  template <bool use_indices, bool use_prefetch, bool use_hessians>
+  void ConstructHistogramInner(const data_size_t* data_indices,
+                               data_size_t start, data_size_t end,
+                               const score_t* gradients,
+                               const score_t* hessians, hist_t* out) const {
     data_size_t i = start;
     if (use_prefetch) {
       const data_size_t pf_offset = 32 / sizeof(VAL_T);
@@ -105,7 +127,8 @@ class MultiValSparseBin : public MultiValBin {
 
       for (; i < pf_end; ++i) {
         const auto idx = use_indices ? data_indices[i] : i;
-        const auto pf_idx = use_indices ? data_indices[i + pf_offset] : i + pf_offset;
+        const auto pf_idx =
+            use_indices ? data_indices[i + pf_offset] : i + pf_offset;
         PREFETCH_T0(gradients + pf_idx);
         if (use_hessians) {
           PREFETCH_T0(hessians + pf_idx);
@@ -138,66 +161,153 @@ class MultiValSparseBin : public MultiValBin {
       }
     }
   }
-  #undef ACC_GH
+#undef ACC_GH
 
-  void ConstructHistogram(const data_size_t* data_indices, data_size_t start, data_size_t end,
-    const score_t* gradients, const score_t* hessians,
-    hist_t* out) const override {
-    ConstructHistogramInner<true, true, true>(data_indices, start, end, gradients, hessians, out);
+  void ConstructHistogram(const data_size_t* data_indices, data_size_t start,
+                          data_size_t end, const score_t* gradients,
+                          const score_t* hessians, hist_t* out) const override {
+    ConstructHistogramInner<true, true, true>(data_indices, start, end,
+                                              gradients, hessians, out);
   }
 
   void ConstructHistogram(data_size_t start, data_size_t end,
-    const score_t* gradients, const score_t* hessians,
-    hist_t* out) const override {
-    ConstructHistogramInner<false, false, true>(nullptr, start, end, gradients, hessians, out);
+                          const score_t* gradients, const score_t* hessians,
+                          hist_t* out) const override {
+    ConstructHistogramInner<false, false, true>(nullptr, start, end, gradients,
+                                                hessians, out);
   }
 
-  void ConstructHistogram(const data_size_t* data_indices, data_size_t start, data_size_t end,
-    const score_t* gradients,
-    hist_t* out) const override {
-    ConstructHistogramInner<true, true, false>(data_indices, start, end, gradients, nullptr, out);
+  void ConstructHistogram(const data_size_t* data_indices, data_size_t start,
+                          data_size_t end, const score_t* gradients,
+                          hist_t* out) const override {
+    ConstructHistogramInner<true, true, false>(data_indices, start, end,
+                                               gradients, nullptr, out);
   }
 
   void ConstructHistogram(data_size_t start, data_size_t end,
-    const score_t* gradients,
-    hist_t* out) const override {
-    ConstructHistogramInner<false, false, false>(nullptr, start, end, gradients, nullptr, out);
+                          const score_t* gradients,
+                          hist_t* out) const override {
+    ConstructHistogramInner<false, false, false>(nullptr, start, end, gradients,
+                                                 nullptr, out);
   }
 
-  void CopySubset(const Bin * full_bin, const data_size_t * used_indices, data_size_t num_used_indices) override {
+  void CopySubset(const Bin* full_bin, const data_size_t* used_indices,
+                  data_size_t num_used_indices) override {
     auto other_bin = dynamic_cast<const MultiValSparseBin<VAL_T>*>(full_bin);
     row_ptr_.resize(num_data_ + 1, 0);
+    data_size_t estimate_num_data =
+        static_cast<data_size_t>(num_data_ * estimate_element_per_row_ * 1.5);
     data_.clear();
+    data_.reserve(estimate_num_data);
     for (data_size_t i = 0; i < num_used_indices; ++i) {
-      for (data_size_t j = other_bin->row_ptr_[used_indices[i]]; j < other_bin->row_ptr_[used_indices[i] + 1]; ++j) {
+      for (data_size_t j = other_bin->row_ptr_[used_indices[i]];
+           j < other_bin->row_ptr_[used_indices[i] + 1]; ++j) {
         data_.push_back(other_bin->data_[j]);
       }
-      row_ptr_[i + 1] = row_ptr_[i] + other_bin->row_ptr_[used_indices[i] + 1] - other_bin->row_ptr_[used_indices[i]];
+      row_ptr_[i + 1] = row_ptr_[i] + other_bin->row_ptr_[used_indices[i] + 1] -
+                        other_bin->row_ptr_[used_indices[i]];
+    }
+  }
+
+  MultiValBin* CreateLike(int num_bin, int,
+                          double estimate_element_per_row) const override {
+    return new MultiValSparseBin<VAL_T>(num_data_, num_bin,
+                                        estimate_element_per_row);
+  }
+
+  void ReSizeForSubFeature(int num_bin, int,
+                           double estimate_element_per_row) override {
+    num_bin_ = num_bin;
+    estimate_element_per_row_ = estimate_element_per_row;
+    data_size_t estimate_num_data =
+        static_cast<data_size_t>(num_data_ * estimate_element_per_row_ * 1.1);
+    size_t npart = 1 + t_data_.size();
+    data_size_t avg_num_data =
+        static_cast<data_size_t>(estimate_num_data / npart);
+    if (static_cast<data_size_t>(data_.size()) < avg_num_data) {
+      data_.resize(avg_num_data, 0);
+    }
+    for (size_t i = 0; i < t_data_.size(); ++i) {
+      if (static_cast<data_size_t>(t_data_[i].size()) < avg_num_data) {
+        t_data_[i].resize(avg_num_data, 0);
+      }
     }
   }
 
-  inline data_size_t RowPtr(data_size_t idx) const {
-    return row_ptr_[idx];
+  void CopySubFeature(const MultiValBin* full_bin, const std::vector<int>&,
+                      const std::vector<uint32_t>& lower,
+                      const std::vector<uint32_t>& upper,
+                      const std::vector<uint32_t>& delta) override {
+    const auto other =
+        reinterpret_cast<const MultiValSparseBin<VAL_T>*>(full_bin);
+    int num_threads = 1;
+#pragma omp parallel
+#pragma omp master
+    { num_threads = omp_get_num_threads(); }
+
+    const int min_block_size = 1024;
+    const int n_block = std::min(
+        num_threads, (num_data_ + min_block_size - 1) / min_block_size);
+    const data_size_t block_size = (num_data_ + n_block - 1) / n_block;
+    std::vector<data_size_t> sizes(t_data_.size() + 1, 0);
+    const int pre_alloc_size = 50;
+#pragma omp parallel for schedule(static, 1)
+    for (int tid = 0; tid < n_block; ++tid) {
+      data_size_t start = tid * block_size;
+      data_size_t end = std::min(num_data_, start + block_size);
+      auto& buf = (tid == 0) ? data_ : t_data_[tid - 1];
+      data_size_t size = 0;
+      for (data_size_t i = start; i < end; ++i) {
+        const auto j_start = other->RowPtr(i);
+        const auto j_end = other->RowPtr(i + 1);
+        if (size + (j_end - j_start) > static_cast<data_size_t>(buf.size())) {
+          buf.resize(size + (j_end - j_start) * pre_alloc_size);
+        }
+        int k = 0;
+        const data_size_t pre_size = size;
+        for (auto j = j_start; j < j_end; ++j) {
+          auto val = other->data_[j];
+          while (val >= upper[k]) {
+            ++k;
+          }
+          if (val >= lower[k]) {
+            buf[size++] = static_cast<VAL_T>(val - delta[k]);
+          }
+        }
+        row_ptr_[i + 1] = size - pre_size;
+      }
+      sizes[tid] = size;
+    }
+    MergeData(sizes.data());
   }
 
+  inline data_size_t RowPtr(data_size_t idx) const { return row_ptr_[idx]; }
+
   MultiValSparseBin<VAL_T>* Clone() override;
 
  private:
   data_size_t num_data_;
   int num_bin_;
+  double estimate_element_per_row_;
   std::vector<VAL_T, Common::AlignmentAllocator<VAL_T, 32>> data_;
-  std::vector<data_size_t, Common::AlignmentAllocator<data_size_t, 32>> row_ptr_;
-  std::vector<std::vector<VAL_T>> t_data_;
+  std::vector<data_size_t, Common::AlignmentAllocator<data_size_t, 32>>
+      row_ptr_;
+  std::vector<std::vector<VAL_T, Common::AlignmentAllocator<VAL_T, 32>>>
+      t_data_;
+  std::vector<data_size_t> t_size_;
 
-  MultiValSparseBin<VAL_T>(const MultiValSparseBin<VAL_T> & other)
-    : num_data_(other.num_data_), num_bin_(other.num_bin_), data_(other.data_), row_ptr_(other.row_ptr_) {
-  }
+  MultiValSparseBin<VAL_T>(const MultiValSparseBin<VAL_T>& other)
+      : num_data_(other.num_data_),
+        num_bin_(other.num_bin_),
+        estimate_element_per_row_(other.estimate_element_per_row_),
+        data_(other.data_),
+        row_ptr_(other.row_ptr_) {}
 };
 
-template<typename VAL_T>
+template <typename VAL_T>
 MultiValSparseBin<VAL_T>* MultiValSparseBin<VAL_T>::Clone() {
   return new MultiValSparseBin<VAL_T>(*this);
 }
 
 }  // namespace LightGBM
-#endif   // LIGHTGBM_IO_MULTI_VAL_SPARSE_BIN_HPP_
+#endif  // LIGHTGBM_IO_MULTI_VAL_SPARSE_BIN_HPP_

src/treelearner/gpu_tree_learner.cpp

@@ -974,7 +974,7 @@ void GPUTreeLearner::ConstructHistograms(const std::vector<int8_t>& is_feature_u
     smaller_leaf_splits_->data_indices(), smaller_leaf_splits_->num_data_in_leaf(),
     gradients_, hessians_,
     ordered_gradients_.data(), ordered_hessians_.data(), is_constant_hessian_,
-    multi_val_bin_.get(), is_hist_colwise_,
+    is_hist_colwise_, temp_state_.get(),
     ptr_smaller_leaf_hist_data);
   // wait for GPU to finish, only if GPU is actually used
   if (is_gpu_used) {
@@ -1039,7 +1039,7 @@ void GPUTreeLearner::ConstructHistograms(const std::vector<int8_t>& is_feature_u
       larger_leaf_splits_->data_indices(), larger_leaf_splits_->num_data_in_leaf(),
       gradients_, hessians_,
       ordered_gradients_.data(), ordered_hessians_.data(), is_constant_hessian_,
-      multi_val_bin_.get(), is_hist_colwise_,
+      is_hist_colwise_, temp_state_.get(),
       ptr_larger_leaf_hist_data);
     // wait for GPU to finish, only if GPU is actually used
     if (is_gpu_used) {

src/treelearner/serial_tree_learner.cpp

@@ -80,15 +80,18 @@ void SerialTreeLearner::Init(const Dataset* train_data, bool is_constant_hessian
 void SerialTreeLearner::GetMultiValBin(const Dataset* dataset, bool is_first_time) {
   if (is_first_time) {
     auto used_feature = GetUsedFeatures(true);
-    multi_val_bin_.reset(dataset->TestMultiThreadingMethod(ordered_gradients_.data(), ordered_hessians_.data(), used_feature,
+    temp_state_.reset(dataset->TestMultiThreadingMethod(
+        ordered_gradients_.data(), ordered_hessians_.data(), used_feature,
       is_constant_hessian_, config_->force_col_wise, config_->force_row_wise, &is_hist_colwise_));
   } else {
     // cannot change is_hist_col_wise during training
-    multi_val_bin_.reset(dataset->TestMultiThreadingMethod(ordered_gradients_.data(), ordered_hessians_.data(), is_feature_used_,
+    temp_state_.reset(dataset->TestMultiThreadingMethod(
+        ordered_gradients_.data(), ordered_hessians_.data(), is_feature_used_,
       is_constant_hessian_, is_hist_colwise_, !is_hist_colwise_, &is_hist_colwise_));
   }
 }
 
+// Todo: optimized bagging for multi-val bin
 void SerialTreeLearner::ResetTrainingData(const Dataset* train_data) {
   train_data_ = train_data;
   num_data_ = train_data_->num_data();
@@ -288,7 +291,7 @@ void SerialTreeLearner::BeforeTrain() {
       is_feature_used_[i] = 1;
     }
   }
-
+  train_data_->InitTrain(is_feature_used_, is_hist_colwise_, temp_state_.get());
   // initialize data partition
   data_partition_->Init();
 
@@ -375,22 +378,21 @@ void SerialTreeLearner::ConstructHistograms(const std::vector<int8_t>& is_featur
   Common::FunctionTimer fun_timer("SerialTreeLearner::ConstructHistograms", global_timer);
   // construct smaller leaf
   hist_t* ptr_smaller_leaf_hist_data = smaller_leaf_histogram_array_[0].RawData() - kHistOffset;
-  train_data_->ConstructHistograms(is_feature_used,
-                                   smaller_leaf_splits_->data_indices(), smaller_leaf_splits_->num_data_in_leaf(),
-                                   gradients_, hessians_,
-                                   ordered_gradients_.data(), ordered_hessians_.data(), is_constant_hessian_,
-                                   multi_val_bin_.get(), is_hist_colwise_,
-                                   ptr_smaller_leaf_hist_data);
+  train_data_->ConstructHistograms(
+      is_feature_used, smaller_leaf_splits_->data_indices(),
+      smaller_leaf_splits_->num_data_in_leaf(), gradients_, hessians_,
+      ordered_gradients_.data(), ordered_hessians_.data(), is_constant_hessian_,
+      is_hist_colwise_, temp_state_.get(), ptr_smaller_leaf_hist_data);
 
   if (larger_leaf_histogram_array_ != nullptr && !use_subtract) {
     // construct larger leaf
     hist_t* ptr_larger_leaf_hist_data = larger_leaf_histogram_array_[0].RawData() - kHistOffset;
-    train_data_->ConstructHistograms(is_feature_used,
-                                     larger_leaf_splits_->data_indices(), larger_leaf_splits_->num_data_in_leaf(),
-                                     gradients_, hessians_,
-                                     ordered_gradients_.data(), ordered_hessians_.data(), is_constant_hessian_,
-                                     multi_val_bin_.get(), is_hist_colwise_,
-                                     ptr_larger_leaf_hist_data);
+    train_data_->ConstructHistograms(
+        is_feature_used, larger_leaf_splits_->data_indices(),
+        larger_leaf_splits_->num_data_in_leaf(), gradients_, hessians_,
+        ordered_gradients_.data(), ordered_hessians_.data(),
+        is_constant_hessian_, is_hist_colwise_, temp_state_.get(),
+        ptr_larger_leaf_hist_data);
   }
 }
 

src/treelearner/serial_tree_learner.h

@@ -188,7 +188,7 @@ class SerialTreeLearner: public TreeLearner {
   int num_threads_;
   std::vector<int> ordered_bin_indices_;
   bool is_constant_hessian_;
-  std::unique_ptr<MultiValBin> multi_val_bin_;
+  std::unique_ptr<TrainingTempState> temp_state_;
   bool is_hist_colwise_;
   std::unique_ptr<CostEfficientGradientBoosting> cegb_;
 };