1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
|
/**
* Copyright 2023 by XGBoost Contributors
*/
#include <gtest/gtest.h>
#include <thrust/copy.h> // copy
#include <thrust/sequence.h> // sequence
#include <thrust/sort.h> // is_sorted
#include <algorithm> // is_sorted
#include <cstddef> // size_t
#include "../../../src/common/algorithm.cuh"
#include "../../../src/common/device_helpers.cuh"
#include "../helpers.h" // MakeCUDACtx
namespace xgboost {
namespace common {
void TestSegmentedArgSort() {
auto ctx = MakeCUDACtx(0);
size_t constexpr kElements = 100, kGroups = 3;
dh::device_vector<size_t> sorted_idx(kElements, 0);
dh::device_vector<size_t> offset_ptr(kGroups + 1, 0);
offset_ptr[0] = 0;
offset_ptr[1] = 2;
offset_ptr[2] = 78;
offset_ptr[kGroups] = kElements;
auto d_offset_ptr = dh::ToSpan(offset_ptr);
auto d_sorted_idx = dh::ToSpan(sorted_idx);
dh::LaunchN(sorted_idx.size(), [=] XGBOOST_DEVICE(size_t idx) {
auto group = dh::SegmentId(d_offset_ptr, idx);
d_sorted_idx[idx] = idx - d_offset_ptr[group];
});
dh::device_vector<float> values(kElements, 0.0f);
thrust::sequence(values.begin(), values.end(), 0.0f);
SegmentedArgSort<false, true>(&ctx, dh::ToSpan(values), d_offset_ptr, d_sorted_idx);
std::vector<size_t> h_sorted_index(sorted_idx.size());
thrust::copy(sorted_idx.begin(), sorted_idx.end(), h_sorted_index.begin());
for (size_t i = 1; i < kGroups + 1; ++i) {
auto group_sorted_idx = common::Span<size_t>(h_sorted_index)
.subspan(offset_ptr[i - 1], offset_ptr[i] - offset_ptr[i - 1]);
ASSERT_TRUE(std::is_sorted(group_sorted_idx.begin(), group_sorted_idx.end(), std::greater<>{}));
ASSERT_EQ(group_sorted_idx.back(), 0);
for (auto j : group_sorted_idx) {
ASSERT_LT(j, group_sorted_idx.size());
}
}
}
TEST(Algorithm, SegmentedArgSort) { TestSegmentedArgSort(); }
TEST(Algorithm, GpuArgSort) {
auto ctx = MakeCUDACtx(0);
dh::device_vector<float> values(20);
dh::Iota(dh::ToSpan(values), ctx.CUDACtx()->Stream()); // accending
dh::device_vector<size_t> sorted_idx(20);
ArgSort<false>(&ctx, dh::ToSpan(values), dh::ToSpan(sorted_idx)); // sort to descending
ASSERT_TRUE(thrust::is_sorted(ctx.CUDACtx()->CTP(), sorted_idx.begin(), sorted_idx.end(),
thrust::greater<size_t>{}));
dh::Iota(dh::ToSpan(values), ctx.CUDACtx()->Stream());
dh::device_vector<size_t> groups(3);
groups[0] = 0;
groups[1] = 10;
groups[2] = 20;
SegmentedArgSort<false, false>(&ctx, dh::ToSpan(values), dh::ToSpan(groups),
dh::ToSpan(sorted_idx));
ASSERT_FALSE(thrust::is_sorted(thrust::device, sorted_idx.begin(), sorted_idx.end(),
thrust::greater<size_t>{}));
ASSERT_TRUE(
thrust::is_sorted(sorted_idx.begin(), sorted_idx.begin() + 10, thrust::greater<size_t>{}));
ASSERT_TRUE(
thrust::is_sorted(sorted_idx.begin() + 10, sorted_idx.end(), thrust::greater<size_t>{}));
}
TEST(Algorithm, SegmentedSequence) {
dh::device_vector<std::size_t> idx(16);
dh::device_vector<std::size_t> ptr(3);
Context ctx = MakeCUDACtx(0);
ptr[0] = 0;
ptr[1] = 4;
ptr[2] = idx.size();
SegmentedSequence(&ctx, dh::ToSpan(ptr), dh::ToSpan(idx));
ASSERT_EQ(idx[0], 0);
ASSERT_EQ(idx[4], 0);
ASSERT_EQ(idx[3], 3);
ASSERT_EQ(idx[15], 11);
}
} // namespace common
} // namespace xgboost
|
