Cleanup.

dmlc · Sep 18, 2024 · 8320640 · 8320640
1 parent 1158014
commit 8320640
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 15 deletions.
diff --git a/src/data/ellpack_page.cu b/src/data/ellpack_page.cu
@@ -406,7 +406,10 @@ EllpackPageImpl::EllpackPageImpl(Context const* ctx, GHistIndexMatrix const& pag
  this->monitor_.Stop("CopyGHistToEllpack");
 }
 
-EllpackPageImpl::~EllpackPageImpl() noexcept(false) { dh::DefaultStream().Sync(); }
+EllpackPageImpl::~EllpackPageImpl() noexcept(false) {
+ // Sync the stream to make sure all running CUDA kernels finish before deallocation.
+ dh::DefaultStream().Sync();
+}
 
 // A functor that copies the data from one EllpackPage to another.
 struct CopyPage {
@@ -420,15 +423,15 @@ struct CopyPage {
  : cbw{dst->NumSymbols()},
  dst_data_d{dst->gidx_buffer.data()},
  src_iterator_d{src->gidx_buffer.data(), src->NumSymbols()},
- offset(offset) {}
+ offset{offset} {}
 
  __device__ void operator()(size_t element_id) {
  cbw.AtomicWriteSymbol(dst_data_d, src_iterator_d[element_id], element_id + offset);
  }
 };
 
 // Copy the data from the given EllpackPage to the current page.
-size_t EllpackPageImpl::Copy(Context const* ctx, EllpackPageImpl const* page, bst_idx_t offset) {
+bst_idx_t EllpackPageImpl::Copy(Context const* ctx, EllpackPageImpl const* page, bst_idx_t offset) {
  monitor_.Start(__func__);
  bst_idx_t num_elements = page->n_rows * page->row_stride;
  CHECK_EQ(this->row_stride, page->row_stride);

diff --git a/src/data/extmem_quantile_dmatrix.cc b/src/data/extmem_quantile_dmatrix.cc
@@ -55,14 +55,6 @@ ExtMemQuantileDMatrix::~ExtMemQuantileDMatrix() {
  DeleteCacheFiles(cache_info_);
 }
 
-[[nodiscard]] bst_idx_t ExtMemQuantileDMatrix::EllpackPageFetchCount() const {
- return std::visit([](auto &&ptr) { return ptr->FetchCount(); }, this->ellpack_page_source_);
-}
-
-[[nodiscard]] bst_idx_t ExtMemQuantileDMatrix::GradientIndexFetchCount() const {
- return this->ghist_index_source_->FetchCount();
-}
-
 BatchSet<ExtSparsePage> ExtMemQuantileDMatrix::GetExtBatches(Context const *, BatchParam const &) {
  LOG(FATAL) << "Not implemented";
  auto begin_iter =

diff --git a/src/data/extmem_quantile_dmatrix.h b/src/data/extmem_quantile_dmatrix.h
@@ -35,10 +35,6 @@ class ExtMemQuantileDMatrix : public QuantileDMatrix {
 
  [[nodiscard]] std::int32_t NumBatches() const override { return n_batches_; }
 
- // For testing, getter for the number of fetches.
- [[nodiscard]] bst_idx_t EllpackPageFetchCount() const;
- [[nodiscard]] bst_idx_t GradientIndexFetchCount() const;
-
  private:
  void InitFromCPU(
  Context const *ctx,