Use multiple buffers.

trivialfis · Oct 27, 2021 · 2849f57 · 2849f57
1 parent 38b3c48
commit 2849f57
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 21 deletions.
diff --git a/src/common/hist_util.cc b/src/common/hist_util.cc
@@ -9,14 +9,13 @@
 #include <numeric>
 #include <vector>
 
-#include "xgboost/base.h"
 #include "../common/common.h"
-#include "hist_util.h"
-#include "random.h"
+#include "../data/gradient_index.h"
 #include "column_matrix.h"
+#include "hist_util.h"
 #include "quantile.h"
-#include "./../tree/updater_quantile_hist.h"
-#include "../data/gradient_index.h"
+#include "random.h"
+#include "xgboost/base.h"
 
 #if defined(XGBOOST_MM_PREFETCH_PRESENT)
   #include <xmmintrin.h>

diff --git a/src/tree/hist/histogram.h b/src/tree/hist/histogram.h
@@ -25,7 +25,7 @@ template <typename GradientSumT, typename ExpandEntry> class HistogramBuilder {
   /*! \brief culmulative local parent histogram of gradients. */
   common::HistCollection<GradientSumT> hist_local_worker_;
   common::GHistBuilder<GradientSumT> builder_;
-  common::ParallelGHistBuilder<GradientSumT> buffer_;
+  std::vector<common::ParallelGHistBuilder<GradientSumT>> buffer_;
   rabit::Reducer<GradientPairT, GradientPairT::Reduce> reducer_;
   BatchParam param_;
   int32_t n_threads_ {-1};
@@ -50,7 +50,10 @@ template <typename GradientSumT, typename ExpandEntry> class HistogramBuilder {
     param_ = p;
     hist_.Init(total_bins);
     hist_local_worker_.Init(total_bins);
-    buffer_.Init(total_bins);
+    buffer_.resize(n_batches);
+    for (auto& v : buffer_) {
+      v.Init(total_bins);
+    }
     builder_ = common::GHistBuilder<GradientSumT>(total_bins);
     is_distributed_ = is_distributed;
   }
@@ -71,10 +74,7 @@ template <typename GradientSumT, typename ExpandEntry> class HistogramBuilder {
       const int32_t nid = nodes_for_explicit_hist_build[i].nid;
       target_hists[i] = hist_[nid];
     }
-    if (page_idx == 0) {
-      // FIXME: Handle different size of space.
-      buffer_.Reset(this->n_threads_, n_nodes, space, target_hists);
-    }
+    buffer_[page_idx].Reset(this->n_threads_, n_nodes, space, target_hists);
 
     // Parallel processing by nodes and data in each node
     common::ParallelFor2d(
@@ -86,7 +86,7 @@ template <typename GradientSumT, typename ExpandEntry> class HistogramBuilder {
           auto end_of_row_set = std::min(r.end(), elem.Size());
           auto rid_set = common::RowSetCollection::Elem(
               elem.begin + start_of_row_set, elem.begin + end_of_row_set, nid);
-          auto hist = buffer_.GetInitializedHist(tid, nid_in_set);
+          auto hist = buffer_[page_idx].GetInitializedHist(tid, nid_in_set);
           if (rid_set.Size() != 0) {
             builder_.template BuildHist<any_missing>(gpair_h, rid_set, gidx,
                                                      hist);
@@ -141,11 +141,11 @@ template <typename GradientSumT, typename ExpandEntry> class HistogramBuilder {
     }
 
     if (is_distributed_) {
-      this->SyncHistogramDistributed(p_tree, nodes_for_explicit_hist_build,
-                                     nodes_for_subtraction_trick,
-                                     starting_index, sync_count);
+      this->SyncHistogramDistributed(
+          page_id, p_tree, nodes_for_explicit_hist_build,
+          nodes_for_subtraction_trick, starting_index, sync_count);
     } else {
-      this->SyncHistogramLocal(p_tree, nodes_for_explicit_hist_build,
+      this->SyncHistogramLocal(page_id, p_tree, nodes_for_explicit_hist_build,
                                nodes_for_subtraction_trick, starting_index,
                                sync_count);
     }
@@ -171,6 +171,7 @@ template <typename GradientSumT, typename ExpandEntry> class HistogramBuilder {
   }
 
   void SyncHistogramDistributed(
+      size_t page_idx,
       RegTree *p_tree,
       std::vector<ExpandEntry> const &nodes_for_explicit_hist_build,
       std::vector<ExpandEntry> const &nodes_for_subtraction_trick,
@@ -184,7 +185,7 @@ template <typename GradientSumT, typename ExpandEntry> class HistogramBuilder {
           const auto &entry = nodes_for_explicit_hist_build[node];
           auto this_hist = this->hist_[entry.nid];
           // Merging histograms from each thread into once
-          buffer_.ReduceHist(node, r.begin(), r.end());
+          buffer_[page_idx].ReduceHist(node, r.begin(), r.end());
           // Store posible parent node
           auto this_local = hist_local_worker_[entry.nid];
           common::CopyHist(this_local, this_hist, r.begin(), r.end());
@@ -217,6 +218,7 @@ template <typename GradientSumT, typename ExpandEntry> class HistogramBuilder {
   }
 
   void SyncHistogramLocal(
+      size_t page_idx,
       RegTree *p_tree,
       std::vector<ExpandEntry> const &nodes_for_explicit_hist_build,
       std::vector<ExpandEntry> const &nodes_for_subtraction_trick,
@@ -231,7 +233,7 @@ template <typename GradientSumT, typename ExpandEntry> class HistogramBuilder {
           const auto &entry = nodes_for_explicit_hist_build[node];
           auto this_hist = this->hist_[entry.nid];
           // Merging histograms from each thread into once
-          this->buffer_.ReduceHist(node, r.begin(), r.end());
+          this->buffer_[page_idx].ReduceHist(node, r.begin(), r.end());
 
           if (!(*p_tree)[entry.nid].IsRoot()) {
             const size_t parent_id = (*p_tree)[entry.nid].Parent();

diff --git a/tests/cpp/tree/hist/test_histogram.cc b/tests/cpp/tree/hist/test_histogram.cc
@@ -175,14 +175,14 @@ void TestSyncHist(bool is_distributed) {
     }
   }
 
-  histogram.Buffer().Reset(1, n_nodes, space, target_hists);
+  histogram.Buffer()[0].Reset(1, n_nodes, space, target_hists);
   // sync hist
   if (is_distributed) {
-    histogram.SyncHistogramDistributed(&tree, nodes_for_explicit_hist_build_,
+    histogram.SyncHistogramDistributed(0, &tree, nodes_for_explicit_hist_build_,
                                        nodes_for_subtraction_trick_,
                                        starting_index, sync_count);
   } else {
-    histogram.SyncHistogramLocal(&tree, nodes_for_explicit_hist_build_,
+    histogram.SyncHistogramLocal(0, &tree, nodes_for_explicit_hist_build_,
                                  nodes_for_subtraction_trick_, starting_index,
                                  sync_count);
   }