Skip to content

Commit

Permalink
Format.
Browse files Browse the repository at this point in the history
  • Loading branch information
@trivialfis
trivialfis committed Oct 29, 2021
1 parent d84700b commit dde3ec4
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 51 deletions.
74 changes: 34 additions & 40 deletions include/xgboost/linalg.h
View file
Expand Up @@ -6,10 +6,12 @@
#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>
Expand All @@ -24,8 +26,7 @@ constexpr size_t Offset(S (&strides)[D], size_t n, size_t dim, Head head) {
}

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) {
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...);
}
Expand All @@ -36,7 +37,8 @@ struct IntTag {};
/**
* \brief Calculate the dimension of sliced tensor.
*/
template <typename T> constexpr int32_t CalcSliceDim() {
template <typename T>
constexpr int32_t CalcSliceDim() {
return std::is_same<T, IntTag>::value ? 0 : 1;
}

Expand All @@ -49,8 +51,7 @@ 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>;
using IndexToTag = std::conditional_t<std::is_integral<RemoveCRType<S>>::value, IntTag, AllTag>;

template <int32_t n, typename Fn>
constexpr XGBOOST_DEVICE auto UnrollLoop(Fn fn) {
Expand All @@ -74,7 +75,8 @@ constexpr detail::AllTag All() { return {}; }
* \brief A tensor view with static type and shape, it implements indexing and
* slicing.
*/
template <typename T, int32_t kDim = 5> class TensorView {
template <typename T, int32_t kDim = 5>
class TensorView {
using ShapeT = size_t[kDim];
using StrideT = ShapeT;

Expand Down Expand Up @@ -102,40 +104,34 @@ template <typename T, int32_t kDim = 5> class TensorView {
};

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 {
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],
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...);
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],
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 {
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...);
auto res = MakeSliceDim<D>(old_dim + 1, new_dim, new_shape, new_stride, slices...);
return {res.old_dim, res.new_dim, res.offset + offset};
}

Expand All @@ -148,8 +144,7 @@ template <typename T, int32_t kDim = 5> class TensorView {
* \brief Create a tensor with data and shape.
*/
template <typename I, int32_t D>
XGBOOST_DEVICE TensorView(common::Span<T> data, I const (&shape)[D],
int32_t device)
XGBOOST_DEVICE TensorView(common::Span<T> data, I const (&shape)[D], int32_t device)
: data_{data}, device_{device} {
static_assert(D > 0 && D <= kDim, "Invalid shape.");
// shape
Expand All @@ -170,8 +165,8 @@ template <typename T, int32_t kDim = 5> class TensorView {
* 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)
XGBOOST_DEVICE TensorView(common::Span<T> data, I const (&shape)[D], I const (&stride)[D],
int32_t device)
: data_{data}, device_{device} {
static_assert(D == kDim, "Invalid shape & stride.");
detail::UnrollLoop<D>([&](auto i) {
Expand Down Expand Up @@ -202,7 +197,8 @@ template <typename T, int32_t kDim = 5> class TensorView {
*
* \endcode
*/
template <typename... Index> XGBOOST_DEVICE T &operator()(Index &&...index) {
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 data_[offset];
Expand All @@ -229,29 +225,27 @@ template <typename T, int32_t kDim = 5> class TensorView {
*
* \endcode
*/
template <typename... S> XGBOOST_DEVICE auto Slice(S &&...slices) const {
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...);
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_};
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_};
}
XGBOOST_DEVICE auto Shape() const { return common::Span<size_t const, kDim>{shape_}; }
XGBOOST_DEVICE auto Shape(size_t i) const { return shape_[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 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_; }
Expand All @@ -273,6 +267,6 @@ using VectorView = TensorView<T, 1>;
*/
template <typename T>
using MatrixView = TensorView<T, 2>;
} // namespace linalg
} // namespace xgboost
} // namespace linalg
} // namespace xgboost
#endif // XGBOOST_LINALG_H_
18 changes: 7 additions & 11 deletions tests/cpp/common/test_linalg.cc
View file
@@ -1,19 +1,18 @@
#include <gtest/gtest.h>
#include <numeric>
#include <xgboost/host_device_vector.h>
#include <xgboost/linalg.h>

#include <numeric>

namespace xgboost {
namespace linalg {
auto MakeMatrixFromTest(HostDeviceVector<float> *storage, size_t n_rows,
size_t n_cols) {
auto MakeMatrixFromTest(HostDeviceVector<float> *storage, size_t n_rows, size_t n_cols) {
storage->Resize(n_rows * n_cols);
auto &h_storage = storage->HostVector();

std::iota(h_storage.begin(), h_storage.end(), 0);

auto m = linalg::TensorView<float, 2>{
h_storage, {n_rows, static_cast<size_t>(n_cols)}, -1};
auto m = linalg::TensorView<float, 2>{h_storage, {n_rows, static_cast<size_t>(n_cols)}, -1};
return m;
}

Expand Down Expand Up @@ -57,10 +56,7 @@ TEST(Linalg, Tensor) {
ASSERT_EQ(s.Shape()[1], 4);

std::vector<std::vector<double>> sol{
{12.0, 13.0, 14.0, 15.0},
{16.0, 17.0, 18.0, 19.0},
{20.0, 21.0, 22.0, 23.0}
};
{12.0, 13.0, 14.0, 15.0}, {16.0, 17.0, 18.0, 19.0}, {20.0, 21.0, 22.0, 23.0}};
for (size_t i = 0; i < s.Shape()[0]; ++i) {
for (size_t j = 0; j < s.Shape()[1]; ++j) {
ASSERT_EQ(s(i, j), sol[i][j]);
Expand Down Expand Up @@ -113,5 +109,5 @@ TEST(Linalg, Empty) {
ASSERT_EQ(s.Shape(0), 0);
}
}
} // namespace linalg
} // namespace xgboost
} // namespace linalg
} // namespace xgboost

0 comments on commit dde3ec4

Please sign in to comment.