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>
30 class PartitionBuilder {
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) {
148 const common::DenseColumn<BinIdxType, any_missing>& column =
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);
159 const common::SparseColumn<BinIdxType>& column
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
249 void CalculateRowOffsets() {
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
287 size_t n_offset_left;
288 size_t n_offset_right;
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