Prune GPU sketch after each merge

Author: RAMitchell

src/common/quantile.cu

@@ -304,18 +304,8 @@ void MergeImpl(Context const *ctx, Span<SketchEntry const> const &d_x,
 void SketchContainer::Push(Context const *ctx, Span<Entry const> entries, Span<size_t> columns_ptr,
                            common::Span<OffsetT> cuts_ptr, size_t total_cuts, Span<float> weights) {
   curt::SetDevice(ctx->Ordinal());
-  auto &current = this->entries_;
-  auto &columns_ptr_out = this->columns_ptr_;
-  Span<SketchEntry> out;
-  dh::device_vector<SketchEntry> cuts;
-  bool first_window = current.empty();
-  if (!first_window) {
-    cuts.resize(total_cuts);
-    out = dh::ToSpan(cuts);
-  } else {
-    current.resize(total_cuts);
-    out = dh::ToSpan(current);
-  }
+  dh::device_vector<SketchEntry> cuts(total_cuts);
+  auto out = dh::ToSpan(cuts);
   auto ft = this->feature_types_.ConstDeviceSpan();
   if (weights.empty()) {
     auto to_sketch_entry = [] __device__(size_t sample_idx, Span<Entry const> const &column,
@@ -340,19 +330,10 @@ void SketchContainer::Push(Context const *ctx, Span<Entry const> entries, Span<s
     PruneImpl<Entry>(cuts_ptr, entries, columns_ptr, ft, out, to_sketch_entry);
   }
   auto n_uniques = this->ScanInput(ctx, out, cuts_ptr);
-
-  if (!first_window) {
-    CHECK_EQ(columns_ptr_out.Size(), cuts_ptr.size());
-    out = out.subspan(0, n_uniques);
-    this->Merge(ctx, cuts_ptr, out);
-  } else {
-    current.resize(n_uniques);
-    columns_ptr_out.SetDevice(ctx->Device());
-    columns_ptr_out.Resize(cuts_ptr.size());
-
-    auto d_cuts_ptr = columns_ptr_out.DeviceSpan();
-    CopyTo(d_cuts_ptr, cuts_ptr);
-  }
+  CHECK_EQ(this->columns_ptr_.Size(), cuts_ptr.size());
+  this->Merge(ctx, cuts_ptr, out.subspan(0, n_uniques));
+  auto intermediate_num_cuts = static_cast<bst_idx_t>(num_bins_ * kFactor);
+  this->Prune(ctx, intermediate_num_cuts);
 }
 
 size_t SketchContainer::ScanInput(Context const *ctx, Span<SketchEntry> entries,