Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Some comments for row partitioner. (#4832)

Browse Source
This commit is contained in:
Jiaming Yuan
2019-09-06 03:01:42 -04:00
committed by GitHub
parent a5f232feb8
commit f90e7f9aa8
3 changed files with 45 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
src/tree/gpu_hist/row_partitioner.cuh
View File

@@ -30,19 +30,32 @@ __forceinline__ __device__ void AtomicIncrement(int64_t* d_count, bool increment
* partition training rows into different leaf nodes. */
class RowPartitioner {
public:
using TreePositionT = int;
using TreePositionT = int32_t;
using RowIndexT = bst_uint;
struct Segment;
private:
int device_idx;
/*! \brief Range of rows for each node. */
/*! \brief In here if you want to find the rows belong to a node nid, first you need to
* get the indices segment from ridx_segments[nid], then get the row index that
* represents position of row in input data X. `RowPartitioner::GetRows` would be a
* good starting place to get a sense what are these vector storing.
*
* node id -> segment -> indices of rows belonging to node
*/
/*! \brief Range of row index for each node, pointers into ridx below. */
std::vector<Segment> ridx_segments;
dh::caching_device_vector<RowIndexT> ridx_a;
dh::caching_device_vector<RowIndexT> ridx_b;
dh::caching_device_vector<TreePositionT> position_a;
dh::caching_device_vector<TreePositionT> position_b;
/*! \brief mapping for node id -> rows.
* This looks like:
* node id | 1 | 2 |
* rows idx | 3, 5, 1 | 13, 31 |
*/
dh::DoubleBuffer<RowIndexT> ridx;
/*! \brief mapping for row -> node id. */
dh::DoubleBuffer<TreePositionT> position;
dh::caching_device_vector<int64_t>
left_counts; // Useful to keep a bunch of zeroed memory for sort position
@@ -95,20 +108,22 @@ class RowPartitioner {
void UpdatePosition(TreePositionT nidx, TreePositionT left_nidx,
TreePositionT right_nidx, UpdatePositionOpT op) {
dh::safe_cuda(cudaSetDevice(device_idx));
Segment segment = ridx_segments.at(nidx);
Segment segment = ridx_segments.at(nidx); // rows belongs to node nidx
auto d_ridx = ridx.CurrentSpan();
auto d_position = position.CurrentSpan();
if (left_counts.size() <= nidx) {
left_counts.resize((nidx * 2) + 1);
thrust::fill(left_counts.begin(), left_counts.end(), 0);
}
// Now we divide the row segment into left and right node.
int64_t* d_left_count = left_counts.data().get() + nidx;
// Launch 1 thread for each row
dh::LaunchN<1, 128>(device_idx, segment.Size(), [=] __device__(size_t idx) {
// LaunchN starts from zero, so we restore the row index by adding segment.begin
idx += segment.begin;
RowIndexT ridx = d_ridx[idx];
// Missing value
TreePositionT new_position = op(ridx);
TreePositionT new_position = op(ridx); // new node id
KERNEL_CHECK(new_position == left_nidx || new_position == right_nidx);
AtomicIncrement(d_left_count, new_position == left_nidx);
d_position[idx] = new_position;