optimise some memroy reallocations

razdoburdin · Dec 16, 2024 · b1884b2 · b1884b2
1 parent f4f3bd4
commit b1884b2
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/src/tree/common_row_partitioner.h b/src/tree/common_row_partitioner.h
@@ -131,12 +131,19 @@ class CommonRowPartitioner {
   CommonRowPartitioner(Context const* ctx, bst_idx_t num_row, bst_idx_t _base_rowid,
                        bool is_col_split)
       : base_rowid{_base_rowid}, is_col_split_{is_col_split} {
-    row_set_collection_.Clear();
+    Reset(ctx, num_row, _base_rowid, is_col_split);
+  }
+
+  void Reset(Context const* ctx, bst_idx_t num_row, bst_idx_t _base_rowid, bool is_col_split) {
+    base_rowid = _base_rowid;
+
     std::vector<bst_idx_t>& row_indices = *row_set_collection_.Data();
     row_indices.resize(num_row);
 
     bst_idx_t* p_row_indices = row_indices.data();
-    common::Iota(ctx, p_row_indices, p_row_indices + row_indices.size(), base_rowid);
+    common::Iota(ctx, p_row_indices, p_row_indices + num_row, base_rowid);
+
+    row_set_collection_.Clear();
     row_set_collection_.Init();
 
     if (is_col_split_) {

diff --git a/src/tree/updater_quantile_hist.cc b/src/tree/updater_quantile_hist.cc
@@ -346,15 +346,21 @@ class HistUpdater {
   void InitData(DMatrix *fmat, RegTree const *p_tree) {
     monitor_->Start(__func__);
     bst_bin_t n_total_bins{0};
-    partitioner_.clear();
+    size_t page_idx = 0;
     for (auto const &page : fmat->GetBatches<GHistIndexMatrix>(ctx_, HistBatch(param_))) {
       if (n_total_bins == 0) {
         n_total_bins = page.cut.TotalBins();
       } else {
         CHECK_EQ(n_total_bins, page.cut.TotalBins());
       }
-      partitioner_.emplace_back(this->ctx_, page.Size(), page.base_rowid,
-                                fmat->Info().IsColumnSplit());
+      if (page_idx < partitioner_.size()) {
+        partitioner_[page_idx].Reset(this->ctx_, page.Size(), page.base_rowid,
+                                     fmat->Info().IsColumnSplit());
+      } else {
+        partitioner_.emplace_back(this->ctx_, page.Size(), page.base_rowid,
+                                  fmat->Info().IsColumnSplit());
+      }
+      page_idx++;
     }
     histogram_builder_->Reset(ctx_, n_total_bins, 1, HistBatch(param_), collective::IsDistributed(),
                               fmat->Info().IsColumnSplit(), hist_param_);