Move GHistIndex into DMatrix. (#7064)

src/data/gradient_index.cc

@@ -0,0 +1,165 @@
+/*!
+ * Copyright 2017-2021 by Contributors
+ * \brief Data type for fast histogram aggregation.
+ */
+#include <algorithm>
+#include <limits>
+#include "gradient_index.h"
+#include "../common/hist_util.h"
+
+namespace xgboost {
+void GHistIndexMatrix::Init(DMatrix* p_fmat, int max_bins) {
+  cut = common::SketchOnDMatrix(p_fmat, max_bins);
+
+  max_num_bins = max_bins;
+  const int32_t nthread = omp_get_max_threads();
+  const uint32_t nbins = cut.Ptrs().back();
+  hit_count.resize(nbins, 0);
+  hit_count_tloc_.resize(nthread * nbins, 0);
+
+  this->p_fmat = p_fmat;
+  size_t new_size = 1;
+  for (const auto &batch : p_fmat->GetBatches<SparsePage>()) {
+    new_size += batch.Size();
+  }
+
+  row_ptr.resize(new_size);
+  row_ptr[0] = 0;
+
+  size_t rbegin = 0;
+  size_t prev_sum = 0;
+  const bool isDense = p_fmat->IsDense();
+  this->isDense_ = isDense;
+
+  for (const auto &batch : p_fmat->GetBatches<SparsePage>()) {
+    // The number of threads is pegged to the batch size. If the OMP
+    // block is parallelized on anything other than the batch/block size,
+    // it should be reassigned
+    const size_t batch_threads = std::max(
+        size_t(1),
+        std::min(batch.Size(), static_cast<size_t>(omp_get_max_threads())));
+    auto page = batch.GetView();
+    common::MemStackAllocator<size_t, 128> partial_sums(batch_threads);
+    size_t* p_part = partial_sums.Get();
+
+    size_t block_size =  batch.Size() / batch_threads;
+
+    dmlc::OMPException exc;
+    #pragma omp parallel num_threads(batch_threads)
+    {
+      #pragma omp for
+      for (omp_ulong tid = 0; tid < batch_threads; ++tid) {
+        exc.Run([&]() {
+          size_t ibegin = block_size * tid;
+          size_t iend = (tid == (batch_threads-1) ? batch.Size() : (block_size * (tid+1)));
+
+          size_t sum = 0;
+          for (size_t i = ibegin; i < iend; ++i) {
+            sum += page[i].size();
+            row_ptr[rbegin + 1 + i] = sum;
+          }
+        });
+      }
+
+      #pragma omp single
+      {
+        exc.Run([&]() {
+          p_part[0] = prev_sum;
+          for (size_t i = 1; i < batch_threads; ++i) {
+            p_part[i] = p_part[i - 1] + row_ptr[rbegin + i*block_size];
+          }
+        });
+      }
+
+      #pragma omp for
+      for (omp_ulong tid = 0; tid < batch_threads; ++tid) {
+        exc.Run([&]() {
+          size_t ibegin = block_size * tid;
+          size_t iend = (tid == (batch_threads-1) ? batch.Size() : (block_size * (tid+1)));
+
+          for (size_t i = ibegin; i < iend; ++i) {
+            row_ptr[rbegin + 1 + i] += p_part[tid];
+          }
+        });
+      }
+    }
+    exc.Rethrow();
+
+    const size_t n_offsets = cut.Ptrs().size() - 1;
+    const size_t n_index = row_ptr[rbegin + batch.Size()];
+    ResizeIndex(n_index, isDense);
+
+    CHECK_GT(cut.Values().size(), 0U);
+
+    uint32_t* offsets = nullptr;
+    if (isDense) {
+      index.ResizeOffset(n_offsets);
+      offsets = index.Offset();
+      for (size_t i = 0; i < n_offsets; ++i) {
+        offsets[i] = cut.Ptrs()[i];
+      }
+    }
+
+    if (isDense) {
+      common::BinTypeSize curent_bin_size = index.GetBinTypeSize();
+      if (curent_bin_size == common::kUint8BinsTypeSize) {
+        common::Span<uint8_t> index_data_span = {index.data<uint8_t>(),
+                                                 n_index};
+        SetIndexData(index_data_span, batch_threads, batch, rbegin, nbins,
+                     [offsets](auto idx, auto j) {
+                       return static_cast<uint8_t>(idx - offsets[j]);
+                     });
+
+      } else if (curent_bin_size == common::kUint16BinsTypeSize) {
+        common::Span<uint16_t> index_data_span = {index.data<uint16_t>(),
+                                                  n_index};
+        SetIndexData(index_data_span, batch_threads, batch, rbegin, nbins,
+                     [offsets](auto idx, auto j) {
+                       return static_cast<uint16_t>(idx - offsets[j]);
+                     });
+      } else {
+        CHECK_EQ(curent_bin_size, common::kUint32BinsTypeSize);
+        common::Span<uint32_t> index_data_span = {index.data<uint32_t>(),
+                                                  n_index};
+        SetIndexData(index_data_span, batch_threads, batch, rbegin, nbins,
+                     [offsets](auto idx, auto j) {
+                       return static_cast<uint32_t>(idx - offsets[j]);
+                     });
+      }
+
+    /* For sparse DMatrix we have to store index of feature for each bin
+       in index field to chose right offset. So offset is nullptr and index is not reduced */
+    } else {
+      common::Span<uint32_t> index_data_span = {index.data<uint32_t>(), n_index};
+      SetIndexData(index_data_span, batch_threads, batch, rbegin, nbins,
+                   [](auto idx, auto) { return idx; });
+    }
+
+    common::ParallelFor(bst_omp_uint(nbins), nthread, [&](bst_omp_uint idx) {
+      for (int32_t tid = 0; tid < nthread; ++tid) {
+        hit_count[idx] += hit_count_tloc_[tid * nbins + idx];
+        hit_count_tloc_[tid * nbins + idx] = 0;  // reset for next batch
+      }
+    });
+
+    prev_sum = row_ptr[rbegin + batch.Size()];
+    rbegin += batch.Size();
+  }
+}
+
+
+void GHistIndexMatrix::ResizeIndex(const size_t n_index,
+                                   const bool isDense) {
+  if ((max_num_bins - 1 <= static_cast<int>(std::numeric_limits<uint8_t>::max())) && isDense) {
+    index.SetBinTypeSize(common::kUint8BinsTypeSize);
+    index.Resize((sizeof(uint8_t)) * n_index);
+  } else if ((max_num_bins - 1 > static_cast<int>(std::numeric_limits<uint8_t>::max())  &&
+    max_num_bins - 1 <= static_cast<int>(std::numeric_limits<uint16_t>::max())) && isDense) {
+    index.SetBinTypeSize(common::kUint16BinsTypeSize);
+    index.Resize((sizeof(uint16_t)) * n_index);
+  } else {
+    index.SetBinTypeSize(common::kUint32BinsTypeSize);
+    index.Resize((sizeof(uint32_t)) * n_index);
+  }
+}
+}  // namespace xgboost

src/data/gradient_index.h

@@ -0,0 +1,86 @@
+/*!
+ * Copyright 2017-2021 by Contributors
+ * \brief Data type for fast histogram aggregation.
+ */
+#ifndef XGBOOST_DATA_GRADIENT_INDEX_H_
+#define XGBOOST_DATA_GRADIENT_INDEX_H_
+#include <vector>
+#include "xgboost/base.h"
+#include "xgboost/data.h"
+#include "../common/hist_util.h"
+#include "../common/threading_utils.h"
+
+namespace xgboost {
+/*!
+ * \brief preprocessed global index matrix, in CSR format
+ *
+ *  Transform floating values to integer index in histogram This is a global histogram
+ *  index for CPU histogram.  On GPU ellpack page is used.
+ */
+class GHistIndexMatrix {
+ public:
+  /*! \brief row pointer to rows by element position */
+  std::vector<size_t> row_ptr;
+  /*! \brief The index data */
+  common::Index index;
+  /*! \brief hit count of each index */
+  std::vector<size_t> hit_count;
+  /*! \brief The corresponding cuts */
+  common::HistogramCuts cut;
+  DMatrix* p_fmat;
+  size_t max_num_bins;
+
+  GHistIndexMatrix(DMatrix* x, int32_t max_bin) {
+    this->Init(x, max_bin);
+  }
+  // Create a global histogram matrix, given cut
+  void Init(DMatrix* p_fmat, int max_num_bins);
+
+  // specific method for sparse data as no possibility to reduce allocated memory
+  template <typename BinIdxType, typename GetOffset>
+  void SetIndexData(common::Span<BinIdxType> index_data_span,
+                    size_t batch_threads, const SparsePage &batch,
+                    size_t rbegin, size_t nbins, GetOffset get_offset) {
+    const xgboost::Entry *data_ptr = batch.data.HostVector().data();
+    const std::vector<bst_row_t> &offset_vec = batch.offset.HostVector();
+    const size_t batch_size = batch.Size();
+    CHECK_LT(batch_size, offset_vec.size());
+    BinIdxType* index_data = index_data_span.data();
+    common::ParallelFor(omp_ulong(batch_size), batch_threads, [&](omp_ulong i) {
+      const int tid = omp_get_thread_num();
+      size_t ibegin = row_ptr[rbegin + i];
+      size_t iend = row_ptr[rbegin + i + 1];
+      const size_t size = offset_vec[i + 1] - offset_vec[i];
+      SparsePage::Inst inst = {data_ptr + offset_vec[i], size};
+      CHECK_EQ(ibegin + inst.size(), iend);
+      for (bst_uint j = 0; j < inst.size(); ++j) {
+        uint32_t idx = cut.SearchBin(inst[j]);
+        index_data[ibegin + j] = get_offset(idx, j);
+        ++hit_count_tloc_[tid * nbins + idx];
+      }
+    });
+  }
+
+  void ResizeIndex(const size_t n_index,
+                   const bool isDense);
+
+  inline void GetFeatureCounts(size_t* counts) const {
+    auto nfeature = cut.Ptrs().size() - 1;
+    for (unsigned fid = 0; fid < nfeature; ++fid) {
+      auto ibegin = cut.Ptrs()[fid];
+      auto iend = cut.Ptrs()[fid + 1];
+      for (auto i = ibegin; i < iend; ++i) {
+        counts[fid] += hit_count[i];
+      }
+    }
+  }
+  inline bool IsDense() const {
+    return isDense_;
+  }
+
+ private:
+  std::vector<size_t> hit_count_tloc_;
+  bool isDense_;
+};
+}      // namespace xgboost
+#endif  // XGBOOST_DATA_GRADIENT_INDEX_H_

src/data/iterative_device_dmatrix.h

@@ -58,6 +58,10 @@ class IterativeDeviceDMatrix : public DMatrix {
     LOG(FATAL) << "Not implemented.";
     return BatchSet<SortedCSCPage>(BatchIterator<SortedCSCPage>(nullptr));
   }
+  BatchSet<GHistIndexMatrix> GetGradientIndex(const BatchParam&) override {
+    LOG(FATAL) << "Not implemented.";
+    return BatchSet<GHistIndexMatrix>(BatchIterator<GHistIndexMatrix>(nullptr));
+  }
 
   BatchSet<EllpackPage> GetEllpackBatches(const BatchParam& param) override;
 

src/data/proxy_dmatrix.h

@@ -97,6 +97,10 @@ class DMatrixProxy : public DMatrix {
     LOG(FATAL) << "Not implemented.";
     return BatchSet<EllpackPage>(BatchIterator<EllpackPage>(nullptr));
   }
+  BatchSet<GHistIndexMatrix> GetGradientIndex(const BatchParam&) override {
+    LOG(FATAL) << "Not implemented.";
+    return BatchSet<GHistIndexMatrix>(BatchIterator<GHistIndexMatrix>(nullptr));
+  }
 
   dmlc::any Adapter() const {
     return batch_;

src/data/simple_dmatrix.cc

@@ -17,6 +17,7 @@
 #include "../common/random.h"
 #include "../common/threading_utils.h"
 #include "adapter.h"
+#include "gradient_index.h"
 
 namespace xgboost {
 namespace data {
@@ -89,6 +90,20 @@ BatchSet<EllpackPage> SimpleDMatrix::GetEllpackBatches(const BatchParam& param)
   return BatchSet<EllpackPage>(begin_iter);
 }
 
+BatchSet<GHistIndexMatrix> SimpleDMatrix::GetGradientIndex(const BatchParam& param) {
+  if (!(batch_param_ != BatchParam{})) {
+    CHECK(param != BatchParam{}) << "Batch parameter is not initialized.";
+  }
+  if (!gradient_index_  || (batch_param_ != param && param != BatchParam{})) {
+    CHECK_GE(param.max_bin, 2);
+    gradient_index_.reset(new GHistIndexMatrix(this, param.max_bin));
+    batch_param_ = param;
+  }
+  auto begin_iter = BatchIterator<GHistIndexMatrix>(
+      new SimpleBatchIteratorImpl<GHistIndexMatrix>(gradient_index_.get()));
+  return BatchSet<GHistIndexMatrix>(begin_iter);
+}
+
 template <typename AdapterT>
 SimpleDMatrix::SimpleDMatrix(AdapterT* adapter, float missing, int nthread) {
   std::vector<uint64_t> qids;

src/data/simple_dmatrix.h

@@ -13,6 +13,7 @@
 #include <memory>
 #include <string>
 
+#include "gradient_index.h"
 
 namespace xgboost {
 namespace data {
@@ -43,12 +44,14 @@ class SimpleDMatrix : public DMatrix {
   BatchSet<CSCPage> GetColumnBatches() override;
   BatchSet<SortedCSCPage> GetSortedColumnBatches() override;
   BatchSet<EllpackPage> GetEllpackBatches(const BatchParam& param) override;
+  BatchSet<GHistIndexMatrix> GetGradientIndex(const BatchParam& param) override;
 
   MetaInfo info_;
   SparsePage sparse_page_;  // Primary storage type
   std::unique_ptr<CSCPage> column_page_;
   std::unique_ptr<SortedCSCPage> sorted_column_page_;
   std::unique_ptr<EllpackPage> ellpack_page_;
+  std::unique_ptr<GHistIndexMatrix> gradient_index_;
   BatchParam batch_param_;
 
   bool EllpackExists() const override {

src/data/sparse_page_dmatrix.h

@@ -47,6 +47,10 @@ class SparsePageDMatrix : public DMatrix {
   BatchSet<CSCPage> GetColumnBatches() override;
   BatchSet<SortedCSCPage> GetSortedColumnBatches() override;
   BatchSet<EllpackPage> GetEllpackBatches(const BatchParam& param) override;
+  BatchSet<GHistIndexMatrix> GetGradientIndex(const BatchParam&) override {
+    LOG(FATAL) << "Not implemented.";
+    return BatchSet<GHistIndexMatrix>(BatchIterator<GHistIndexMatrix>(nullptr));
+  }
 
   // source data pointers.
   std::unique_ptr<SparsePageSource> row_source_;