[EM] Return a full DMatrix instead of a Ellpack from the GPU sampler. (dmlc#10753)

Author: trivialfis

.gitignore

@@ -63,6 +63,7 @@ java/xgboost4j-demo/data/
 java/xgboost4j-demo/tmp/
 java/xgboost4j-demo/model/
 nb-configuration*
+
 # Eclipse
 .project
 .cproject
@@ -154,3 +155,6 @@ model*.json
 *.rds
 Rplots.pdf
 *.zip
+
+# nsys
+*.nsys-rep

include/xgboost/data.h

@@ -110,8 +110,15 @@ class MetaInfo {
    * @brief Validate all metainfo.
    */
   void Validate(DeviceOrd device) const;
-
-  MetaInfo Slice(common::Span<int32_t const> ridxs) const;
+  /**
+   * @brief Slice the meta info.
+   *
+   * The device of ridxs is specified by the ctx object.
+   *
+   * @param ridxs Index of selected rows.
+   * @param nnz   The number of non-missing values.
+   */
+  MetaInfo Slice(Context const* ctx, common::Span<bst_idx_t const> ridxs, bst_idx_t nnz) const;
 
   MetaInfo Copy() const;
   /**

src/common/device_helpers.cuh

@@ -508,6 +508,11 @@ xgboost::common::Span<T> ToSpan(DeviceUVector<T> &vec) {
   return {vec.data(), vec.size()};
 }
 
+template <typename T>
+xgboost::common::Span<std::add_const_t<T>> ToSpan(DeviceUVector<T> const &vec) {
+  return {vec.data(), vec.size()};
+}
+
 // thrust begin, similiar to std::begin
 template <typename T>
 thrust::device_ptr<T> tbegin(xgboost::HostDeviceVector<T>& vector) {  // NOLINT

src/common/linalg_op.cuh

@@ -76,7 +76,7 @@ struct IterOp {
 // returns a thrust iterator for a tensor view.
 template <typename T, std::int32_t kDim>
 auto tcbegin(TensorView<T, kDim> v) {  // NOLINT
-  return dh::MakeTransformIterator<T>(
+  return thrust::make_transform_iterator(
       thrust::make_counting_iterator(0ul),
       detail::IterOp<std::add_const_t<std::remove_const_t<T>>, kDim>{v});
 }
@@ -85,5 +85,16 @@ template <typename T, std::int32_t kDim>
 auto tcend(TensorView<T, kDim> v) {  // NOLINT
   return tcbegin(v) + v.Size();
 }
+
+template <typename T, std::int32_t kDim>
+auto tbegin(TensorView<T, kDim> v) {  // NOLINT
+  return thrust::make_transform_iterator(thrust::make_counting_iterator(0ul),
+                                         detail::IterOp<std::remove_const_t<T>, kDim>{v});
+}
+
+template <typename T, std::int32_t kDim>
+auto tend(TensorView<T, kDim> v) {  // NOLINT
+  return tbegin(v) + v.Size();
+}
 }  // namespace xgboost::linalg
 #endif  // XGBOOST_COMMON_LINALG_OP_CUH_

src/data/data.cc

@@ -351,8 +351,10 @@ void MetaInfo::LoadBinary(dmlc::Stream *fi) {
   this->has_categorical_ = LoadFeatureType(feature_type_names, &feature_types.HostVector());
 }
 
+namespace {
 template <typename T>
-std::vector<T> Gather(const std::vector<T> &in, common::Span<int const> ridxs, size_t stride = 1) {
+std::vector<T> Gather(const std::vector<T>& in, common::Span<bst_idx_t const> ridxs,
+                      size_t stride = 1) {
   if (in.empty()) {
     return {};
   }
@@ -361,16 +363,56 @@ std::vector<T> Gather(const std::vector<T> &in, common::Span<int const> ridxs, s
   for (auto i = 0ull; i < size; i++) {
     auto ridx = ridxs[i];
     for (size_t j = 0; j < stride; ++j) {
-      out[i * stride +j] = in[ridx * stride + j];
+      out[i * stride + j] = in[ridx * stride + j];
     }
   }
   return out;
 }
+}  // namespace
+
+namespace cuda_impl {
+void SliceMetaInfo(Context const* ctx, MetaInfo const& info, common::Span<bst_idx_t const> ridx,
+                   MetaInfo* p_out);
+#if !defined(XGBOOST_USE_CUDA)
+void SliceMetaInfo(Context const*, MetaInfo const&, common::Span<bst_idx_t const>, MetaInfo*) {
+  common::AssertGPUSupport();
+}
+#endif
+}  // namespace cuda_impl
 
-MetaInfo MetaInfo::Slice(common::Span<int32_t const> ridxs) const {
+MetaInfo MetaInfo::Slice(Context const* ctx, common::Span<bst_idx_t const> ridxs,
+                         bst_idx_t nnz) const {
+  /**
+   * Shape
+   */
   MetaInfo out;
   out.num_row_ = ridxs.size();
   out.num_col_ = this->num_col_;
+  out.num_nonzero_ = nnz;
+
+  /**
+   * Feature Info
+   */
+  out.feature_weights.SetDevice(ctx->Device());
+  out.feature_weights.Resize(this->feature_weights.Size());
+  out.feature_weights.Copy(this->feature_weights);
+
+  out.feature_names = this->feature_names;
+
+  out.feature_types.SetDevice(ctx->Device());
+  out.feature_types.Resize(this->feature_types.Size());
+  out.feature_types.Copy(this->feature_types);
+
+  out.feature_type_names = this->feature_type_names;
+
+  /**
+   * Sample Info
+   */
+  if (ctx->IsCUDA()) {
+    cuda_impl::SliceMetaInfo(ctx, *this, ridxs, &out);
+    return out;
+  }
+
   // Groups is maintained by a higher level Python function.  We should aim at deprecating
   // the slice function.
   if (this->labels.Size() != this->num_row_) {
@@ -386,13 +428,11 @@ MetaInfo MetaInfo::Slice(common::Span<int32_t const> ridxs) const {
     });
   }
 
-  out.labels_upper_bound_.HostVector() =
-      Gather(this->labels_upper_bound_.HostVector(), ridxs);
-  out.labels_lower_bound_.HostVector() =
-      Gather(this->labels_lower_bound_.HostVector(), ridxs);
+  out.labels_upper_bound_.HostVector() = Gather(this->labels_upper_bound_.HostVector(), ridxs);
+  out.labels_lower_bound_.HostVector() = Gather(this->labels_lower_bound_.HostVector(), ridxs);
   // weights
   if (this->weights_.Size() + 1 == this->group_ptr_.size()) {
-    auto& h_weights =  out.weights_.HostVector();
+    auto& h_weights = out.weights_.HostVector();
     // Assuming all groups are available.
     out.weights_.HostVector() = h_weights;
   } else {
@@ -414,14 +454,6 @@ MetaInfo MetaInfo::Slice(common::Span<int32_t const> ridxs) const {
     });
   }
 
-  out.feature_weights.Resize(this->feature_weights.Size());
-  out.feature_weights.Copy(this->feature_weights);
-
-  out.feature_names = this->feature_names;
-  out.feature_types.Resize(this->feature_types.Size());
-  out.feature_types.Copy(this->feature_types);
-  out.feature_type_names = this->feature_type_names;
-
   return out;
 }
 

src/data/data.cu

@@ -1,9 +1,11 @@
 /**
- * Copyright 2019-2022 by XGBoost Contributors
+ * Copyright 2019-2024, XGBoost Contributors
  *
  * \file data.cu
  * \brief Handles setting metainfo from array interface.
  */
+#include <thrust/gather.h>  // for gather
+
 #include "../common/cuda_context.cuh"
 #include "../common/device_helpers.cuh"
 #include "../common/linalg_op.cuh"
@@ -169,6 +171,62 @@ void MetaInfo::SetInfoFromCUDA(Context const& ctx, StringView key, Json array) {
   }
 }
 
+namespace {
+void Gather(Context const* ctx, linalg::MatrixView<float const> in,
+            common::Span<bst_idx_t const> ridx, linalg::Matrix<float>* p_out) {
+  if (in.Empty()) {
+    return;
+  }
+  auto& out = *p_out;
+  out.Reshape(ridx.size(), in.Shape(1));
+  auto d_out = out.View(ctx->Device());
+
+  auto cuctx = ctx->CUDACtx();
+  auto map_it = thrust::make_transform_iterator(thrust::make_counting_iterator(0ull),
+                                                [=] XGBOOST_DEVICE(bst_idx_t i) {
+                                                  auto [r, c] = linalg::UnravelIndex(i, in.Shape());
+                                                  return (ridx[r] * in.Shape(1)) + c;
+                                                });
+  CHECK_NE(in.Shape(1), 0);
+  thrust::gather(cuctx->TP(), map_it, map_it + out.Size(), linalg::tcbegin(in),
+                 linalg::tbegin(d_out));
+}
+
+template <typename T>
+void Gather(Context const* ctx, HostDeviceVector<T> const& in, common::Span<bst_idx_t const> ridx,
+            HostDeviceVector<T>* p_out) {
+  if (in.Empty()) {
+    return;
+  }
+  in.SetDevice(ctx->Device());
+
+  auto& out = *p_out;
+  out.SetDevice(ctx->Device());
+  out.Resize(ridx.size());
+  auto d_out = out.DeviceSpan();
+
+  auto cuctx = ctx->CUDACtx();
+  auto d_in = in.ConstDeviceSpan();
+  thrust::gather(cuctx->TP(), dh::tcbegin(ridx), dh::tcend(ridx), dh::tcbegin(d_in),
+                 dh::tbegin(d_out));
+}
+}  // anonymous namespace
+
+namespace cuda_impl {
+void SliceMetaInfo(Context const* ctx, MetaInfo const& info, common::Span<bst_idx_t const> ridx,
+                   MetaInfo* p_out) {
+  auto& out = *p_out;
+
+  Gather(ctx, info.labels.View(ctx->Device()), ridx, &p_out->labels);
+  Gather(ctx, info.base_margin_.View(ctx->Device()), ridx, &p_out->base_margin_);
+
+  Gather(ctx, info.labels_lower_bound_, ridx, &out.labels_lower_bound_);
+  Gather(ctx, info.labels_upper_bound_, ridx, &out.labels_upper_bound_);
+
+  Gather(ctx, info.weights_, ridx, &out.weights_);
+}
+}  // namespace cuda_impl
+
 template <typename AdapterT>
 DMatrix* DMatrix::Create(AdapterT* adapter, float missing, int nthread,
                          const std::string& cache_prefix, DataSplitMode data_split_mode) {

src/data/ellpack_page.cu

@@ -1,12 +1,13 @@
 /**
  * Copyright 2019-2024, XGBoost contributors
  */
+#include <cuda/functional>  // for proclaim_return_type
 #include <thrust/iterator/discard_iterator.h>
 #include <thrust/iterator/transform_output_iterator.h>
 
-#include <algorithm>  // for copy
-#include <utility>    // for move
-#include <vector>     // for vector
+#include <algorithm>        // for copy
+#include <utility>          // for move
+#include <vector>           // for vector
 
 #include "../common/categorical.h"
 #include "../common/cuda_context.cuh"
@@ -576,4 +577,17 @@ EllpackDeviceAccessor EllpackPageImpl::GetHostAccessor(
           common::CompressedIterator<uint32_t>(h_gidx_buffer->data(), NumSymbols()),
           feature_types};
 }
+
+[[nodiscard]] bst_idx_t EllpackPageImpl::NumNonMissing(
+    Context const* ctx, common::Span<FeatureType const> feature_types) const {
+  auto d_acc = this->GetDeviceAccessor(ctx->Device(), feature_types);
+  using T = typename decltype(d_acc.gidx_iter)::value_type;
+  auto it = thrust::make_transform_iterator(
+      thrust::make_counting_iterator(0ull),
+      cuda::proclaim_return_type<T>([=] __device__(std::size_t i) { return d_acc.gidx_iter[i]; }));
+  auto nnz = thrust::count_if(ctx->CUDACtx()->CTP(), it, it + d_acc.row_stride * d_acc.n_rows,
+                              cuda::proclaim_return_type<bool>(
+                                  [=] __device__(T gidx) { return gidx != d_acc.NullValue(); }));
+  return nnz;
+}
 }  // namespace xgboost

src/data/ellpack_page.cuh

@@ -236,6 +236,11 @@ class EllpackPageImpl {
   [[nodiscard]] EllpackDeviceAccessor GetHostAccessor(
       Context const* ctx, std::vector<common::CompressedByteT>* h_gidx_buffer,
       common::Span<FeatureType const> feature_types = {}) const;
+  /**
+   * @brief Calculate the number of non-missing values.
+   */
+  [[nodiscard]] bst_idx_t NumNonMissing(Context const* ctx,
+                                        common::Span<FeatureType const> feature_types) const;
 
  private:
   /**

src/data/iterative_dmatrix.cu

@@ -101,6 +101,17 @@ void IterativeDMatrix::InitFromCUDA(Context const* ctx, BatchParam const& p,
   // Synchronise worker columns
 }
 
+IterativeDMatrix::IterativeDMatrix(std::shared_ptr<EllpackPage> ellpack, MetaInfo const& info,
+                                   BatchParam batch) {
+  this->ellpack_ = ellpack;
+  CHECK_EQ(this->Info().num_row_, 0);
+  CHECK_EQ(this->Info().num_col_, 0);
+  this->Info().Extend(info, true, true);
+  this->Info().num_nonzero_ = info.num_nonzero_;
+  CHECK_EQ(this->Info().num_row_, info.num_row_);
+  this->batch_ = batch;
+}
+
 BatchSet<EllpackPage> IterativeDMatrix::GetEllpackBatches(Context const* ctx,
                                                           BatchParam const& param) {
   if (param.Initialized()) {

src/data/iterative_dmatrix.h

@@ -48,6 +48,11 @@ class IterativeDMatrix : public QuantileDMatrix {
                             std::shared_ptr<DMatrix> ref, DataIterResetCallback *reset,
                             XGDMatrixCallbackNext *next, float missing, int nthread,
                             bst_bin_t max_bin);
+  /**
+   * @param Directly construct a QDM from an existing one.
+   */
+  IterativeDMatrix(std::shared_ptr<EllpackPage> ellpack, MetaInfo const &info, BatchParam batch);
+
   ~IterativeDMatrix() override = default;
 
   bool EllpackExists() const override { return static_cast<bool>(ellpack_); }