Sketching from adapters (#5365)

* Sketching from adapters

* Add weights test

src/common/hist_util.h

@@ -179,16 +179,14 @@ class DenseCuts  : public CutsBuilder {
   void Build(DMatrix* p_fmat, uint32_t max_num_bins) override;
 };
 
-// FIXME(trivialfis): Merge this into generic cut builder.
-/*! \brief Builds the cut matrix on the GPU.
- *
- *  \return The row stride across the entire dataset.
- */
-size_t DeviceSketch(int device,
-                    int max_bin,
-                    int gpu_batch_nrows,
-                    DMatrix* dmat,
-                    HistogramCuts* hmat);
+
+HistogramCuts DeviceSketch(int device, DMatrix* dmat, int max_bins,
+                           size_t sketch_batch_num_elements = 10000000);
+
+template <typename AdapterT>
+HistogramCuts AdapterDeviceSketch(AdapterT* adapter, int num_bins,
+                                  float missing,
+                                  size_t sketch_batch_num_elements = 10000000);
 
 /*!
  * \brief preprocessed global index matrix, in CSR format

src/data/adapter.h

@@ -71,6 +71,7 @@ namespace data {
 constexpr size_t kAdapterUnknownSize = std::numeric_limits<size_t >::max();
 
 struct COOTuple {
+  COOTuple() = default;
   XGBOOST_DEVICE COOTuple(size_t row_idx, size_t column_idx, float value)
       : row_idx(row_idx), column_idx(column_idx), value(value) {}
 

src/data/ellpack_page.cu

@@ -78,6 +78,20 @@ EllpackPageImpl::EllpackPageImpl(int device, EllpackInfo info, size_t n_rows) {
   monitor_.StopCuda("InitCompressedData");
 }
 
+size_t GetRowStride(DMatrix* dmat) {
+  if (dmat->IsDense()) return dmat->Info().num_col_;
+
+  size_t row_stride = 0;
+  for (const auto& batch : dmat->GetBatches<SparsePage>()) {
+    const auto& row_offset = batch.offset.ConstHostVector();
+    for (auto i = 1ull; i < row_offset.size(); i++) {
+      row_stride = std::max(
+        row_stride, static_cast<size_t>(row_offset[i] - row_offset[i - 1]));
+    }
+  }
+  return row_stride;
+}
+
 // Construct an ELLPACK matrix in memory.
 EllpackPageImpl::EllpackPageImpl(DMatrix* dmat, const BatchParam& param) {
   monitor_.Init("ellpack_page");
@@ -87,13 +101,13 @@ EllpackPageImpl::EllpackPageImpl(DMatrix* dmat, const BatchParam& param) {
 
   monitor_.StartCuda("Quantiles");
   // Create the quantile sketches for the dmatrix and initialize HistogramCuts.
-  common::HistogramCuts hmat;
-  size_t row_stride =
-      common::DeviceSketch(param.gpu_id, param.max_bin, param.gpu_batch_nrows, dmat, &hmat);
+  size_t row_stride = GetRowStride(dmat);
+  auto cuts = common::DeviceSketch(param.gpu_id, dmat, param.max_bin,
+                                   param.gpu_batch_nrows);
   monitor_.StopCuda("Quantiles");
 
   monitor_.StartCuda("InitEllpackInfo");
-  InitInfo(param.gpu_id, dmat->IsDense(), row_stride, hmat);
+  InitInfo(param.gpu_id, dmat->IsDense(), row_stride, cuts);
   monitor_.StopCuda("InitEllpackInfo");
 
   monitor_.StartCuda("InitCompressedData");

src/data/ellpack_page_source.cu

@@ -70,6 +70,20 @@ const EllpackPage& EllpackPageSource::Value() const {
   return impl_->Value();
 }
 
+size_t GetRowStride(DMatrix* dmat) {
+  if (dmat->IsDense()) return dmat->Info().num_col_;
+
+  size_t row_stride = 0;
+  for (const auto& batch : dmat->GetBatches<SparsePage>()) {
+    const auto& row_offset = batch.offset.ConstHostVector();
+    for (auto i = 1ull; i < row_offset.size(); i++) {
+      row_stride = std::max(
+          row_stride, static_cast<size_t>(row_offset[i] - row_offset[i - 1]));
+    }
+  }
+  return row_stride;
+}
+
 // Build the quantile sketch across the whole input data, then use the histogram cuts to compress
 // each CSR page, and write the accumulated ELLPACK pages to disk.
 EllpackPageSourceImpl::EllpackPageSourceImpl(DMatrix* dmat,
@@ -85,13 +99,13 @@ EllpackPageSourceImpl::EllpackPageSourceImpl(DMatrix* dmat,
   dh::safe_cuda(cudaSetDevice(device_));
 
   monitor_.StartCuda("Quantiles");
-  common::HistogramCuts hmat;
-  size_t row_stride =
-      common::DeviceSketch(device_, param.max_bin, param.gpu_batch_nrows, dmat, &hmat);
+  size_t row_stride = GetRowStride(dmat);
+  auto cuts = common::DeviceSketch(param.gpu_id, dmat, param.max_bin,
+                                   param.gpu_batch_nrows);
   monitor_.StopCuda("Quantiles");
 
   monitor_.StartCuda("CreateEllpackInfo");
-  ellpack_info_ = EllpackInfo(device_, dmat->IsDense(), row_stride, hmat, &ba_);
+  ellpack_info_ = EllpackInfo(device_, dmat->IsDense(), row_stride, cuts, &ba_);
   monitor_.StopCuda("CreateEllpackInfo");
 
   monitor_.StartCuda("WriteEllpackPages");

tests/cpp/common/test_hist_util.cc

@@ -261,7 +261,25 @@ TEST(hist_util, DenseCutsAccuracyTest) {
       HistogramCuts cuts;
       DenseCuts dense(&cuts);
       dense.Build(dmat.get(), num_bins);
-      ValidateCuts(cuts, x, num_rows, num_columns, num_bins);
+      ValidateCuts(cuts, dmat.get(), num_bins);
+    }
+  }
+}
+
+TEST(hist_util, DenseCutsAccuracyTestWeights) {
+  int bin_sizes[] = {2, 16, 256, 512};
+  int sizes[] = {100, 1000, 1500};
+  int num_columns = 5;
+  for (auto num_rows : sizes) {
+    auto x = GenerateRandom(num_rows, num_columns);
+    auto dmat = GetDMatrixFromData(x, num_rows, num_columns);
+    auto w = GenerateRandomWeights(num_rows);
+    dmat->Info().weights_.HostVector() = w;
+    for (auto num_bins : bin_sizes) {
+      HistogramCuts cuts;
+      DenseCuts dense(&cuts);
+      dense.Build(dmat.get(), num_bins);
+      ValidateCuts(cuts, dmat.get(), num_bins);
     }
   }
 }
@@ -279,7 +297,7 @@ TEST(hist_util, DenseCutsExternalMemory) {
       HistogramCuts cuts;
       DenseCuts dense(&cuts);
       dense.Build(dmat.get(), num_bins);
-      ValidateCuts(cuts, x, num_rows, num_columns, num_bins);
+      ValidateCuts(cuts, dmat.get(), num_bins);
     }
   }
 }
@@ -295,7 +313,7 @@ TEST(hist_util, SparseCutsAccuracyTest) {
       HistogramCuts cuts;
       SparseCuts sparse(&cuts);
       sparse.Build(dmat.get(), num_bins);
-      ValidateCuts(cuts, x, num_rows, num_columns, num_bins);
+      ValidateCuts(cuts, dmat.get(), num_bins);
     }
   }
 }
@@ -335,7 +353,7 @@ TEST(hist_util, SparseCutsExternalMemory) {
       HistogramCuts cuts;
       SparseCuts dense(&cuts);
       dense.Build(dmat.get(), num_bins);
-      ValidateCuts(cuts, x, num_rows, num_columns, num_bins);
+      ValidateCuts(cuts, dmat.get(), num_bins);
     }
   }
 }