forked from dmlc/xgboost
-
Notifications
You must be signed in to change notification settings - Fork 0
/
linalg.h
301 lines (269 loc) · 9.33 KB
/
linalg.h
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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
/*!
* Copyright 2021 by XGBoost Contributors
* \file linalg.h
* \brief Linear algebra related utilities.
*/
#ifndef XGBOOST_LINALG_H_
#define XGBOOST_LINALG_H_
#include <xgboost/base.h>
#include <xgboost/generic_parameters.h>
#include <xgboost/span.h>
#include <algorithm>
#include <cassert>
#include <type_traits>
#include <utility>
#include <vector>
namespace xgboost {
namespace linalg {
namespace detail {
template <typename S, typename Head, size_t D>
constexpr size_t Offset(S (&strides)[D], size_t n, size_t dim, Head head) {
assert(dim < D);
return n + head * strides[dim];
}
template <typename S, size_t D, typename Head, typename... Tail>
constexpr size_t Offset(S (&strides)[D], size_t n, size_t dim, Head head, Tail &&...rest) {
assert(dim < D);
return Offset(strides, n + (head * strides[dim]), dim + 1, rest...);
}
struct AllTag {};
struct IntTag {};
/**
* \brief Calculate the dimension of sliced tensor.
*/
template <typename T>
constexpr int32_t CalcSliceDim() {
return std::is_same<T, IntTag>::value ? 0 : 1;
}
template <typename T, typename... S>
constexpr std::enable_if_t<sizeof...(S) != 0, int32_t> CalcSliceDim() {
return CalcSliceDim<T>() + CalcSliceDim<S...>();
}
template <int32_t D>
constexpr size_t CalcSize(size_t (&shape)[D]) {
size_t size = 1;
for (auto d : shape) {
size *= d;
}
return size;
}
template <typename S>
using RemoveCRType = std::remove_const_t<std::remove_reference_t<S>>;
template <typename S>
using IndexToTag = std::conditional_t<std::is_integral<RemoveCRType<S>>::value, IntTag, AllTag>;
template <int32_t n, typename Fn>
XGBOOST_DEVICE constexpr auto UnrollLoop(Fn fn) {
#if defined __CUDA_ARCH__
#pragma unroll n
#endif // defined __CUDA_ARCH__
for (int32_t i = 0; i < n; ++i) {
fn(i);
}
}
} // namespace detail
/**
* \brief Specify all elements in the axis is used for slice.
*/
constexpr detail::AllTag All() { return {}; }
/**
* \brief A tensor view with static type and shape. It implements indexing and slicing.
*
* Most of the algorithms in XGBoost are implemented for both CPU and GPU without using
* much linear algebra routines, this class is a helper intended to ease some high level
* operations like indexing into prediction tensor or gradient matrix. It can be passed
* into CUDA kernel as normal argument for GPU algorithms.
*/
template <typename T, int32_t kDim = 5>
class TensorView {
public:
using ShapeT = size_t[kDim];
using StrideT = ShapeT;
private:
StrideT stride_{1};
ShapeT shape_{0};
common::Span<T> data_;
T* ptr_{nullptr}; // pointer of data_ to avoid bound check.
size_t size_{0};
int32_t device_{-1};
// Unlike `Tensor`, the data_ can have arbitrary size since this is just a view.
XGBOOST_DEVICE void CalcSize() {
if (data_.empty()) {
size_ = 0;
} else {
size_ = detail::CalcSize(shape_);
}
}
struct SliceHelper {
size_t old_dim;
size_t new_dim;
size_t offset;
};
template <int32_t D, typename... S>
XGBOOST_DEVICE SliceHelper MakeSliceDim(size_t old_dim, size_t new_dim, size_t new_shape[D],
size_t new_stride[D], detail::AllTag) const {
new_stride[new_dim] = stride_[old_dim];
new_shape[new_dim] = shape_[old_dim];
return {old_dim + 1, new_dim + 1, 0};
}
template <int32_t D, typename... S>
XGBOOST_DEVICE SliceHelper MakeSliceDim(size_t old_dim, size_t new_dim, size_t new_shape[D],
size_t new_stride[D], detail::AllTag,
S &&...slices) const {
new_stride[new_dim] = stride_[old_dim];
new_shape[new_dim] = shape_[old_dim];
return MakeSliceDim<D>(old_dim + 1, new_dim + 1, new_shape, new_stride, slices...);
}
template <int32_t D, typename Index>
XGBOOST_DEVICE SliceHelper MakeSliceDim(size_t old_dim, size_t new_dim, size_t new_shape[D],
size_t new_stride[D], Index i) const {
return {old_dim + 1, new_dim, stride_[old_dim] * i};
}
template <int32_t D, typename Index, typename... S>
XGBOOST_DEVICE std::enable_if_t<std::is_integral<Index>::value, SliceHelper> MakeSliceDim(
size_t old_dim, size_t new_dim, size_t new_shape[D], size_t new_stride[D], Index i,
S &&...slices) const {
auto offset = stride_[old_dim] * i;
auto res = MakeSliceDim<D>(old_dim + 1, new_dim, new_shape, new_stride, slices...);
return {res.old_dim, res.new_dim, res.offset + offset};
}
public:
size_t constexpr static kValueSize = sizeof(T);
size_t constexpr static kDimension = kDim;
public:
/**
* \brief Create a tensor with data and shape.
*
* \tparam I Type of the shape array element.
* \tparam D Size of the shape array, can be lesser than or equal to tensor dimension.
*
* \param data Raw data input, can be const if this tensor has const type in its
* template parameter.
* \param shape shape of the tensor
* \param device Device ordinal
*/
template <typename I, int32_t D>
XGBOOST_DEVICE TensorView(common::Span<T> data, I const (&shape)[D], int32_t device)
: data_{data}, ptr_{data_.data()}, device_{device} {
static_assert(D > 0 && D <= kDim, "Invalid shape.");
// shape
detail::UnrollLoop<D>([&](auto i) { shape_[i] = shape[i]; });
for (auto i = D; i < kDim; ++i) {
shape_[i] = 1;
}
// stride
stride_[kDim - 1] = 1;
for (int32_t s = kDim - 2; s >= 0; --s) {
stride_[s] = shape_[s + 1] * stride_[s + 1];
}
this->CalcSize();
};
/**
* \brief Create a tensor with data, shape and strides. Don't use this constructor if
* stride can be calculated from shape.
*/
template <typename I, int32_t D>
XGBOOST_DEVICE TensorView(common::Span<T> data, I const (&shape)[D], I const (&stride)[D],
int32_t device)
: data_{data}, ptr_{data_.data()}, device_{device} {
static_assert(D == kDim, "Invalid shape & stride.");
detail::UnrollLoop<D>([&](auto i) {
shape_[i] = shape[i];
stride_[i] = stride[i];
});
this->CalcSize();
};
XGBOOST_DEVICE TensorView(TensorView const &that)
: data_{that.data_}, ptr_{data_.data()}, size_{that.size_}, device_{that.device_} {
detail::UnrollLoop<kDim>([&](auto i) {
stride_[i] = that.stride_[i];
shape_[i] = that.shape_[i];
});
}
/**
* \brief Index the tensor to obtain a scalar value.
*
* \code
*
* // Create a 3-dim tensor.
* Tensor<float, 3> t {data, shape, 0};
* float pi = 3.14159;
* t(1, 2, 3) = pi;
* ASSERT_EQ(t(1, 2, 3), pi);
*
* \endcode
*/
template <typename... Index>
XGBOOST_DEVICE T &operator()(Index &&...index) {
static_assert(sizeof...(index) <= kDim, "Invalid index.");
size_t offset = detail::Offset(stride_, 0ul, 0ul, index...);
return ptr_[offset];
}
/**
* \brief Index the tensor to obtain a scalar value.
*/
template <typename... Index>
XGBOOST_DEVICE T const &operator()(Index &&...index) const {
static_assert(sizeof...(index) <= kDim, "Invalid index.");
size_t offset = detail::Offset(stride_, 0ul, 0ul, index...);
return ptr_[offset];
}
/**
* \brief Slice the tensor. The returned tensor has inferred dim and shape.
*
* \code
*
* // Create a 3-dim tensor.
* Tensor<float, 3> t {data, shape, 0};
* // s has 2 dimensions (matrix)
* auto s = t.Slice(1, All(), All());
*
* \endcode
*/
template <typename... S>
XGBOOST_DEVICE auto Slice(S &&...slices) const {
static_assert(sizeof...(slices) <= kDim, "Invalid slice.");
int32_t constexpr kNewDim{detail::CalcSliceDim<detail::IndexToTag<S>...>()};
size_t new_shape[kNewDim];
size_t new_stride[kNewDim];
auto res = MakeSliceDim<kNewDim>(size_t(0), size_t(0), new_shape, new_stride, slices...);
// ret is a different type due to changed dimension, so we can not access its private
// fields.
TensorView<T, kNewDim> ret{data_.subspan(data_.empty() ? 0 : res.offset), new_shape, new_stride,
device_};
return ret;
}
XGBOOST_DEVICE auto Shape() const { return common::Span<size_t const, kDim>{shape_}; }
/**
* Get the shape for i^th dimension
*/
XGBOOST_DEVICE auto Shape(size_t i) const { return shape_[i]; }
XGBOOST_DEVICE auto Stride() const { return common::Span<size_t const, kDim>{stride_}; }
/**
* Get the stride for i^th dimension, stride is specified as number of items instead of bytes.
*/
XGBOOST_DEVICE auto Stride(size_t i) const { return stride_[i]; }
XGBOOST_DEVICE auto cbegin() const { return data_.cbegin(); } // NOLINT
XGBOOST_DEVICE auto cend() const { return data_.cend(); } // NOLINT
XGBOOST_DEVICE auto begin() { return data_.begin(); } // NOLINT
XGBOOST_DEVICE auto end() { return data_.end(); } // NOLINT
XGBOOST_DEVICE size_t Size() const { return size_; }
XGBOOST_DEVICE auto Values() const { return data_; }
XGBOOST_DEVICE auto DeviceIdx() const { return device_; }
};
/**
* \brief A view over a vector, specialization of Tensor
*
* \tparam T data type of vector
*/
template <typename T>
using VectorView = TensorView<T, 1>;
/**
* \brief A view over a matrix, specialization of Tensor.
*
* \tparam T data type of matrix
*/
template <typename T>
using MatrixView = TensorView<T, 2>;
} // namespace linalg
} // namespace xgboost
#endif // XGBOOST_LINALG_H_