[EM] Avoid writing cut matrix to cache. (dmlc#10444)

Author: trivialfis

include/xgboost/data.h

@@ -113,7 +113,10 @@ class MetaInfo {
   MetaInfo Slice(common::Span<int32_t const> ridxs) const;
 
   MetaInfo Copy() const;
-
+  /**
+   * @brief Whether the matrix is dense.
+   */
+  bool IsDense() const { return num_col_ * num_row_ == num_nonzero_; }
   /*!
    * \brief Get weight of each instances.
    * \param i Instance index.
@@ -538,10 +541,10 @@ class DMatrix {
   /*! \brief virtual destructor */
   virtual ~DMatrix();
 
-  /*! \brief Whether the matrix is dense. */
-  [[nodiscard]] bool IsDense() const {
-    return Info().num_nonzero_ == Info().num_row_ * Info().num_col_;
-  }
+  /**
+   * @brief Whether the matrix is dense.
+   */
+  [[nodiscard]] bool IsDense() const { return this->Info().IsDense(); }
 
   /**
    * \brief Load DMatrix from URI.

src/collective/aggregator.h

@@ -9,7 +9,6 @@
 #include <limits>
 #include <string>
 #include <utility>
-#include <vector>
 
 #include "allreduce.h"
 #include "broadcast.h"

src/common/hist_util.h

@@ -162,6 +162,17 @@ class HistogramCuts {
     }
     return vals[bin_idx - 1];
   }
+
+  void SetDevice(DeviceOrd d) const {
+    this->cut_ptrs_.SetDevice(d);
+    this->cut_ptrs_.ConstDevicePointer();
+
+    this->cut_values_.SetDevice(d);
+    this->cut_values_.ConstDevicePointer();
+
+    this->min_vals_.SetDevice(d);
+    this->min_vals_.ConstDevicePointer();
+  }
 };
 
 /**

src/data/ellpack_page.cc

@@ -1,21 +1,23 @@
 /**
- * Copyright 2019-2023, XGBoost contributors
+ * Copyright 2019-2024, XGBoost contributors
  */
 #ifndef XGBOOST_USE_CUDA
 
 #include "ellpack_page.h"
 
 #include <xgboost/data.h>
 
+#include <memory>  // for shared_ptr
+
 // dummy implementation of EllpackPage in case CUDA is not used
 namespace xgboost {
 
 class EllpackPageImpl {
-  common::HistogramCuts cuts_;
+  std::shared_ptr<common::HistogramCuts> cuts_;
 
  public:
-  [[nodiscard]] common::HistogramCuts& Cuts() { return cuts_; }
-  [[nodiscard]] common::HistogramCuts const& Cuts() const { return cuts_; }
+  [[nodiscard]] common::HistogramCuts const& Cuts() const { return *cuts_; }
+  [[nodiscard]] std::shared_ptr<common::HistogramCuts const> CutsShared() const { return cuts_; }
 };
 
 EllpackPage::EllpackPage() = default;
@@ -40,12 +42,6 @@ size_t EllpackPage::Size() const {
   return 0;
 }
 
-[[nodiscard]] common::HistogramCuts& EllpackPage::Cuts() {
-  LOG(FATAL) << "Internal Error: XGBoost is not compiled with CUDA but "
-                "EllpackPage is required";
-  return impl_->Cuts();
-}
-
 [[nodiscard]] common::HistogramCuts const& EllpackPage::Cuts() const {
   LOG(FATAL) << "Internal Error: XGBoost is not compiled with CUDA but "
                 "EllpackPage is required";

src/data/ellpack_page.cu

@@ -12,8 +12,8 @@
 #include "../common/cuda_context.cuh"
 #include "../common/hist_util.cuh"
 #include "../common/transform_iterator.h"  // MakeIndexTransformIter
-#include "./ellpack_page.cuh"
-#include "device_adapter.cuh"  // for NoInfInData
+#include "device_adapter.cuh"              // for NoInfInData
+#include "ellpack_page.cuh"
 #include "ellpack_page.h"
 #include "gradient_index.h"
 #include "xgboost/data.h"
@@ -33,11 +33,6 @@ size_t EllpackPage::Size() const { return impl_->Size(); }
 
 void EllpackPage::SetBaseRowId(std::size_t row_id) { impl_->SetBaseRowId(row_id); }
 
-[[nodiscard]] common::HistogramCuts& EllpackPage::Cuts() {
-  CHECK(impl_);
-  return impl_->Cuts();
-}
-
 [[nodiscard]] common::HistogramCuts const& EllpackPage::Cuts() const {
   CHECK(impl_);
   return impl_->Cuts();
@@ -94,7 +89,8 @@ __global__ void CompressBinEllpackKernel(
 }
 
 // Construct an ELLPACK matrix with the given number of empty rows.
-EllpackPageImpl::EllpackPageImpl(DeviceOrd device, common::HistogramCuts cuts, bool is_dense,
+EllpackPageImpl::EllpackPageImpl(DeviceOrd device,
+                                 std::shared_ptr<common::HistogramCuts const> cuts, bool is_dense,
                                  size_t row_stride, size_t n_rows)
     : is_dense(is_dense), cuts_(std::move(cuts)), row_stride(row_stride), n_rows(n_rows) {
   monitor_.Init("ellpack_page");
@@ -105,12 +101,11 @@ EllpackPageImpl::EllpackPageImpl(DeviceOrd device, common::HistogramCuts cuts, b
   monitor_.Stop("InitCompressedData");
 }
 
-EllpackPageImpl::EllpackPageImpl(DeviceOrd device, common::HistogramCuts cuts,
-                                 const SparsePage &page, bool is_dense,
-                                 size_t row_stride,
+EllpackPageImpl::EllpackPageImpl(DeviceOrd device,
+                                 std::shared_ptr<common::HistogramCuts const> cuts,
+                                 const SparsePage& page, bool is_dense, size_t row_stride,
                                  common::Span<FeatureType const> feature_types)
-    : cuts_(std::move(cuts)), is_dense(is_dense), n_rows(page.Size()),
-      row_stride(row_stride) {
+    : cuts_(std::move(cuts)), is_dense(is_dense), n_rows(page.Size()), row_stride(row_stride) {
   this->InitCompressedData(device);
   this->CreateHistIndices(device, page, feature_types);
 }
@@ -127,9 +122,10 @@ EllpackPageImpl::EllpackPageImpl(Context const* ctx, DMatrix* dmat, const BatchP
   // Create the quantile sketches for the dmatrix and initialize HistogramCuts.
   row_stride = GetRowStride(dmat);
   if (!param.hess.empty()) {
-    cuts_ = common::DeviceSketchWithHessian(ctx, dmat, param.max_bin, param.hess);
+    cuts_ = std::make_shared<common::HistogramCuts>(
+        common::DeviceSketchWithHessian(ctx, dmat, param.max_bin, param.hess));
   } else {
-    cuts_ = common::DeviceSketch(ctx, dmat, param.max_bin);
+    cuts_ = std::make_shared<common::HistogramCuts>(common::DeviceSketch(ctx, dmat, param.max_bin));
   }
   monitor_.Stop("Quantiles");
 
@@ -297,7 +293,7 @@ template <typename AdapterBatch>
 EllpackPageImpl::EllpackPageImpl(AdapterBatch batch, float missing, DeviceOrd device, bool is_dense,
                                  common::Span<size_t> row_counts_span,
                                  common::Span<FeatureType const> feature_types, size_t row_stride,
-                                 size_t n_rows, common::HistogramCuts const& cuts) {
+                                 size_t n_rows, std::shared_ptr<common::HistogramCuts const> cuts) {
   dh::safe_cuda(cudaSetDevice(device.ordinal));
 
   *this = EllpackPageImpl(device, cuts, is_dense, row_stride, n_rows);
@@ -309,7 +305,7 @@ EllpackPageImpl::EllpackPageImpl(AdapterBatch batch, float missing, DeviceOrd de
   template EllpackPageImpl::EllpackPageImpl(                                               \
       __BATCH_T batch, float missing, DeviceOrd device, bool is_dense,                     \
       common::Span<size_t> row_counts_span, common::Span<FeatureType const> feature_types, \
-      size_t row_stride, size_t n_rows, common::HistogramCuts const& cuts);
+      size_t row_stride, size_t n_rows, std::shared_ptr<common::HistogramCuts const> cuts);
 
 ELLPACK_BATCH_SPECIALIZE(data::CudfAdapterBatch)
 ELLPACK_BATCH_SPECIALIZE(data::CupyAdapterBatch)
@@ -359,7 +355,11 @@ void CopyGHistToEllpack(GHistIndexMatrix const& page, common::Span<size_t const>
 
 EllpackPageImpl::EllpackPageImpl(Context const* ctx, GHistIndexMatrix const& page,
                                  common::Span<FeatureType const> ft)
-    : is_dense{page.IsDense()}, base_rowid{page.base_rowid}, n_rows{page.Size()}, cuts_{page.cut} {
+    : is_dense{page.IsDense()},
+      base_rowid{page.base_rowid},
+      n_rows{page.Size()},
+      // This makes a copy of the cut values.
+      cuts_{std::make_shared<common::HistogramCuts>(page.cut)} {
   auto it = common::MakeIndexTransformIter(
       [&](size_t i) { return page.row_ptr[i + 1] - page.row_ptr[i]; });
   row_stride = *std::max_element(it, it + page.Size());

src/data/ellpack_page.cuh

@@ -23,20 +23,20 @@ struct EllpackDeviceAccessor {
   bool is_dense;
   /*! \brief Row length for ELLPACK, equal to number of features. */
   size_t row_stride;
-  size_t base_rowid{};
-  size_t n_rows{};
-  common::CompressedIterator<uint32_t> gidx_iter;
+  bst_idx_t base_rowid{0};
+  bst_idx_t n_rows{0};
+  common::CompressedIterator<std::uint32_t> gidx_iter;
   /*! \brief Minimum value for each feature. Size equals to number of features. */
-  common::Span<const bst_float> min_fvalue;
+  common::Span<const float> min_fvalue;
   /*! \brief Histogram cut pointers. Size equals to (number of features + 1). */
-  common::Span<const uint32_t> feature_segments;
+  common::Span<const std::uint32_t> feature_segments;
   /*! \brief Histogram cut values. Size equals to (bins per feature * number of features). */
-  common::Span<const bst_float> gidx_fvalue_map;
+  common::Span<const float> gidx_fvalue_map;
 
   common::Span<const FeatureType> feature_types;
 
-  EllpackDeviceAccessor(DeviceOrd device, const common::HistogramCuts& cuts, bool is_dense,
-                        size_t row_stride, size_t base_rowid, size_t n_rows,
+  EllpackDeviceAccessor(DeviceOrd device, std::shared_ptr<const common::HistogramCuts> cuts,
+                        bool is_dense, size_t row_stride, size_t base_rowid, size_t n_rows,
                         common::CompressedIterator<uint32_t> gidx_iter,
                         common::Span<FeatureType const> feature_types)
       : is_dense(is_dense),
@@ -46,16 +46,16 @@ struct EllpackDeviceAccessor {
         gidx_iter(gidx_iter),
         feature_types{feature_types} {
     if (device.IsCPU()) {
-      gidx_fvalue_map = cuts.cut_values_.ConstHostSpan();
-      feature_segments = cuts.cut_ptrs_.ConstHostSpan();
-      min_fvalue = cuts.min_vals_.ConstHostSpan();
+      gidx_fvalue_map = cuts->cut_values_.ConstHostSpan();
+      feature_segments = cuts->cut_ptrs_.ConstHostSpan();
+      min_fvalue = cuts->min_vals_.ConstHostSpan();
     } else {
-      cuts.cut_values_.SetDevice(device);
-      cuts.cut_ptrs_.SetDevice(device);
-      cuts.min_vals_.SetDevice(device);
-      gidx_fvalue_map = cuts.cut_values_.ConstDeviceSpan();
-      feature_segments = cuts.cut_ptrs_.ConstDeviceSpan();
-      min_fvalue = cuts.min_vals_.ConstDeviceSpan();
+      cuts->cut_values_.SetDevice(device);
+      cuts->cut_ptrs_.SetDevice(device);
+      cuts->min_vals_.SetDevice(device);
+      gidx_fvalue_map = cuts->cut_values_.ConstDeviceSpan();
+      feature_segments = cuts->cut_ptrs_.ConstDeviceSpan();
+      min_fvalue = cuts->min_vals_.ConstDeviceSpan();
     }
   }
   // Get a matrix element, uses binary search for look up Return NaN if missing
@@ -142,13 +142,14 @@ class EllpackPageImpl {
    * This is used in the sampling case. The ELLPACK page is constructed from an existing EllpackInfo
    * and the given number of rows.
    */
-  EllpackPageImpl(DeviceOrd device, common::HistogramCuts cuts, bool is_dense, size_t row_stride,
-                  size_t n_rows);
+  EllpackPageImpl(DeviceOrd device, std::shared_ptr<common::HistogramCuts const> cuts,
+                  bool is_dense, size_t row_stride, size_t n_rows);
   /*!
    * \brief Constructor used for external memory.
    */
-  EllpackPageImpl(DeviceOrd device, common::HistogramCuts cuts, const SparsePage& page,
-                  bool is_dense, size_t row_stride, common::Span<FeatureType const> feature_types);
+  EllpackPageImpl(DeviceOrd device, std::shared_ptr<common::HistogramCuts const> cuts,
+                  const SparsePage& page, bool is_dense, size_t row_stride,
+                  common::Span<FeatureType const> feature_types);
 
   /*!
    * \brief Constructor from an existing DMatrix.
@@ -162,7 +163,7 @@ class EllpackPageImpl {
   explicit EllpackPageImpl(AdapterBatch batch, float missing, DeviceOrd device, bool is_dense,
                            common::Span<size_t> row_counts_span,
                            common::Span<FeatureType const> feature_types, size_t row_stride,
-                           size_t n_rows, common::HistogramCuts const& cuts);
+                           size_t n_rows, std::shared_ptr<common::HistogramCuts const> cuts);
   /**
    * \brief Constructor from an existing CPU gradient index.
    */
@@ -194,16 +195,17 @@ class EllpackPageImpl {
     base_rowid = row_id;
   }
 
-  [[nodiscard]] common::HistogramCuts& Cuts() { return cuts_; }
-  [[nodiscard]] common::HistogramCuts const& Cuts() const { return cuts_; }
+  [[nodiscard]] common::HistogramCuts const& Cuts() const { return *cuts_; }
+  [[nodiscard]] std::shared_ptr<common::HistogramCuts const> CutsShared() const { return cuts_; }
+  void SetCuts(std::shared_ptr<common::HistogramCuts const> cuts) { cuts_ = cuts; }
 
   /*! \return Estimation of memory cost of this page. */
   static size_t MemCostBytes(size_t num_rows, size_t row_stride, const common::HistogramCuts&cuts) ;
 
 
   /*! \brief Return the total number of symbols (total number of bins plus 1 for
    * not found). */
-  [[nodiscard]] std::size_t NumSymbols() const { return cuts_.TotalBins() + 1; }
+  [[nodiscard]] std::size_t NumSymbols() const { return cuts_->TotalBins() + 1; }
 
   [[nodiscard]] EllpackDeviceAccessor GetDeviceAccessor(
       DeviceOrd device, common::Span<FeatureType const> feature_types = {}) const;
@@ -225,19 +227,18 @@ class EllpackPageImpl {
    */
   void InitCompressedData(DeviceOrd device);
 
-
-public:
+ public:
   /*! \brief Whether or not if the matrix is dense. */
   bool is_dense;
   /*! \brief Row length for ELLPACK. */
   size_t row_stride;
-  size_t base_rowid{0};
-  size_t n_rows{};
+  bst_idx_t base_rowid{0};
+  bst_idx_t n_rows{};
   /*! \brief global index of histogram, which is stored in ELLPACK format. */
   HostDeviceVector<common::CompressedByteT> gidx_buffer;
 
  private:
-  common::HistogramCuts cuts_;
+  std::shared_ptr<common::HistogramCuts const> cuts_;
   common::Monitor monitor_;
 };
 

src/data/ellpack_page.h

@@ -49,7 +49,6 @@ class EllpackPage {
   [[nodiscard]] const EllpackPageImpl* Impl() const { return impl_.get(); }
   EllpackPageImpl* Impl() { return impl_.get(); }
 
-  [[nodiscard]] common::HistogramCuts& Cuts();
   [[nodiscard]] common::HistogramCuts const& Cuts() const;
 
  private: