xgboost
partition_builder.h
Go to the documentation of this file.
1 
7 #ifndef XGBOOST_COMMON_PARTITION_BUILDER_H_
8 #define XGBOOST_COMMON_PARTITION_BUILDER_H_
9 
10 #include <xgboost/data.h>
11 
12 #include <algorithm>
13 #include <memory>
14 #include <utility>
15 #include <vector>
16 
17 #include "categorical.h"
18 #include "column_matrix.h"
19 #include "xgboost/tree_model.h"
20 
21 namespace xgboost {
22 namespace common {
23 
24 // The builder is required for samples partition to left and rights children for set of nodes
25 // Responsible for:
26 // 1) Effective memory allocation for intermediate results for multi-thread work
27 // 2) Merging partial results produced by threads into original row set (row_set_collection_)
28 // BlockSize is template to enable memory alignment easily with C++11 'alignas()' feature
29 template<size_t BlockSize>
31  public:
32  template<typename Func>
33  void Init(const size_t n_tasks, size_t n_nodes, Func funcNTask) {
34  left_right_nodes_sizes_.resize(n_nodes);
35  blocks_offsets_.resize(n_nodes+1);
36 
37  blocks_offsets_[0] = 0;
38  for (size_t i = 1; i < n_nodes+1; ++i) {
39  blocks_offsets_[i] = blocks_offsets_[i-1] + funcNTask(i-1);
40  }
41 
42  if (n_tasks > max_n_tasks_) {
43  mem_blocks_.resize(n_tasks);
44  max_n_tasks_ = n_tasks;
45  }
46  }
47 
48  // split row indexes (rid_span) to 2 parts (left_part, right_part) depending
49  // on comparison of indexes values (idx_span) and split point (split_cond)
50  // Handle dense columns
51  // Analog of std::stable_partition, but in no-inplace manner
52  template <bool default_left, bool any_missing, typename ColumnType, typename Predicate>
53  inline std::pair<size_t, size_t> PartitionKernel(const ColumnType& column,
54  common::Span<const size_t> row_indices,
55  common::Span<size_t> left_part,
56  common::Span<size_t> right_part,
57  size_t base_rowid, Predicate&& pred) {
58  size_t* p_left_part = left_part.data();
59  size_t* p_right_part = right_part.data();
60  size_t nleft_elems = 0;
61  size_t nright_elems = 0;
62  auto state = column.GetInitialState(row_indices.front() - base_rowid);
63 
64  auto p_row_indices = row_indices.data();
65  auto n_samples = row_indices.size();
66 
67  for (size_t i = 0; i < n_samples; ++i) {
68  auto rid = p_row_indices[i];
69  const int32_t bin_id = column.GetBinIdx(rid - base_rowid, &state);
70  if (any_missing && bin_id == ColumnType::kMissingId) {
71  if (default_left) {
72  p_left_part[nleft_elems++] = rid;
73  } else {
74  p_right_part[nright_elems++] = rid;
75  }
76  } else {
77  if (pred(rid, bin_id)) {
78  p_left_part[nleft_elems++] = rid;
79  } else {
80  p_right_part[nright_elems++] = rid;
81  }
82  }
83  }
84 
85  return {nleft_elems, nright_elems};
86  }
87 
88  template <typename Pred>
89  inline std::pair<size_t, size_t> PartitionRangeKernel(common::Span<const size_t> ridx,
90  common::Span<size_t> left_part,
91  common::Span<size_t> right_part,
92  Pred pred) {
93  size_t* p_left_part = left_part.data();
94  size_t* p_right_part = right_part.data();
95  size_t nleft_elems = 0;
96  size_t nright_elems = 0;
97  for (auto row_id : ridx) {
98  if (pred(row_id)) {
99  p_left_part[nleft_elems++] = row_id;
100  } else {
101  p_right_part[nright_elems++] = row_id;
102  }
103  }
104  return {nleft_elems, nright_elems};
105  }
106 
107  template <typename BinIdxType, bool any_missing, bool any_cat>
108  void Partition(const size_t node_in_set, const size_t nid, const common::Range1d range,
109  const int32_t split_cond, GHistIndexMatrix const& gmat,
110  const ColumnMatrix& column_matrix, const RegTree& tree, const size_t* rid) {
111  common::Span<const size_t> rid_span(rid + range.begin(), rid + range.end());
112  common::Span<size_t> left = GetLeftBuffer(node_in_set, range.begin(), range.end());
113  common::Span<size_t> right = GetRightBuffer(node_in_set, range.begin(), range.end());
114  const bst_uint fid = tree[nid].SplitIndex();
115  const bool default_left = tree[nid].DefaultLeft();
116  const auto column_ptr = column_matrix.GetColumn<BinIdxType, any_missing>(fid);
117 
118  bool is_cat = tree.GetSplitTypes()[nid] == FeatureType::kCategorical;
119  auto node_cats = tree.NodeCats(nid);
120 
121  auto const& index = gmat.index;
122  auto const& cut_values = gmat.cut.Values();
123  auto const& cut_ptrs = gmat.cut.Ptrs();
124 
125  auto pred = [&](auto ridx, auto bin_id) {
126  if (any_cat && is_cat) {
127  auto begin = gmat.RowIdx(ridx);
128  auto end = gmat.RowIdx(ridx + 1);
129  auto f_begin = cut_ptrs[fid];
130  auto f_end = cut_ptrs[fid + 1];
131  // bypassing the column matrix as we need the cut value instead of bin idx for categorical
132  // features.
133  auto gidx = BinarySearchBin(begin, end, index, f_begin, f_end);
134  bool go_left;
135  if (gidx == -1) {
136  go_left = default_left;
137  } else {
138  go_left = Decision(node_cats, cut_values[gidx], default_left);
139  }
140  return go_left;
141  } else {
142  return bin_id <= split_cond;
143  }
144  };
145 
146  std::pair<size_t, size_t> child_nodes_sizes;
147  if (column_ptr->GetType() == xgboost::common::kDenseColumn) {
149  static_cast<const common::DenseColumn<BinIdxType, any_missing>& >(*(column_ptr.get()));
150  if (default_left) {
151  child_nodes_sizes = PartitionKernel<true, any_missing>(column, rid_span, left, right,
152  gmat.base_rowid, pred);
153  } else {
154  child_nodes_sizes = PartitionKernel<false, any_missing>(column, rid_span, left, right,
155  gmat.base_rowid, pred);
156  }
157  } else {
158  CHECK_EQ(any_missing, true);
160  = static_cast<const common::SparseColumn<BinIdxType>& >(*(column_ptr.get()));
161  if (default_left) {
162  child_nodes_sizes = PartitionKernel<true, any_missing>(column, rid_span, left, right,
163  gmat.base_rowid, pred);
164  } else {
165  child_nodes_sizes = PartitionKernel<false, any_missing>(column, rid_span, left, right,
166  gmat.base_rowid, pred);
167  }
168  }
169 
170  const size_t n_left = child_nodes_sizes.first;
171  const size_t n_right = child_nodes_sizes.second;
172 
173  SetNLeftElems(node_in_set, range.begin(), range.end(), n_left);
174  SetNRightElems(node_in_set, range.begin(), range.end(), n_right);
175  }
176 
191  template <typename Pred>
192  void PartitionRange(const size_t node_in_set, const size_t nid, common::Range1d range,
193  bst_feature_t fidx, common::RowSetCollection* p_row_set_collection,
194  Pred pred) {
195  auto& row_set_collection = *p_row_set_collection;
196  const size_t* p_ridx = row_set_collection[nid].begin;
197  common::Span<const size_t> ridx(p_ridx + range.begin(), p_ridx + range.end());
198  common::Span<size_t> left = this->GetLeftBuffer(node_in_set, range.begin(), range.end());
199  common::Span<size_t> right = this->GetRightBuffer(node_in_set, range.begin(), range.end());
200  std::pair<size_t, size_t> child_nodes_sizes = PartitionRangeKernel(ridx, left, right, pred);
201 
202  const size_t n_left = child_nodes_sizes.first;
203  const size_t n_right = child_nodes_sizes.second;
204 
205  this->SetNLeftElems(node_in_set, range.begin(), range.end(), n_left);
206  this->SetNRightElems(node_in_set, range.begin(), range.end(), n_right);
207  }
208 
209  // allocate thread local memory, should be called for each specific task
210  void AllocateForTask(size_t id) {
211  if (mem_blocks_[id].get() == nullptr) {
212  BlockInfo* local_block_ptr = new BlockInfo;
213  CHECK_NE(local_block_ptr, (BlockInfo*)nullptr);
214  mem_blocks_[id].reset(local_block_ptr);
215  }
216  }
217 
218  common::Span<size_t> GetLeftBuffer(int nid, size_t begin, size_t end) {
219  const size_t task_idx = GetTaskIdx(nid, begin);
220  return { mem_blocks_.at(task_idx)->Left(), end - begin };
221  }
222 
223  common::Span<size_t> GetRightBuffer(int nid, size_t begin, size_t end) {
224  const size_t task_idx = GetTaskIdx(nid, begin);
225  return { mem_blocks_.at(task_idx)->Right(), end - begin };
226  }
227 
228  void SetNLeftElems(int nid, size_t begin, size_t end, size_t n_left) {
229  size_t task_idx = GetTaskIdx(nid, begin);
230  mem_blocks_.at(task_idx)->n_left = n_left;
231  }
232 
233  void SetNRightElems(int nid, size_t begin, size_t end, size_t n_right) {
234  size_t task_idx = GetTaskIdx(nid, begin);
235  mem_blocks_.at(task_idx)->n_right = n_right;
236  }
237 
238 
239  size_t GetNLeftElems(int nid) const {
240  return left_right_nodes_sizes_[nid].first;
241  }
242 
243  size_t GetNRightElems(int nid) const {
244  return left_right_nodes_sizes_[nid].second;
245  }
246 
247  // Each thread has partial results for some set of tree-nodes
248  // The function decides order of merging partial results into final row set
250  for (size_t i = 0; i < blocks_offsets_.size()-1; ++i) {
251  size_t n_left = 0;
252  for (size_t j = blocks_offsets_[i]; j < blocks_offsets_[i+1]; ++j) {
253  mem_blocks_[j]->n_offset_left = n_left;
254  n_left += mem_blocks_[j]->n_left;
255  }
256  size_t n_right = 0;
257  for (size_t j = blocks_offsets_[i]; j < blocks_offsets_[i+1]; ++j) {
258  mem_blocks_[j]->n_offset_right = n_left + n_right;
259  n_right += mem_blocks_[j]->n_right;
260  }
261  left_right_nodes_sizes_[i] = {n_left, n_right};
262  }
263  }
264 
265  void MergeToArray(int nid, size_t begin, size_t* rows_indexes) {
266  size_t task_idx = GetTaskIdx(nid, begin);
267 
268  size_t* left_result = rows_indexes + mem_blocks_[task_idx]->n_offset_left;
269  size_t* right_result = rows_indexes + mem_blocks_[task_idx]->n_offset_right;
270 
271  const size_t* left = mem_blocks_[task_idx]->Left();
272  const size_t* right = mem_blocks_[task_idx]->Right();
273 
274  std::copy_n(left, mem_blocks_[task_idx]->n_left, left_result);
275  std::copy_n(right, mem_blocks_[task_idx]->n_right, right_result);
276  }
277 
278  size_t GetTaskIdx(int nid, size_t begin) {
279  return blocks_offsets_[nid] + begin / BlockSize;
280  }
281 
282  protected:
283  struct BlockInfo{
284  size_t n_left;
285  size_t n_right;
286 
289 
290  size_t* Left() {
291  return &left_data_[0];
292  }
293 
294  size_t* Right() {
295  return &right_data_[0];
296  }
297  private:
298  size_t left_data_[BlockSize];
299  size_t right_data_[BlockSize];
300  };
301  std::vector<std::pair<size_t, size_t>> left_right_nodes_sizes_;
302  std::vector<size_t> blocks_offsets_;
303  std::vector<std::shared_ptr<BlockInfo>> mem_blocks_;
304  size_t max_n_tasks_ = 0;
305 };
306 
307 } // namespace common
308 } // namespace xgboost
309 
310 #endif // XGBOOST_COMMON_PARTITION_BUILDER_H_
xgboost::common::PartitionBuilder::PartitionKernel
std::pair< size_t, size_t > PartitionKernel(const ColumnType &column, common::Span< const size_t > row_indices, common::Span< size_t > left_part, common::Span< size_t > right_part, size_t base_rowid, Predicate &&pred)
Definition: partition_builder.h:53
xgboost::common::PartitionBuilder::MergeToArray
void MergeToArray(int nid, size_t begin, size_t *rows_indexes)
Definition: partition_builder.h:265
xgboost::common::PartitionBuilder::BlockInfo::Right
size_t * Right()
Definition: partition_builder.h:294
xgboost::common::PartitionBuilder
Definition: partition_builder.h:30
xgboost::common::PartitionBuilder::GetRightBuffer
common::Span< size_t > GetRightBuffer(int nid, size_t begin, size_t end)
Definition: partition_builder.h:223
categorical.h
xgboost::RegTree::GetSplitTypes
const std::vector< FeatureType > & GetSplitTypes() const
Get split types for all nodes.
Definition: tree_model.h:605
xgboost::common::PartitionBuilder::PartitionRange
void PartitionRange(const size_t node_in_set, const size_t nid, common::Range1d range, bst_feature_t fidx, common::RowSetCollection *p_row_set_collection, Pred pred)
Partition tree nodes with specific range of row indices.
Definition: partition_builder.h:192
xgboost::common::PartitionBuilder::max_n_tasks_
size_t max_n_tasks_
Definition: partition_builder.h:304
xgboost::common::Range1d::begin
size_t begin() const
Definition: threading_utils.h:41
column_matrix.h
Utility for fast column-wise access.
xgboost::common::PartitionBuilder::mem_blocks_
std::vector< std::shared_ptr< BlockInfo > > mem_blocks_
Definition: partition_builder.h:303
xgboost::RegTree::NodeCats
common::Span< uint32_t const > NodeCats(bst_node_t nidx) const
Get the bit storage for categories.
Definition: tree_model.h:610
xgboost::common::ColumnMatrix
a collection of columns, with support for construction from GHistIndexMatrix.
Definition: column_matrix.h:134
xgboost::common::kDenseColumn
@ kDenseColumn
Definition: column_matrix.h:26
xgboost::common::PartitionBuilder::BlockInfo::n_offset_right
size_t n_offset_right
Definition: partition_builder.h:288
xgboost::common::RowSetCollection
collection of rowset
Definition: row_set.h:20
xgboost::bst_feature_t
uint32_t bst_feature_t
Type for data column (feature) index.
Definition: base.h:123
xgboost::common::PartitionBuilder::BlockInfo::n_offset_left
size_t n_offset_left
Definition: partition_builder.h:287
xgboost::common::RowSetCollection::begin
std::vector< Elem >::const_iterator begin() const
Definition: row_set.h:47
xgboost::common::PartitionBuilder::BlockInfo::n_left
size_t n_left
Definition: partition_builder.h:284
xgboost::common::PartitionBuilder::PartitionRangeKernel
std::pair< size_t, size_t > PartitionRangeKernel(common::Span< const size_t > ridx, common::Span< size_t > left_part, common::Span< size_t > right_part, Pred pred)
Definition: partition_builder.h:89
xgboost::common::BinarySearchBin
int32_t XGBOOST_HOST_DEV_INLINE BinarySearchBin(size_t begin, size_t end, GradientIndex const &data, uint32_t const fidx_begin, uint32_t const fidx_end)
Definition: hist_util.h:299
xgboost::common::PartitionBuilder::Partition
void Partition(const size_t node_in_set, const size_t nid, const common::Range1d range, const int32_t split_cond, GHistIndexMatrix const &gmat, const ColumnMatrix &column_matrix, const RegTree &tree, const size_t *rid)
Definition: partition_builder.h:108
xgboost::common::PartitionBuilder::left_right_nodes_sizes_
std::vector< std::pair< size_t, size_t > > left_right_nodes_sizes_
Definition: partition_builder.h:301
xgboost::common::DenseColumn
Definition: column_matrix.h:107
xgboost::common::SparseColumn
Definition: column_matrix.h:67
xgboost::RegTree
define regression tree to be the most common tree model. This is the data structure used in xgboost's...
Definition: tree_model.h:131
xgboost::common::Span::front
XGBOOST_DEVICE reference front() const
Definition: span.h:531
xgboost::common::PartitionBuilder::GetLeftBuffer
common::Span< size_t > GetLeftBuffer(int nid, size_t begin, size_t end)
Definition: partition_builder.h:218
xgboost::common::PartitionBuilder::BlockInfo
Definition: partition_builder.h:283
xgboost::common::PartitionBuilder::BlockInfo::n_right
size_t n_right
Definition: partition_builder.h:285
xgboost::common::PartitionBuilder::BlockInfo::Left
size_t * Left()
Definition: partition_builder.h:290
xgboost::common::Span< const size_t >
xgboost::common::PartitionBuilder::GetTaskIdx
size_t GetTaskIdx(int nid, size_t begin)
Definition: partition_builder.h:278
xgboost::common::ColumnType
ColumnType
column type
Definition: column_matrix.h:26
data.h
The input data structure of xgboost.
xgboost::common::PartitionBuilder::SetNRightElems
void SetNRightElems(int nid, size_t begin, size_t end, size_t n_right)
Definition: partition_builder.h:233
xgboost::common::PartitionBuilder::AllocateForTask
void AllocateForTask(size_t id)
Definition: partition_builder.h:210
xgboost::common::PartitionBuilder::GetNLeftElems
size_t GetNLeftElems(int nid) const
Definition: partition_builder.h:239
xgboost::common::Range1d
Definition: threading_utils.h:35
xgboost::common::PartitionBuilder::blocks_offsets_
std::vector< size_t > blocks_offsets_
Definition: partition_builder.h:302
xgboost::common::PartitionBuilder::CalculateRowOffsets
void CalculateRowOffsets()
Definition: partition_builder.h:249
xgboost::common::Span::size
constexpr XGBOOST_DEVICE index_type size() const __span_noexcept
Definition: span.h:553
xgboost::common::ColumnMatrix::GetColumn
std::unique_ptr< const Column< BinIdxType > > GetColumn(unsigned fid) const
Definition: column_matrix.h:240
tree_model.h
model structure for tree
xgboost::common::Span::data
constexpr XGBOOST_DEVICE pointer data() const __span_noexcept
Definition: span.h:548
xgboost::common::Decision
XGBOOST_DEVICE bool Decision(common::Span< uint32_t const > cats, float cat, bool dft_left)
Definition: categorical.h:56
xgboost::bst_uint
uint32_t bst_uint
unsigned integer type used for feature index.
Definition: base.h:113
xgboost::common::Range1d::end
size_t end() const
Definition: threading_utils.h:45
xgboost::get
auto get(U &json) -> decltype(detail::GetImpl(*Cast< T >(&json.GetValue())))&
Get Json value.
Definition: json.h:583
xgboost::FeatureType::kCategorical
@ kCategorical
xgboost::common::PartitionBuilder::SetNLeftElems
void SetNLeftElems(int nid, size_t begin, size_t end, size_t n_left)
Definition: partition_builder.h:228
xgboost::common::PartitionBuilder::GetNRightElems
size_t GetNRightElems(int nid) const
Definition: partition_builder.h:243
xgboost
namespace of xgboost
Definition: base.h:110
xgboost::common::PartitionBuilder::Init
void Init(const size_t n_tasks, size_t n_nodes, Func funcNTask)
Definition: partition_builder.h:33