[EM] Add GPU version of the external memory QDM. (dmlc#10689)

Author: trivialfis

include/xgboost/data.h

@@ -641,7 +641,7 @@ class DMatrix {
             typename XGDMatrixCallbackNext>
   static DMatrix* Create(DataIterHandle iter, DMatrixHandle proxy, std::shared_ptr<DMatrix> ref,
                          DataIterResetCallback* reset, XGDMatrixCallbackNext* next, float missing,
-                         std::int32_t nthread, bst_bin_t max_bin, std::string cache);
+                         std::int32_t nthread, bst_bin_t max_bin, std::string cache, bool on_host);
 
   virtual DMatrix *Slice(common::Span<int32_t const> ridxs) = 0;
 

src/common/device_helpers.cuh

@@ -116,6 +116,13 @@ inline int32_t CurrentDevice() {
   return device;
 }
 
+// Helper function to get a device from a potentially CPU context.
+inline auto GetDevice(xgboost::Context const *ctx) {
+  auto d = (ctx->IsCUDA()) ? ctx->Device() : xgboost::DeviceOrd::CUDA(dh::CurrentDevice());
+  CHECK(!d.IsCPU());
+  return d;
+}
+
 inline size_t TotalMemory(int device_idx) {
   size_t device_free = 0;
   size_t device_total = 0;

src/data/data.cc

@@ -914,9 +914,9 @@ template <typename DataIterHandle, typename DMatrixHandle, typename DataIterRese
           typename XGDMatrixCallbackNext>
 DMatrix* DMatrix::Create(DataIterHandle iter, DMatrixHandle proxy, std::shared_ptr<DMatrix> ref,
                          DataIterResetCallback* reset, XGDMatrixCallbackNext* next, float missing,
-                         std::int32_t nthread, bst_bin_t max_bin, std::string cache) {
+                         std::int32_t nthread, bst_bin_t max_bin, std::string cache, bool on_host) {
   return new data::ExtMemQuantileDMatrix{
-      iter, proxy, ref, reset, next, missing, nthread, std::move(cache), max_bin};
+      iter, proxy, ref, reset, next, missing, nthread, std::move(cache), max_bin, on_host};
 }
 
 template DMatrix* DMatrix::Create<DataIterHandle, DMatrixHandle, DataIterResetCallback,
@@ -935,7 +935,7 @@ template DMatrix* DMatrix::Create<DataIterHandle, DMatrixHandle, DataIterResetCa
 template DMatrix*
 DMatrix::Create<DataIterHandle, DMatrixHandle, DataIterResetCallback, XGDMatrixCallbackNext>(
     DataIterHandle, DMatrixHandle, std::shared_ptr<DMatrix>, DataIterResetCallback*,
-    XGDMatrixCallbackNext*, float, std::int32_t, bst_bin_t, std::string);
+    XGDMatrixCallbackNext*, float, std::int32_t, bst_bin_t, std::string, bool);
 
 template <typename AdapterT>
 DMatrix* DMatrix::Create(AdapterT* adapter, float missing, int nthread, const std::string&,

src/data/ellpack_page.cc

@@ -47,6 +47,18 @@ bst_idx_t EllpackPage::Size() const {
                 "EllpackPage is required";
   return impl_->Cuts();
 }
+
+[[nodiscard]] bst_idx_t EllpackPage::BaseRowId() const {
+  LOG(FATAL) << "Internal Error: XGBoost is not compiled with CUDA but "
+                "EllpackPage is required";
+  return 0;
+}
+
+[[nodiscard]] bool EllpackPage::IsDense() const {
+  LOG(FATAL) << "Internal Error: XGBoost is not compiled with CUDA but "
+                "EllpackPage is required";
+  return false;
+}
 }  // namespace xgboost
 
 #endif  // XGBOOST_USE_CUDA

src/data/ellpack_page.cu

@@ -39,6 +39,9 @@ void EllpackPage::SetBaseRowId(std::size_t row_id) { impl_->SetBaseRowId(row_id)
   return impl_->Cuts();
 }
 
+[[nodiscard]] bst_idx_t EllpackPage::BaseRowId() const { return this->Impl()->base_rowid; }
+[[nodiscard]] bool EllpackPage::IsDense() const { return this->Impl()->IsDense(); }
+
 // Bin each input data entry, store the bin indices in compressed form.
 __global__ void CompressBinEllpackKernel(
     common::CompressedBufferWriter wr,
@@ -397,7 +400,7 @@ struct CopyPage {
 size_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(row_stride, page->row_stride);
+  CHECK_EQ(this->row_stride, page->row_stride);
   CHECK_EQ(NumSymbols(), page->NumSymbols());
   CHECK_GE(n_rows * row_stride, offset + num_elements);
   if (page == this) {

src/data/ellpack_page.cuh

@@ -203,6 +203,7 @@ class EllpackPageImpl {
   [[nodiscard]] std::shared_ptr<common::HistogramCuts const> CutsShared() const { return cuts_; }
   void SetCuts(std::shared_ptr<common::HistogramCuts const> cuts) { cuts_ = cuts; }
 
+  [[nodiscard]] bool IsDense() const { return is_dense; }
   /** @return Estimation of memory cost of this page. */
   static size_t MemCostBytes(size_t num_rows, size_t row_stride, const common::HistogramCuts&cuts) ;
 

src/data/ellpack_page.h

@@ -42,6 +42,7 @@ class EllpackPage {
 
   /*! \return Number of instances in the page. */
   [[nodiscard]] bst_idx_t Size() const;
+  [[nodiscard]] bool IsDense() const;
 
   /*! \brief Set the base row id for this page. */
   void SetBaseRowId(std::size_t row_id);
@@ -50,6 +51,7 @@ class EllpackPage {
   EllpackPageImpl* Impl() { return impl_.get(); }
 
   [[nodiscard]] common::HistogramCuts const& Cuts() const;
+  [[nodiscard]] bst_idx_t BaseRowId() const;
 
  private:
   std::unique_ptr<EllpackPageImpl> impl_;

src/data/ellpack_page_source.cu

@@ -15,7 +15,8 @@
 #include "ellpack_page.cuh"                   // for EllpackPageImpl
 #include "ellpack_page.h"                     // for EllpackPage
 #include "ellpack_page_source.h"
-#include "xgboost/base.h"  // for bst_idx_t
+#include "proxy_dmatrix.cuh"  // for Dispatch
+#include "xgboost/base.h"     // for bst_idx_t
 
 namespace xgboost::data {
 struct EllpackHostCache {
@@ -182,4 +183,51 @@ template void
 EllpackPageSourceImpl<EllpackCacheStreamPolicy<EllpackPage, EllpackFormatPolicy>>::Fetch();
 template void
 EllpackPageSourceImpl<EllpackMmapStreamPolicy<EllpackPage, EllpackFormatPolicy>>::Fetch();
+
+/**
+ * ExtEllpackPageSourceImpl
+ */
+template <typename F>
+void ExtEllpackPageSourceImpl<F>::Fetch() {
+  dh::safe_cuda(cudaSetDevice(this->Device().ordinal));
+  if (!this->ReadCache()) {
+    auto iter = this->source_->Iter();
+    CHECK_EQ(this->count_, iter);
+    ++(*this->source_);
+    CHECK_GE(this->source_->Iter(), 1);
+    cuda_impl::Dispatch(proxy_, [this](auto const& value) {
+      proxy_->Info().feature_types.SetDevice(dh::GetDevice(this->ctx_));
+      auto d_feature_types = proxy_->Info().feature_types.ConstDeviceSpan();
+      auto n_samples = value.NumRows();
+
+      dh::device_vector<size_t> row_counts(n_samples + 1, 0);
+      common::Span<size_t> row_counts_span(row_counts.data().get(), row_counts.size());
+      cuda_impl::Dispatch(proxy_, [=](auto const& value) {
+        return GetRowCounts(value, row_counts_span, dh::GetDevice(this->ctx_), this->missing_);
+      });
+
+      this->page_.reset(new EllpackPage{});
+      *this->page_->Impl() = EllpackPageImpl{this->ctx_,
+                                             value,
+                                             this->missing_,
+                                             this->info_->IsDense(),
+                                             row_counts_span,
+                                             d_feature_types,
+                                             this->ext_info_.row_stride,
+                                             n_samples,
+                                             this->GetCuts()};
+      this->info_->Extend(proxy_->Info(), false, true);
+    });
+    this->page_->SetBaseRowId(this->ext_info_.base_rows.at(iter));
+    this->WriteCache();
+  }
+}
+
+// Instantiation
+template void
+ExtEllpackPageSourceImpl<DefaultFormatStreamPolicy<EllpackPage, EllpackFormatPolicy>>::Fetch();
+template void
+ExtEllpackPageSourceImpl<EllpackCacheStreamPolicy<EllpackPage, EllpackFormatPolicy>>::Fetch();
+template void
+ExtEllpackPageSourceImpl<EllpackMmapStreamPolicy<EllpackPage, EllpackFormatPolicy>>::Fetch();
 }  // namespace xgboost::data

src/data/ellpack_page_source.h

@@ -8,6 +8,7 @@
 #include <cstdint>  // for int32_t
 #include <memory>   // for shared_ptr
 #include <utility>  // for move
+#include <vector>   // for vector
 
 #include "../common/cuda_rt_utils.h"  // for SupportsPageableMem
 #include "../common/hist_util.h"      // for HistogramCuts
@@ -169,6 +170,51 @@ using EllpackPageHostSource =
 using EllpackPageSource =
     EllpackPageSourceImpl<EllpackMmapStreamPolicy<EllpackPage, EllpackFormatPolicy>>;
 
+template <typename FormatCreatePolicy>
+class ExtEllpackPageSourceImpl : public ExtQantileSourceMixin<EllpackPage, FormatCreatePolicy> {
+  using Super = ExtQantileSourceMixin<EllpackPage, FormatCreatePolicy>;
+
+  Context const* ctx_;
+  BatchParam p_;
+  DMatrixProxy* proxy_;
+  MetaInfo* info_;
+  ExternalDataInfo ext_info_;
+
+  std::vector<bst_idx_t> base_rows_;
+
+ public:
+  ExtEllpackPageSourceImpl(
+      Context const* ctx, float missing, MetaInfo* info, ExternalDataInfo ext_info,
+      std::shared_ptr<Cache> cache, BatchParam param, std::shared_ptr<common::HistogramCuts> cuts,
+      std::shared_ptr<DataIterProxy<DataIterResetCallback, XGDMatrixCallbackNext>> source,
+      DMatrixProxy* proxy, std::vector<bst_idx_t> base_rows)
+      : Super{missing,
+              ctx->Threads(),
+              static_cast<bst_feature_t>(info->num_col_),
+              ext_info.n_batches,
+              source,
+              cache},
+        ctx_{ctx},
+        p_{std::move(param)},
+        proxy_{proxy},
+        info_{info},
+        ext_info_{std::move(ext_info)},
+        base_rows_{std::move(base_rows)} {
+    this->SetCuts(std::move(cuts), ctx->Device());
+    this->Fetch();
+  }
+
+  void Fetch() final;
+};
+
+// Cache to host
+using ExtEllpackPageHostSource =
+    ExtEllpackPageSourceImpl<EllpackCacheStreamPolicy<EllpackPage, EllpackFormatPolicy>>;
+
+// Cache to disk
+using ExtEllpackPageSource =
+    ExtEllpackPageSourceImpl<EllpackMmapStreamPolicy<EllpackPage, EllpackFormatPolicy>>;
+
 #if !defined(XGBOOST_USE_CUDA)
 template <typename F>
 inline void EllpackPageSourceImpl<F>::Fetch() {
@@ -177,6 +223,11 @@ inline void EllpackPageSourceImpl<F>::Fetch() {
   (void)(is_dense_);
   common::AssertGPUSupport();
 }
+
+template <typename F>
+inline void ExtEllpackPageSourceImpl<F>::Fetch() {
+  common::AssertGPUSupport();
+}
 #endif  // !defined(XGBOOST_USE_CUDA)
 }  // namespace xgboost::data
 

src/data/extmem_quantile_dmatrix.cc

@@ -24,8 +24,8 @@ ExtMemQuantileDMatrix::ExtMemQuantileDMatrix(DataIterHandle iter_handle, DMatrix
                                              DataIterResetCallback *reset,
                                              XGDMatrixCallbackNext *next, float missing,
                                              std::int32_t n_threads, std::string cache,
-                                             bst_bin_t max_bin)
-    : cache_prefix_{std::move(cache)} {
+                                             bst_bin_t max_bin, bool on_host)
+    : cache_prefix_{std::move(cache)}, on_host_{on_host} {
   auto iter = std::make_shared<DataIterProxy<DataIterResetCallback, XGDMatrixCallbackNext>>(
       iter_handle, reset, next);
   iter->Reset();
@@ -72,13 +72,7 @@ void ExtMemQuantileDMatrix::InitFromCPU(
   common::HistogramCuts cuts;
   ExternalDataInfo ext_info;
   cpu_impl::GetDataShape(ctx, proxy, *iter, missing, &ext_info);
-
-  // From here on Info() has the correct data shape
-  this->Info().num_row_ = ext_info.accumulated_rows;
-  this->Info().num_col_ = ext_info.n_features;
-  this->Info().num_nonzero_ = ext_info.nnz;
-  this->Info().SynchronizeNumberOfColumns(ctx);
-  ext_info.Validate();
+  ext_info.SetInfo(ctx, &this->info_);
 
   /**
    * Generate quantiles
@@ -110,7 +104,7 @@ void ExtMemQuantileDMatrix::InitFromCPU(
   CHECK_EQ(n_total_samples, ext_info.accumulated_rows);
 }
 
-BatchSet<GHistIndexMatrix> ExtMemQuantileDMatrix::GetGradientIndexImpl() {
+[[nodiscard]] BatchSet<GHistIndexMatrix> ExtMemQuantileDMatrix::GetGradientIndexImpl() {
   return BatchSet{BatchIterator<GHistIndexMatrix>{this->ghist_index_source_}};
 }
 
@@ -148,5 +142,13 @@ BatchSet<EllpackPage> ExtMemQuantileDMatrix::GetEllpackBatches(Context const *,
                  this->ellpack_page_source_);
   return batch_set;
 }
+
+BatchSet<EllpackPage> ExtMemQuantileDMatrix::GetEllpackPageImpl() {
+  common::AssertGPUSupport();
+  auto batch_set =
+      std::visit([this](auto &&ptr) { return BatchSet{BatchIterator<EllpackPage>{ptr}}; },
+                 this->ellpack_page_source_);
+  return batch_set;
+}
 #endif
 }  // namespace xgboost::data