Unify the partitioner for hist and approx.

Co-authored-by: dmitry.razdoburdin <[email protected]>
Co-authored-by: jiamingy <[email protected]>

Author: 3 people

src/common/column_matrix.h

@@ -103,15 +103,18 @@ class SparseColumnIter : public Column<BinIdxT> {
 
 template <typename BinIdxT, bool any_missing>
 class DenseColumnIter : public Column<BinIdxT> {
+ public:
+  using ByteType = bool;
+
  private:
   using Base = Column<BinIdxT>;
   /* flags for missing values in dense columns */
-  std::vector<bool> const& missing_flags_;
+  std::vector<ByteType> const& missing_flags_;
   size_t feature_offset_;
 
  public:
   explicit DenseColumnIter(common::Span<const BinIdxT> index, bst_bin_t index_base,
-                           std::vector<bool> const& missing_flags, size_t feature_offset)
+                           std::vector<ByteType> const& missing_flags, size_t feature_offset)
       : Base{index, index_base}, missing_flags_{missing_flags}, feature_offset_{feature_offset} {}
   DenseColumnIter(DenseColumnIter const&) = delete;
   DenseColumnIter(DenseColumnIter&&) = default;
@@ -153,6 +156,7 @@ class ColumnMatrix {
   }
 
  public:
+  using ByteType = bool;
   // get number of features
   bst_feature_t GetNumFeature() const { return static_cast<bst_feature_t>(type_.size()); }
 
@@ -195,6 +199,8 @@ class ColumnMatrix {
     }
   }
 
+  bool IsInitialized() const { return !type_.empty(); }
+
   /**
    * \brief Push batch of data for Quantile DMatrix support.
    *
@@ -352,6 +358,13 @@ class ColumnMatrix {
 
     fi->Read(&row_ind_);
     fi->Read(&feature_offsets_);
+
+    std::vector<std::uint8_t> missing;
+    fi->Read(&missing);
+    missing_flags_.resize(missing.size());
+    std::transform(missing.cbegin(), missing.cend(), missing_flags_.begin(),
+                   [](std::uint8_t flag) { return !!flag; });
+
     index_base_ = index_base;
 #if !DMLC_LITTLE_ENDIAN
     std::underlying_type<BinTypeSize>::type v;
@@ -386,6 +399,11 @@ class ColumnMatrix {
 #endif  // !DMLC_LITTLE_ENDIAN
     write_vec(row_ind_);
     write_vec(feature_offsets_);
+    // dmlc can not handle bool vector
+    std::vector<std::uint8_t> missing(missing_flags_.size());
+    std::transform(missing_flags_.cbegin(), missing_flags_.cend(), missing.begin(),
+                   [](bool flag) { return static_cast<std::uint8_t>(flag); });
+    write_vec(missing);
 
 #if !DMLC_LITTLE_ENDIAN
     auto v = static_cast<std::underlying_type<BinTypeSize>::type>(bins_type_size_);
@@ -413,7 +431,7 @@ class ColumnMatrix {
 
   // index_base_[fid]: least bin id for feature fid
   uint32_t const* index_base_;
-  std::vector<bool> missing_flags_;
+  std::vector<ByteType> missing_flags_;
   BinTypeSize bins_type_size_;
   bool any_missing_;
 };

src/common/numeric.h

@@ -4,6 +4,8 @@
 #ifndef XGBOOST_COMMON_NUMERIC_H_
 #define XGBOOST_COMMON_NUMERIC_H_
 
+#include <dmlc/common.h>  // OMPException
+
 #include <algorithm>  // std::max
 #include <iterator>   // std::iterator_traits
 #include <vector>
@@ -106,6 +108,26 @@ inline double Reduce(Context const*, HostDeviceVector<float> const&) {
  * \brief Reduction with summation.
  */
 double Reduce(Context const* ctx, HostDeviceVector<float> const& values);
+
+template <typename It>
+void Iota(Context const* ctx, It first, It last,
+          typename std::iterator_traits<It>::value_type const& value) {
+  auto n = std::distance(first, last);
+  std::int32_t n_threads = ctx->Threads();
+  const size_t block_size = n / n_threads + !!(n % n_threads);
+  dmlc::OMPException exc;
+#pragma omp parallel num_threads(n_threads)
+  {
+    exc.Run([&]() {
+      const size_t tid = omp_get_thread_num();
+      const size_t ibegin = tid * block_size;
+      const size_t iend = std::min(ibegin + block_size, static_cast<size_t>(n));
+      for (size_t i = ibegin; i < iend; ++i) {
+        first[i] = i + value;
+      }
+    });
+  }
+}
 }  // namespace common
 }  // namespace xgboost
 

src/common/partition_builder.h

@@ -17,6 +17,7 @@
 
 #include "categorical.h"
 #include "column_matrix.h"
+#include "../tree/hist/expand_entry.h"
 #include "xgboost/generic_parameters.h"
 #include "xgboost/tree_model.h"
 
@@ -107,34 +108,42 @@ class PartitionBuilder {
   }
 
   template <typename BinIdxType, bool any_missing, bool any_cat>
-  void Partition(const size_t node_in_set, const size_t nid, const common::Range1d range,
+  void Partition(const size_t node_in_set, std::vector<xgboost::tree::CPUExpandEntry> const &nodes,
+                 const common::Range1d range,
                  const bst_bin_t split_cond, GHistIndexMatrix const& gmat,
-                 const ColumnMatrix& column_matrix, const RegTree& tree, const size_t* rid) {
+                 const common::ColumnMatrix& column_matrix,
+                 const RegTree& tree, const size_t* rid) {
     common::Span<const size_t> rid_span(rid + range.begin(), rid + range.end());
     common::Span<size_t> left = GetLeftBuffer(node_in_set, range.begin(), range.end());
     common::Span<size_t> right = GetRightBuffer(node_in_set, range.begin(), range.end());
-    const bst_uint fid = tree[nid].SplitIndex();
-    const bool default_left = tree[nid].DefaultLeft();
+    std::size_t nid = nodes[node_in_set].nid;
+    bst_feature_t fid = tree[nid].SplitIndex();
+    bool default_left = tree[nid].DefaultLeft();
     bool is_cat = tree.GetSplitTypes()[nid] == FeatureType::kCategorical;
     auto node_cats = tree.NodeCats(nid);
 
     auto const& index = gmat.index;
     auto const& cut_values = gmat.cut.Values();
     auto const& cut_ptrs = gmat.cut.Ptrs();
 
-    auto pred = [&](auto ridx, auto bin_id) {
+    auto gidx_calc = [&](auto ridx) {
+      auto begin = gmat.RowIdx(ridx);
+      if (gmat.IsDense()) {
+        return static_cast<bst_bin_t>(index[begin + fid]);
+      }
+      auto end = gmat.RowIdx(ridx + 1);
+      auto f_begin = cut_ptrs[fid];
+      auto f_end = cut_ptrs[fid + 1];
+      // bypassing the column matrix as we need the cut value instead of bin idx for categorical
+      // features.
+      return BinarySearchBin(begin, end, index, f_begin, f_end);
+    };
+
+    auto pred_hist = [&](auto ridx, auto bin_id) {
       if (any_cat && is_cat) {
-        auto begin = gmat.RowIdx(ridx);
-        auto end = gmat.RowIdx(ridx + 1);
-        auto f_begin = cut_ptrs[fid];
-        auto f_end = cut_ptrs[fid + 1];
-        // bypassing the column matrix as we need the cut value instead of bin idx for categorical
-        // features.
-        auto gidx = BinarySearchBin(begin, end, index, f_begin, f_end);
-        bool go_left;
-        if (gidx == -1) {
-          go_left = default_left;
-        } else {
+        auto gidx = gidx_calc(ridx);
+        bool go_left = default_left;
+        if (gidx > -1) {
           go_left = Decision(node_cats, cut_values[gidx], default_left);
         }
         return go_left;
@@ -143,25 +152,43 @@ class PartitionBuilder {
       }
     };
 
-    std::pair<size_t, size_t> child_nodes_sizes;
-    if (column_matrix.GetColumnType(fid) == xgboost::common::kDenseColumn) {
-      auto column = column_matrix.DenseColumn<BinIdxType, any_missing>(fid);
-      if (default_left) {
-        child_nodes_sizes = PartitionKernel<true, any_missing>(&column, rid_span, left, right,
-                                                               gmat.base_rowid, pred);
-      } else {
-        child_nodes_sizes = PartitionKernel<false, any_missing>(&column, rid_span, left, right,
-                                                                gmat.base_rowid, pred);
+    auto pred_approx = [&](auto ridx) {
+      auto gidx = gidx_calc(ridx);
+      bool go_left = default_left;
+      if (gidx > -1) {
+        if (is_cat) {
+          go_left = Decision(node_cats, cut_values[gidx], default_left);
+        } else {
+          go_left = cut_values[gidx] <= nodes[node_in_set].split.split_value;
+        }
       }
+      return go_left;
+    };
+
+    std::pair<size_t, size_t> child_nodes_sizes;
+    if (!column_matrix.IsInitialized()) {
+      child_nodes_sizes = PartitionRangeKernel(rid_span, left, right, pred_approx);
     } else {
-      CHECK_EQ(any_missing, true);
-      auto column = column_matrix.SparseColumn<BinIdxType>(fid, rid_span.front() - gmat.base_rowid);
-      if (default_left) {
-        child_nodes_sizes = PartitionKernel<true, any_missing>(&column, rid_span, left, right,
-                                                               gmat.base_rowid, pred);
+      if (column_matrix.GetColumnType(fid) == xgboost::common::kDenseColumn) {
+        auto column = column_matrix.DenseColumn<BinIdxType, any_missing>(fid);
+        if (default_left) {
+          child_nodes_sizes = PartitionKernel<true, any_missing>(&column, rid_span, left, right,
+                                                                 gmat.base_rowid, pred_hist);
+        } else {
+          child_nodes_sizes = PartitionKernel<false, any_missing>(&column, rid_span, left, right,
+                                                                  gmat.base_rowid, pred_hist);
+        }
       } else {
-        child_nodes_sizes = PartitionKernel<false, any_missing>(&column, rid_span, left, right,
-                                                                gmat.base_rowid, pred);
+        CHECK_EQ(any_missing, true);
+        auto column =
+            column_matrix.SparseColumn<BinIdxType>(fid, rid_span.front() - gmat.base_rowid);
+        if (default_left) {
+          child_nodes_sizes = PartitionKernel<true, any_missing>(&column, rid_span, left, right,
+                                                                 gmat.base_rowid, pred_hist);
+        } else {
+          child_nodes_sizes = PartitionKernel<false, any_missing>(&column, rid_span, left, right,
+                                                                  gmat.base_rowid, pred_hist);
+        }
       }
     }
 
@@ -172,37 +199,6 @@ class PartitionBuilder {
     SetNRightElems(node_in_set, range.begin(), n_right);
   }
 
-  /**
-   * \brief Partition tree nodes with specific range of row indices.
-   *
-   * \tparam Pred       Predicate for whether a row should be partitioned to the left node.
-   *
-   * \param node_in_set The index of node in current batch of nodes.
-   * \param nid         The canonical node index (node index in the tree).
-   * \param range       The range of input row index.
-   * \param fidx        Feature index.
-   * \param p_row_set_collection Pointer to rows that are  being partitioned.
-   * \param pred        A callback function that returns whether current row should be
-   *                    partitioned to the left node, it should accept the row index as
-   *                    input and returns a boolean value.
-   */
-  template <typename Pred>
-  void PartitionRange(const size_t node_in_set, const size_t nid, common::Range1d range,
-                      common::RowSetCollection* p_row_set_collection, Pred pred) {
-    auto& row_set_collection = *p_row_set_collection;
-    const size_t* p_ridx = row_set_collection[nid].begin;
-    common::Span<const size_t> ridx(p_ridx + range.begin(), p_ridx + range.end());
-    common::Span<size_t> left = this->GetLeftBuffer(node_in_set, range.begin(), range.end());
-    common::Span<size_t> right = this->GetRightBuffer(node_in_set, range.begin(), range.end());
-    std::pair<size_t, size_t> child_nodes_sizes = PartitionRangeKernel(ridx, left, right, pred);
-
-    const size_t n_left = child_nodes_sizes.first;
-    const size_t n_right = child_nodes_sizes.second;
-
-    this->SetNLeftElems(node_in_set, range.begin(), n_left);
-    this->SetNRightElems(node_in_set, range.begin(), n_right);
-  }
-
   // allocate thread local memory, should be called for each specific task
   void AllocateForTask(size_t id) {
     if (mem_blocks_[id].get() == nullptr) {