Skip to content

Commit

Permalink
Convert the rest.
Browse files Browse the repository at this point in the history
  • Loading branch information
@trivialfis
trivialfis committed Jul 9, 2024
1 parent 25ee82a commit 738a3d8
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 28 deletions.
53 changes: 27 additions & 26 deletions src/common/partition_builder.h
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/**
* Copyright 2021-2023 by Contributors
* Copyright 2021-2024, XGBoost Contributors
* \file row_set.h
* \brief Quick Utility to compute subset of rows
* \author Philip Cho, Tianqi Chen
Expand Down Expand Up @@ -53,16 +53,16 @@ class PartitionBuilder {
// Handle dense columns
// Analog of std::stable_partition, but in no-inplace manner
template <bool default_left, bool any_missing, typename ColumnType, typename Predicate>
inline std::pair<size_t, size_t> PartitionKernel(ColumnType* p_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) {
std::pair<size_t, size_t> PartitionKernel(ColumnType* p_column,
common::Span<const bst_idx_t> row_indices,
common::Span<bst_idx_t> left_part,
common::Span<bst_idx_t> right_part,
bst_idx_t base_rowid, Predicate&& pred) {
auto& column = *p_column;
size_t* p_left_part = left_part.data();
size_t* p_right_part = right_part.data();
size_t nleft_elems = 0;
size_t nright_elems = 0;
bst_idx_t* p_left_part = left_part.data();
bst_idx_t* p_right_part = right_part.data();
bst_idx_t nleft_elems = 0;
bst_idx_t nright_elems = 0;

auto p_row_indices = row_indices.data();
auto n_samples = row_indices.size();
Expand All @@ -89,14 +89,14 @@ class PartitionBuilder {
}

template <typename Pred>
inline 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,
inline std::pair<size_t, size_t> PartitionRangeKernel(common::Span<const bst_idx_t> ridx,
common::Span<bst_idx_t> left_part,
common::Span<bst_idx_t> right_part,
Pred pred) {
size_t* p_left_part = left_part.data();
size_t* p_right_part = right_part.data();
size_t nleft_elems = 0;
size_t nright_elems = 0;
bst_idx_t* p_left_part = left_part.data();
bst_idx_t* p_right_part = right_part.data();
bst_idx_t nleft_elems = 0;
bst_idx_t nright_elems = 0;
for (auto row_id : ridx) {
if (pred(row_id)) {
p_left_part[nleft_elems++] = row_id;
Expand All @@ -113,8 +113,8 @@ class PartitionBuilder {
GHistIndexMatrix const& gmat, const common::ColumnMatrix& column_matrix,
const RegTree& tree, bst_idx_t const* rid) {
common::Span<bst_idx_t const> 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());
common::Span<bst_idx_t> left = GetLeftBuffer(node_in_set, range.begin(), range.end());
common::Span<bst_idx_t> right = GetRightBuffer(node_in_set, range.begin(), range.end());
std::size_t nid = nodes[node_in_set].nid;
bst_feature_t fid = tree.SplitIndex(nid);
bool default_left = tree.DefaultLeft(nid);
Expand Down Expand Up @@ -183,8 +183,9 @@ class PartitionBuilder {
}

template <bool any_missing, typename ColumnType, typename Predicate>
void MaskKernel(ColumnType* p_column, common::Span<const size_t> row_indices, size_t base_rowid,
BitVector* decision_bits, BitVector* missing_bits, Predicate&& pred) {
void MaskKernel(ColumnType* p_column, common::Span<const bst_idx_t> row_indices,
bst_idx_t base_rowid, BitVector* decision_bits, BitVector* missing_bits,
Predicate&& pred) {
auto& column = *p_column;
for (auto const row_id : row_indices) {
auto const bin_id = column[row_id - base_rowid];
Expand Down Expand Up @@ -265,8 +266,8 @@ class PartitionBuilder {
const RegTree& tree, bst_idx_t const* rid, BitVector const& decision_bits,
BitVector const& missing_bits) {
common::Span<bst_idx_t const> 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());
common::Span<bst_idx_t> left = GetLeftBuffer(node_in_set, range.begin(), range.end());
common::Span<bst_idx_t> right = GetRightBuffer(node_in_set, range.begin(), range.end());
std::size_t nid = nodes[node_in_set].nid;
bool default_left = tree.DefaultLeft(nid);

Expand Down Expand Up @@ -298,12 +299,12 @@ class PartitionBuilder {
}
}

common::Span<size_t> GetLeftBuffer(int nid, size_t begin, size_t end) {
common::Span<bst_idx_t> GetLeftBuffer(int nid, size_t begin, size_t end) {
const size_t task_idx = GetTaskIdx(nid, begin);
return { mem_blocks_.at(task_idx)->Left(), end - begin };
}

common::Span<size_t> GetRightBuffer(int nid, size_t begin, size_t end) {
common::Span<bst_idx_t> GetRightBuffer(int nid, size_t begin, size_t end) {
const size_t task_idx = GetTaskIdx(nid, begin);
return { mem_blocks_.at(task_idx)->Right(), end - begin };
}
Expand Down Expand Up @@ -345,7 +346,7 @@ class PartitionBuilder {
}
}

void MergeToArray(int nid, size_t begin, bst_idx_t* rows_indexes) {
void MergeToArray(bst_node_t nid, size_t begin, bst_idx_t* rows_indexes) {
size_t task_idx = GetTaskIdx(nid, begin);

bst_idx_t* left_result = rows_indexes + mem_blocks_[task_idx]->n_offset_left;
Expand Down
4 changes: 2 additions & 2 deletions tests/cpp/common/test_partition_builder.cc
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/**
* Copyright 2020-2023 by XGBoost contributors
* Copyright 2020-2024, XGBoost contributors
*/
#include <gtest/gtest.h>

Expand Down Expand Up @@ -58,7 +58,7 @@ TEST(PartitionBuilder, BasicTest) {
}
builder.CalculateRowOffsets();

std::vector<size_t> v(*std::max_element(tasks.begin(), tasks.end()) * kBlockSize);
std::vector<bst_idx_t> v(*std::max_element(tasks.begin(), tasks.end()) * kBlockSize);

for(size_t nid = 0; nid < kNodes; ++nid) {

Expand Down

0 comments on commit 738a3d8

Please sign in to comment.