-
-
Notifications
You must be signed in to change notification settings - Fork 8.7k
/
ranking_utils.cuh
84 lines (76 loc) · 2.45 KB
/
ranking_utils.cuh
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
/*!
* Copyright 2021 by XGBoost Contributors
*/
#ifndef XGBOOST_COMMON_RANKING_UTILS_H_
#define XGBOOST_COMMON_RANKING_UTILS_H_
#include <cub/cub.cuh>
#include "xgboost/base.h"
#include "device_helpers.cuh"
#include "./math.h"
namespace xgboost {
namespace common {
/**
* \param n Number of items (length of the base)
* \param h hight
*/
XGBOOST_DEVICE inline size_t DiscreteTrapezoidArea(size_t n, size_t h) {
n -= 1; // without diagonal entries
h = std::min(n, h); // Specific for ranking.
size_t total = ((n - (h - 1)) + n) * h / 2;
return total;
}
/**
* Used for mapping many groups of trapezoid shaped computation onto CUDA blocks. The
* trapezoid must be on upper right corner.
*
* Equivalent to loops like:
*
* \code
* for (size i = 0; i < h; ++i) {
* for (size_t j = i + 1; j < n; ++j) {
* do_something();
* }
* }
* \endcode
*/
template <typename U>
inline size_t
SegmentedTrapezoidThreads(xgboost::common::Span<U> group_ptr,
xgboost::common::Span<size_t> out_group_threads_ptr,
size_t h) {
CHECK_GE(group_ptr.size(), 1);
CHECK_EQ(group_ptr.size(), out_group_threads_ptr.size());
dh::LaunchN(
group_ptr.size(), [=] XGBOOST_DEVICE(size_t idx) {
if (idx == 0) {
out_group_threads_ptr[0] = 0;
return;
}
size_t cnt = static_cast<size_t>(group_ptr[idx] - group_ptr[idx - 1]);
out_group_threads_ptr[idx] = DiscreteTrapezoidArea(cnt, h);
});
dh::InclusiveSum(out_group_threads_ptr.data(), out_group_threads_ptr.data(),
out_group_threads_ptr.size());
size_t total = 0;
dh::safe_cuda(cudaMemcpy(
&total, out_group_threads_ptr.data() + out_group_threads_ptr.size() - 1,
sizeof(total), cudaMemcpyDeviceToHost));
return total;
}
/**
* Called inside kernel to obtain coordinate from trapezoid grid.
*/
XGBOOST_DEVICE inline void UnravelTrapeziodIdx(size_t i_idx, size_t n,
size_t *out_i, size_t *out_j) {
auto &i = *out_i;
auto &j = *out_j;
double idx = static_cast<double>(i_idx);
double N = static_cast<double>(n);
i = std::ceil(-(0.5 - N + std::sqrt(common::Sqr(N - 0.5) + 2.0 * (-idx - 1.0)))) - 1.0;
auto I = static_cast<double>(i);
size_t n_elems = -0.5 * common::Sqr(I) + (N - 0.5) * I;
j = idx - n_elems + i + 1;
}
} // namespace common
} // namespace xgboost
#endif // XGBOOST_COMMON_RANKING_UTILS_H_