Support F order for the tensor type. (#8872)

- Add F order support for tensor and view.
- Use parameter pack for automatic type cast. (avoid excessive static cast for shape).
This commit is contained in:
Jiaming Yuan
2023-03-08 03:27:49 +08:00
committed by GitHub
parent f53055f75e
commit f236640427
9 changed files with 194 additions and 94 deletions

View File

@@ -6,17 +6,18 @@
#include <xgboost/host_device_vector.h>
#include <xgboost/linalg.h>
#include <numeric>
#include <cstddef> // size_t
#include <numeric> // iota
#include <vector>
#include "../../../src/common/linalg_op.h"
namespace xgboost {
namespace linalg {
namespace xgboost::linalg {
namespace {
auto kCpuId = Context::kCpuId;
}
auto MakeMatrixFromTest(HostDeviceVector<float> *storage, size_t n_rows, size_t n_cols) {
auto MakeMatrixFromTest(HostDeviceVector<float> *storage, std::size_t n_rows, std::size_t n_cols) {
storage->Resize(n_rows * n_cols);
auto &h_storage = storage->HostVector();
@@ -48,10 +49,11 @@ TEST(Linalg, VectorView) {
}
TEST(Linalg, TensorView) {
Context ctx;
std::vector<double> data(2 * 3 * 4, 0);
std::iota(data.begin(), data.end(), 0);
auto t = MakeTensorView(data, {2, 3, 4}, -1);
auto t = MakeTensorView(&ctx, data, 2, 3, 4);
ASSERT_EQ(t.Shape()[0], 2);
ASSERT_EQ(t.Shape()[1], 3);
ASSERT_EQ(t.Shape()[2], 4);
@@ -106,12 +108,12 @@ TEST(Linalg, TensorView) {
{
// Don't assign the initial dimension, tensor should be able to deduce the correct dim
// for Slice.
auto t = MakeTensorView(data, {2, 3, 4}, 0);
auto t = MakeTensorView(&ctx, data, 2, 3, 4);
auto s = t.Slice(1, 2, All());
static_assert(decltype(s)::kDimension == 1);
}
{
auto t = MakeTensorView(data, {2, 3, 4}, 0);
auto t = MakeTensorView(&ctx, data, 2, 3, 4);
auto s = t.Slice(1, linalg::All(), 1);
ASSERT_EQ(s(0), 13);
ASSERT_EQ(s(1), 17);
@@ -119,7 +121,7 @@ TEST(Linalg, TensorView) {
}
{
// range slice
auto t = MakeTensorView(data, {2, 3, 4}, 0);
auto t = MakeTensorView(&ctx, data, 2, 3, 4);
auto s = t.Slice(linalg::All(), linalg::Range(1, 3), 2);
static_assert(decltype(s)::kDimension == 2);
std::vector<double> sol{6, 10, 18, 22};
@@ -134,7 +136,7 @@ TEST(Linalg, TensorView) {
}
{
// range slice
auto t = MakeTensorView(data, {2, 3, 4}, 0);
auto t = MakeTensorView(&ctx, data, 2, 3, 4);
auto s = t.Slice(1, linalg::Range(1, 3), linalg::Range(1, 3));
static_assert(decltype(s)::kDimension == 2);
std::vector<double> sol{17, 18, 21, 22};
@@ -149,7 +151,7 @@ TEST(Linalg, TensorView) {
}
{
// same as no slice.
auto t = MakeTensorView(data, {2, 3, 4}, 0);
auto t = MakeTensorView(&ctx, data, 2, 3, 4);
auto s = t.Slice(linalg::All(), linalg::Range(0, 3), linalg::Range(0, 4));
static_assert(decltype(s)::kDimension == 3);
auto all = t.Slice(linalg::All(), linalg::All(), linalg::All());
@@ -166,7 +168,7 @@ TEST(Linalg, TensorView) {
{
// copy and move constructor.
auto t = MakeTensorView(data, {2, 3, 4}, kCpuId);
auto t = MakeTensorView(&ctx, data, 2, 3, 4);
auto from_copy = t;
auto from_move = std::move(t);
for (size_t i = 0; i < t.Shape().size(); ++i) {
@@ -177,7 +179,7 @@ TEST(Linalg, TensorView) {
{
// multiple slices
auto t = MakeTensorView(data, {2, 3, 4}, kCpuId);
auto t = MakeTensorView(&ctx, data, 2, 3, 4);
auto s_0 = t.Slice(linalg::All(), linalg::Range(0, 2), linalg::Range(1, 4));
ASSERT_FALSE(s_0.CContiguous());
auto s_1 = s_0.Slice(1, 1, linalg::Range(0, 2));
@@ -208,7 +210,7 @@ TEST(Linalg, TensorView) {
TEST(Linalg, Tensor) {
{
Tensor<float, 3> t{{2, 3, 4}, kCpuId};
Tensor<float, 3> t{{2, 3, 4}, kCpuId, Order::kC};
auto view = t.View(kCpuId);
auto const &as_const = t;
@@ -227,7 +229,7 @@ TEST(Linalg, Tensor) {
}
{
// Reshape
Tensor<float, 3> t{{2, 3, 4}, kCpuId};
Tensor<float, 3> t{{2, 3, 4}, kCpuId, Order::kC};
t.Reshape(4, 3, 2);
ASSERT_EQ(t.Size(), 24);
ASSERT_EQ(t.Shape(2), 2);
@@ -245,7 +247,7 @@ TEST(Linalg, Tensor) {
TEST(Linalg, Empty) {
{
auto t = TensorView<double, 2>{{}, {0, 3}, kCpuId};
auto t = TensorView<double, 2>{{}, {0, 3}, kCpuId, Order::kC};
for (int32_t i : {0, 1, 2}) {
auto s = t.Slice(All(), i);
ASSERT_EQ(s.Size(), 0);
@@ -254,7 +256,7 @@ TEST(Linalg, Empty) {
}
}
{
auto t = Tensor<double, 2>{{0, 3}, kCpuId};
auto t = Tensor<double, 2>{{0, 3}, kCpuId, Order::kC};
ASSERT_EQ(t.Size(), 0);
auto view = t.View(kCpuId);
@@ -269,7 +271,7 @@ TEST(Linalg, Empty) {
TEST(Linalg, ArrayInterface) {
auto cpu = kCpuId;
auto t = Tensor<double, 2>{{3, 3}, cpu};
auto t = Tensor<double, 2>{{3, 3}, cpu, Order::kC};
auto v = t.View(cpu);
std::iota(v.Values().begin(), v.Values().end(), 0);
auto arr = Json::Load(StringView{ArrayInterfaceStr(v)});
@@ -313,21 +315,48 @@ TEST(Linalg, Popc) {
}
TEST(Linalg, Stack) {
Tensor<float, 3> l{{2, 3, 4}, kCpuId};
Tensor<float, 3> l{{2, 3, 4}, kCpuId, Order::kC};
ElementWiseTransformHost(l.View(kCpuId), omp_get_max_threads(),
[=](size_t i, float) { return i; });
Tensor<float, 3> r_0{{2, 3, 4}, kCpuId};
Tensor<float, 3> r_0{{2, 3, 4}, kCpuId, Order::kC};
ElementWiseTransformHost(r_0.View(kCpuId), omp_get_max_threads(),
[=](size_t i, float) { return i; });
Stack(&l, r_0);
Tensor<float, 3> r_1{{0, 3, 4}, kCpuId};
Tensor<float, 3> r_1{{0, 3, 4}, kCpuId, Order::kC};
Stack(&l, r_1);
ASSERT_EQ(l.Shape(0), 4);
Stack(&r_1, l);
ASSERT_EQ(r_1.Shape(0), l.Shape(0));
}
} // namespace linalg
} // namespace xgboost
TEST(Linalg, FOrder) {
std::size_t constexpr kRows = 16, kCols = 3;
std::vector<float> data(kRows * kCols);
MatrixView<float> mat{data, {kRows, kCols}, Context::kCpuId, Order::kF};
float k{0};
for (std::size_t i = 0; i < kRows; ++i) {
for (std::size_t j = 0; j < kCols; ++j) {
mat(i, j) = k;
k++;
}
}
auto column = mat.Slice(linalg::All(), 1);
ASSERT_TRUE(column.FContiguous());
ASSERT_EQ(column.Stride(0), 1);
ASSERT_TRUE(column.CContiguous());
k = 1;
for (auto it = linalg::cbegin(column); it != linalg::cend(column); ++it) {
ASSERT_EQ(*it, k);
k += kCols;
}
k = 1;
auto ptr = column.Values().data();
for (auto it = ptr; it != ptr + kRows; ++it) {
ASSERT_EQ(*it, k);
k += kCols;
}
}
} // namespace xgboost::linalg

View File

@@ -7,8 +7,7 @@
#include "xgboost/context.h"
#include "xgboost/linalg.h"
namespace xgboost {
namespace linalg {
namespace xgboost::linalg {
namespace {
void TestElementWiseKernel() {
Tensor<float, 3> l{{2, 3, 4}, 0};
@@ -55,8 +54,10 @@ void TestElementWiseKernel() {
}
void TestSlice() {
Context ctx;
ctx.gpu_id = 1;
thrust::device_vector<double> data(2 * 3 * 4);
auto t = MakeTensorView(dh::ToSpan(data), {2, 3, 4}, 0);
auto t = MakeTensorView(&ctx, dh::ToSpan(data), 2, 3, 4);
dh::LaunchN(1, [=] __device__(size_t) {
auto s = t.Slice(linalg::All(), linalg::Range(0, 3), linalg::Range(0, 4));
auto all = t.Slice(linalg::All(), linalg::All(), linalg::All());
@@ -75,5 +76,4 @@ void TestSlice() {
TEST(Linalg, GPUElementWise) { TestElementWiseKernel(); }
TEST(Linalg, GPUTensorView) { TestSlice(); }
} // namespace linalg
} // namespace xgboost
} // namespace xgboost::linalg