6 #ifndef XGBOOST_LINALG_H_
7 #define XGBOOST_LINALG_H_
9 #include <dmlc/endian.h>
19 #include <type_traits>
25 #if defined(__CUDA__) || defined(__NVCC__)
26 #define LINALG_HD __host__ __device__
29 #endif // defined (__CUDA__) || defined(__NVCC__)
39 return (std::is_floating_point<T>::value
41 : (std::is_integral<T>::value ? (std::is_signed<T>::value ?
'i' :
'u') :
'\0'));
45 template <
size_t dim,
typename S,
typename Head,
size_t D>
46 constexpr
size_t Offset(S (&strides)[D],
size_t n, Head head) {
47 static_assert(dim < D,
"");
48 return n + head * strides[dim];
51 template <
size_t dim,
typename S,
size_t D,
typename Head,
typename... Tail>
52 constexpr std::enable_if_t<
sizeof...(Tail) != 0,
size_t>
Offset(S (&strides)[D],
size_t n,
53 Head head, Tail &&...rest) {
54 static_assert(dim < D,
"");
55 return Offset<dim + 1>(strides, n + (head * strides[dim]), std::forward<Tail>(rest)...);
58 template <
int32_t D,
bool f_array = false>
59 constexpr
void CalcStride(
size_t const (&shape)[D],
size_t (&stride)[D]) {
62 for (int32_t s = 1; s < D; ++s) {
63 stride[s] = shape[s - 1] * stride[s - 1];
67 for (int32_t s = D - 2; s >= 0; --s) {
68 stride[s] = shape[s + 1] * stride[s + 1];
89 return std::is_same<T, IntTag>::value ? 0 : 1;
92 template <
typename T,
typename... S>
93 constexpr std::enable_if_t<
sizeof...(S) != 0, int32_t>
CalcSliceDim() {
98 constexpr
size_t CalcSize(
size_t (&shape)[D]) {
100 for (
auto d : shape) {
106 template <
typename S>
109 template <
typename S>
110 using IndexToTag = std::conditional_t<std::is_integral<RemoveCRType<S>>::value,
IntTag, S>;
112 template <
int32_t n,
typename Fn>
114 #if defined __CUDA_ARCH__
116 #endif // defined __CUDA_ARCH__
117 for (int32_t i = 0; i < n; ++i) {
122 template <
typename T>
125 for (; v != 0; v &= v - 1) c++;
130 #if defined(__CUDA_ARCH__)
132 #elif defined(__GNUC__) || defined(__clang__)
133 return __builtin_popcount(v);
134 #elif defined(_MSC_VER)
142 #if defined(__CUDA_ARCH__)
144 #elif defined(__GNUC__) || defined(__clang__)
145 return __builtin_popcountll(v);
146 #elif defined(_MSC_VER)
147 return __popcnt64(v);
153 template <
class T, std::size_t N, std::size_t... Idx>
154 constexpr
auto Arr2Tup(T (&arr)[N], std::index_sequence<Idx...>) {
155 return std::make_tuple(arr[Idx]...);
158 template <
class T, std::
size_t N>
160 return Arr2Tup(arr, std::make_index_sequence<N>{});
166 template <
typename I,
int32_t D>
169 static_assert(std::is_signed<decltype(D)>::value,
170 "Don't change the type without changing the for loop.");
171 for (int32_t dim = D; --dim > 0;) {
172 auto s =
static_cast<std::remove_const_t<std::remove_reference_t<I>
>>(shape[dim]);
175 index[dim] = idx - t * s;
178 index[dim] = idx & (s - 1);
186 template <
size_t dim,
typename I,
int32_t D>
188 static_assert(dim < D,
"");
192 template <
size_t dim, int32_t D,
typename... S,
typename I,
193 std::enable_if_t<
sizeof...(S) != 0> * =
nullptr>
195 static_assert(dim < D,
"");
197 ReshapeImpl<dim + 1>(out_shape, std::forward<S>(rest)...);
200 template <
typename Fn,
typename Tup,
size_t... I>
202 return f(std::get<I>(t)...);
211 template <
typename Fn,
typename Tup>
213 constexpr
auto kSize = std::tuple_size<Tup>::value;
214 return Apply(std::forward<Fn>(f), std::forward<Tup>(t), std::make_index_sequence<kSize>{});
225 template <
typename I>
243 template <
typename T,
int32_t kDim>
267 template <
size_t old_dim,
size_t new_dim,
int32_t D,
typename I>
268 LINALG_HD size_t MakeSliceDim(
size_t new_shape[D],
size_t new_stride[D],
269 detail::RangeTag<I> &&range)
const {
270 static_assert(new_dim < D,
"");
271 static_assert(old_dim < kDim,
"");
272 new_stride[new_dim] = stride_[old_dim];
273 new_shape[new_dim] = range.Size();
274 assert(
static_cast<decltype(shape_[old_dim])
>(range.end) <= shape_[old_dim]);
276 auto offset = stride_[old_dim] * range.beg;
282 template <
size_t old_dim,
size_t new_dim, int32_t D,
typename I,
typename... S>
283 LINALG_HD size_t MakeSliceDim(
size_t new_shape[D],
size_t new_stride[D],
284 detail::RangeTag<I> &&range, S &&...slices)
const {
285 static_assert(new_dim < D,
"");
286 static_assert(old_dim < kDim,
"");
287 new_stride[new_dim] = stride_[old_dim];
288 new_shape[new_dim] = range.Size();
289 assert(
static_cast<decltype(shape_[old_dim])
>(range.end) <= shape_[old_dim]);
291 auto offset = stride_[old_dim] * range.beg;
292 return MakeSliceDim<old_dim + 1, new_dim + 1, D>(new_shape, new_stride,
293 std::forward<S>(slices)...) +
297 template <
size_t old_dim,
size_t new_dim,
int32_t D>
298 LINALG_HD size_t MakeSliceDim(
size_t new_shape[D],
size_t new_stride[D], detail::AllTag)
const {
299 static_assert(new_dim < D,
"");
300 static_assert(old_dim < kDim,
"");
301 new_stride[new_dim] = stride_[old_dim];
302 new_shape[new_dim] = shape_[old_dim];
308 template <
size_t old_dim,
size_t new_dim, int32_t D,
typename... S>
309 LINALG_HD size_t MakeSliceDim(
size_t new_shape[D],
size_t new_stride[D], detail::AllTag,
310 S &&...slices)
const {
311 static_assert(new_dim < D,
"");
312 static_assert(old_dim < kDim,
"");
313 new_stride[new_dim] = stride_[old_dim];
314 new_shape[new_dim] = shape_[old_dim];
315 return MakeSliceDim<old_dim + 1, new_dim + 1, D>(new_shape, new_stride,
316 std::forward<S>(slices)...);
319 template <
size_t old_dim,
size_t new_dim,
int32_t D,
typename Index>
320 LINALG_HD size_t MakeSliceDim(
size_t new_shape[D],
size_t new_stride[D], Index i)
const {
321 static_assert(old_dim < kDim,
"");
322 return stride_[old_dim] * i;
327 template <
size_t old_dim,
size_t new_dim, int32_t D,
typename Index,
typename... S>
328 LINALG_HD std::enable_if_t<std::is_integral<Index>::value,
size_t> MakeSliceDim(
329 size_t new_shape[D],
size_t new_stride[D], Index i, S &&...slices)
const {
330 static_assert(old_dim < kDim,
"");
331 auto offset = stride_[old_dim] * i;
333 MakeSliceDim<old_dim + 1, new_dim, D>(new_shape, new_stride, std::forward<S>(slices)...);
353 template <
typename I,
int32_t D>
355 : data_{data}, ptr_{data_.
data()}, device_{device} {
356 static_assert(D > 0 && D <= kDim,
"Invalid shape.");
358 detail::UnrollLoop<D>([&](
auto i) { shape_[i] = shape[i]; });
359 for (
auto i = D; i < kDim; ++i) {
372 template <
typename I,
int32_t D>
375 : data_{data}, ptr_{data_.
data()}, device_{device} {
376 static_assert(D == kDim,
"Invalid shape & stride.");
377 detail::UnrollLoop<D>([&](
auto i) {
378 shape_[i] = shape[i];
379 stride_[i] = stride[i];
386 std::enable_if_t<common::detail::IsAllowedElementTypeConversion<U, T>::value> * =
nullptr>
388 : data_{that.
Values()}, ptr_{data_.
data()}, size_{that.Size()}, device_{that.DeviceIdx()} {
389 detail::UnrollLoop<kDim>([&](
auto i) {
390 stride_[i] = that.Stride(i);
391 shape_[i] = that.Shape(i);
408 template <
typename... Index>
410 static_assert(
sizeof...(index) <= kDim,
"Invalid index.");
411 size_t offset = detail::Offset<0ul>(stride_, 0ul, std::forward<Index>(index)...);
412 assert(offset < data_.
size() &&
"Out of bound access.");
418 template <
typename... Index>
420 static_assert(
sizeof...(index) <= kDim,
"Invalid index.");
421 size_t offset = detail::Offset<0ul>(stride_, 0ul, std::forward<Index>(index)...);
422 assert(offset < data_.
size() &&
"Out of bound access.");
439 template <
typename... S>
441 static_assert(
sizeof...(slices) <= kDim,
"Invalid slice.");
442 int32_t constexpr kNewDim{detail::CalcSliceDim<detail::IndexToTag<S>...>()};
443 size_t new_shape[kNewDim];
444 size_t new_stride[kNewDim];
445 auto offset = MakeSliceDim<0, 0, kNewDim>(new_shape, new_stride, std::forward<S>(slices)...);
479 static_assert(std::is_same<decltype(stride), decltype(stride_)>::value,
"");
489 static_assert(std::is_same<decltype(stride), decltype(stride_)>::value,
"");
491 detail::CalcStride<kDim, true>(shape_, stride);
507 template <
typename Container,
typename I, int32_t D,
508 std::enable_if_t<!common::detail::IsSpan<Container>::value> * =
nullptr>
510 using T =
typename Container::value_type;
514 template <
typename T,
typename I,
int32_t D>
524 if (idx > std::numeric_limits<uint32_t>::max()) {
525 return detail::UnravelImpl<uint64_t, D>(
static_cast<uint64_t
>(idx), shape);
527 return detail::UnravelImpl<uint32_t, D>(
static_cast<uint32_t
>(idx), shape);
536 template <
typename T>
546 template <
typename T>
547 auto MakeVec(T *ptr,
size_t s, int32_t device = -1) {
551 template <
typename T>
557 template <
typename T>
568 template <
typename T>
577 template <
typename T,
int32_t D>
580 array_interface[
"data"] = std::vector<Json>(2);
581 array_interface[
"data"][0] =
Integer{
reinterpret_cast<int64_t
>(t.
Values().data())};
582 array_interface[
"data"][1] =
Boolean{
true};
585 array_interface[
"stream"] =
Null{};
587 std::vector<Json> shape(t.
Shape().size());
588 std::vector<Json> stride(t.
Stride().size());
589 for (
size_t i = 0; i < t.
Shape().size(); ++i) {
593 array_interface[
"shape"] =
Array{shape};
594 array_interface[
"strides"] =
Array{stride};
595 array_interface[
"version"] = 3;
597 char constexpr kT = detail::ArrayInterfaceHandler::TypeChar<T>();
598 static_assert(kT !=
'\0',
"");
599 if (DMLC_LITTLE_ENDIAN) {
600 array_interface[
"typestr"] =
String{
"<" + (kT + std::to_string(
sizeof(T)))};
602 array_interface[
"typestr"] =
String{
">" + (kT + std::to_string(
sizeof(T)))};
604 return array_interface;
610 template <
typename T,
int32_t D>
614 res[
"data"][1] =
Boolean{
false};
621 template <
typename T,
int32_t D>
628 template <
typename T,
int32_t D>
639 template <
typename T,
int32_t kDim = 5>
649 template <
typename I, std::
int32_t D>
650 void Initialize(I
const (&shape)[D], std::int32_t device) {
651 static_assert(D <= kDim,
"Invalid shape.");
652 std::copy(shape, shape + D, shape_);
653 for (
auto i = D; i < kDim; ++i) {
672 template <
typename I,
int32_t D>
673 explicit Tensor(I
const (&shape)[D], int32_t device) {
675 std::copy(shape, shape + D, shape_);
676 for (
auto i = D; i < kDim; ++i) {
691 template <
typename It,
typename I,
int32_t D>
692 explicit Tensor(It begin, It end, I
const (&shape)[D], int32_t device) {
694 h_vec.insert(h_vec.begin(), begin, end);
696 this->Initialize(shape, device);
699 template <
typename I,
int32_t D>
700 explicit Tensor(std::initializer_list<T> data, I
const (&shape)[D], int32_t device) {
704 this->Initialize(shape, device);
714 return {span, shape_, device};
717 return {span, shape_, device};
724 return {span, shape_, device};
727 return {span, shape_, device};
736 auto Shape(
size_t i)
const {
return shape_[i]; }
747 template <
typename Fn>
751 <<
"Inconsistent size after modification.";
759 template <
typename... S>
761 static_assert(
sizeof...(S) <= kDim,
"Invalid shape.");
762 detail::ReshapeImpl<0>(shape_, std::forward<S>(s)...);
763 auto constexpr kEnd =
sizeof...(S);
764 static_assert(kEnd <= kDim,
"Invalid shape.");
765 std::fill(shape_ + kEnd, shape_ + kDim, 1);
777 static_assert(D <= kDim,
"Invalid shape.");
778 std::copy(shape, shape + D, this->shape_);
779 std::fill(shape_ + D, shape_ + kDim, 1);
792 template <
typename T,
int32_t D>
798 for (
size_t i = 1; i < D; ++i) {
800 shape[i] = r.
Shape(i);
802 CHECK_EQ(shape[i], r.
Shape(i));
812 #if defined(LINALG_HD)
814 #endif // defined(LINALG_HD)
815 #endif // XGBOOST_LINALG_H_