Fix sketch size calculation.

src/common/hist_util.cu

@@ -162,13 +162,18 @@ size_t RequiredMemory(bst_row_t num_rows, bst_feature_t num_columns, size_t nnz,
 }
 
 size_t SketchBatchNumElements(size_t sketch_batch_num_elements,
- bst_row_t num_rows, size_t columns, size_t nnz, int device,
+ bst_row_t num_rows, bst_feature_t columns,
+ size_t nnz, int device,
  size_t num_cuts, bool has_weight) {
  if (sketch_batch_num_elements == 0) {
  auto required_memory = RequiredMemory(num_rows, columns, nnz, num_cuts, has_weight);
  // use up to 80% of available space
- sketch_batch_num_elements = (dh::AvailableMemory(device) -
- required_memory * 0.8);
+ auto avail = dh::AvailableMemory(device) * 0.8;
+ if (required_memory > avail) {
+ sketch_batch_num_elements = avail / BytesPerElement(has_weight);
+ } else {
+ sketch_batch_num_elements = std::min(num_rows * static_cast<size_t>(columns), nnz);
+ }
  }
  return sketch_batch_num_elements;
 }
@@ -196,7 +201,7 @@ void ProcessBatch(int device, const SparsePage &page, size_t begin, size_t end,
  size_t num_columns) {
  dh::XGBCachingDeviceAllocator<char> alloc;
  const auto& host_data = page.data.ConstHostVector();
- dh::caching_device_vector<Entry> sorted_entries(host_data.begin() + begin,
+ dh::device_vector<Entry> sorted_entries(host_data.begin() + begin,
  host_data.begin() + end);
  thrust::sort(thrust::cuda::par(alloc), sorted_entries.begin(),
  sorted_entries.end(), detail::EntryCompareOp());

src/common/hist_util.cuh

@@ -100,7 +100,8 @@ inline size_t constexpr BytesPerElement(bool has_weight) {
  * directly if it's not 0.
  */
 size_t SketchBatchNumElements(size_t sketch_batch_num_elements,
- bst_row_t num_rows, size_t columns, size_t nnz, int device,
+ bst_row_t num_rows, bst_feature_t columns,
+ size_t nnz, int device,
  size_t num_cuts, bool has_weight);
 
 // Compute number of sample cuts needed on local node to maintain accuracy

tests/cpp/common/test_hist_util.cu

@@ -48,6 +48,18 @@ TEST(HistUtil, DeviceSketch) {
  EXPECT_EQ(device_cuts.MinValues(), host_cuts.MinValues());
 }
 
+TEST(HistUtil, SketchBatchNumElements) {
+ size_t constexpr kCols = 10000;
+ int device;
+ dh::safe_cuda(cudaGetDevice(&device));
+ auto avail = static_cast<size_t>(dh::AvailableMemory(device) * 0.8);
+ auto per_elem = detail::BytesPerElement(false);
+ auto avail_elem = avail / per_elem;
+ size_t rows = avail_elem / kCols * 10;
+ auto batch = detail::SketchBatchNumElements(0, rows, kCols, rows * kCols, device, 256, false);
+ ASSERT_EQ(batch, avail_elem);
+}
+
 TEST(HistUtil, DeviceSketchMemory) {
  int num_columns = 100;
  int num_rows = 1000;