[EM] Improve memory estimation for quantile sketching. (#10843)

I- Add basic estimation for RMM.
- Re-estimate after every sub-batch.
- Some debug logs for memory usage.
- Fix the locking mechanism in the memory allocator logger.
This commit is contained in:
Jiaming Yuan
2024-09-25 03:20:09 +08:00
committed by GitHub
parent f3df0d0eb4
commit bc69a3e877
6 changed files with 196 additions and 104 deletions

View File

@@ -65,7 +65,9 @@ TEST(HistUtil, SketchBatchNumElements) {
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);
auto shape = detail::SketchShape{rows, kCols, rows * kCols};
auto batch = detail::SketchBatchNumElements(detail::UnknownSketchNumElements(), shape, device,
256, false, 0);
ASSERT_EQ(batch, avail_elem);
}