Use the new DeviceOrd in the linalg module. (#9527)

This commit is contained in:
Jiaming Yuan
2023-08-29 13:37:29 +08:00
committed by GitHub
parent 942b957eef
commit ddf2e68821
43 changed files with 252 additions and 273 deletions

View File

@@ -3,7 +3,7 @@
*/
#include <gtest/gtest.h>
#include <xgboost/context.h>
#include <xgboost/host_device_vector.h>
#include <xgboost/host_device_vector.h> // for HostDeviceVector
#include <xgboost/linalg.h>
#include <cstddef> // size_t
@@ -14,8 +14,8 @@
namespace xgboost::linalg {
namespace {
auto kCpuId = Context::kCpuId;
}
DeviceOrd CPU() { return DeviceOrd::CPU(); }
} // namespace
auto MakeMatrixFromTest(HostDeviceVector<float> *storage, std::size_t n_rows, std::size_t n_cols) {
storage->Resize(n_rows * n_cols);
@@ -23,7 +23,7 @@ auto MakeMatrixFromTest(HostDeviceVector<float> *storage, std::size_t n_rows, st
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)}, CPU()};
return m;
}
@@ -31,7 +31,7 @@ TEST(Linalg, MatrixView) {
size_t kRows = 31, kCols = 77;
HostDeviceVector<float> storage;
auto m = MakeMatrixFromTest(&storage, kRows, kCols);
ASSERT_EQ(m.DeviceIdx(), kCpuId);
ASSERT_EQ(m.Device(), CPU());
ASSERT_EQ(m(0, 0), 0);
ASSERT_EQ(m(kRows - 1, kCols - 1), storage.Size() - 1);
}
@@ -76,7 +76,7 @@ TEST(Linalg, TensorView) {
{
// as vector
TensorView<double, 1> vec{data, {data.size()}, -1};
TensorView<double, 1> vec{data, {data.size()}, CPU()};
ASSERT_EQ(vec.Size(), data.size());
ASSERT_EQ(vec.Shape(0), data.size());
ASSERT_EQ(vec.Shape().size(), 1);
@@ -87,7 +87,7 @@ TEST(Linalg, TensorView) {
{
// as matrix
TensorView<double, 2> mat(data, {6, 4}, -1);
TensorView<double, 2> mat(data, {6, 4}, CPU());
auto s = mat.Slice(2, All());
ASSERT_EQ(s.Shape().size(), 1);
s = mat.Slice(All(), 1);
@@ -96,7 +96,7 @@ TEST(Linalg, TensorView) {
{
// assignment
TensorView<double, 3> t{data, {2, 3, 4}, 0};
TensorView<double, 3> t{data, {2, 3, 4}, CPU()};
double pi = 3.14159;
auto old = t(1, 2, 3);
t(1, 2, 3) = pi;
@@ -201,7 +201,7 @@ TEST(Linalg, TensorView) {
}
{
// f-contiguous
TensorView<double, 3> t{data, {4, 3, 2}, {1, 4, 12}, kCpuId};
TensorView<double, 3> t{data, {4, 3, 2}, {1, 4, 12}, CPU()};
ASSERT_TRUE(t.Contiguous());
ASSERT_TRUE(t.FContiguous());
ASSERT_FALSE(t.CContiguous());
@@ -210,11 +210,11 @@ TEST(Linalg, TensorView) {
TEST(Linalg, Tensor) {
{
Tensor<float, 3> t{{2, 3, 4}, kCpuId, Order::kC};
auto view = t.View(kCpuId);
Tensor<float, 3> t{{2, 3, 4}, CPU(), Order::kC};
auto view = t.View(CPU());
auto const &as_const = t;
auto k_view = as_const.View(kCpuId);
auto k_view = as_const.View(CPU());
size_t n = 2 * 3 * 4;
ASSERT_EQ(t.Size(), n);
@@ -229,7 +229,7 @@ TEST(Linalg, Tensor) {
}
{
// Reshape
Tensor<float, 3> t{{2, 3, 4}, kCpuId, Order::kC};
Tensor<float, 3> t{{2, 3, 4}, CPU(), Order::kC};
t.Reshape(4, 3, 2);
ASSERT_EQ(t.Size(), 24);
ASSERT_EQ(t.Shape(2), 2);
@@ -247,7 +247,7 @@ TEST(Linalg, Tensor) {
TEST(Linalg, Empty) {
{
auto t = TensorView<double, 2>{{}, {0, 3}, kCpuId, Order::kC};
auto t = TensorView<double, 2>{{}, {0, 3}, CPU(), Order::kC};
for (int32_t i : {0, 1, 2}) {
auto s = t.Slice(All(), i);
ASSERT_EQ(s.Size(), 0);
@@ -256,9 +256,9 @@ TEST(Linalg, Empty) {
}
}
{
auto t = Tensor<double, 2>{{0, 3}, kCpuId, Order::kC};
auto t = Tensor<double, 2>{{0, 3}, CPU(), Order::kC};
ASSERT_EQ(t.Size(), 0);
auto view = t.View(kCpuId);
auto view = t.View(CPU());
for (int32_t i : {0, 1, 2}) {
auto s = view.Slice(All(), i);
@@ -270,7 +270,7 @@ TEST(Linalg, Empty) {
}
TEST(Linalg, ArrayInterface) {
auto cpu = kCpuId;
auto cpu = 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);
@@ -315,16 +315,16 @@ TEST(Linalg, Popc) {
}
TEST(Linalg, Stack) {
Tensor<float, 3> l{{2, 3, 4}, kCpuId, Order::kC};
ElementWiseTransformHost(l.View(kCpuId), omp_get_max_threads(),
Tensor<float, 3> l{{2, 3, 4}, CPU(), Order::kC};
ElementWiseTransformHost(l.View(CPU()), omp_get_max_threads(),
[=](size_t i, float) { return i; });
Tensor<float, 3> r_0{{2, 3, 4}, kCpuId, Order::kC};
ElementWiseTransformHost(r_0.View(kCpuId), omp_get_max_threads(),
Tensor<float, 3> r_0{{2, 3, 4}, CPU(), Order::kC};
ElementWiseTransformHost(r_0.View(CPU()), omp_get_max_threads(),
[=](size_t i, float) { return i; });
Stack(&l, r_0);
Tensor<float, 3> r_1{{0, 3, 4}, kCpuId, Order::kC};
Tensor<float, 3> r_1{{0, 3, 4}, CPU(), Order::kC};
Stack(&l, r_1);
ASSERT_EQ(l.Shape(0), 4);
@@ -335,7 +335,7 @@ TEST(Linalg, Stack) {
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};
MatrixView<float> mat{data, {kRows, kCols}, CPU(), Order::kF};
float k{0};
for (std::size_t i = 0; i < kRows; ++i) {
for (std::size_t j = 0; j < kCols; ++j) {

View File

@@ -11,17 +11,18 @@
namespace xgboost::linalg {
namespace {
void TestElementWiseKernel() {
auto device = DeviceOrd::CUDA(0);
Tensor<float, 3> l{{2, 3, 4}, 0};
{
/**
* Non-contiguous
*/
// GPU view
auto t = l.View(0).Slice(linalg::All(), 1, linalg::All());
auto t = l.View(device).Slice(linalg::All(), 1, linalg::All());
ASSERT_FALSE(t.CContiguous());
ElementWiseTransformDevice(t, [] __device__(size_t i, float) { return i; });
// CPU view
t = l.View(Context::kCpuId).Slice(linalg::All(), 1, linalg::All());
t = l.View(DeviceOrd::CPU()).Slice(linalg::All(), 1, linalg::All());
size_t k = 0;
for (size_t i = 0; i < l.Shape(0); ++i) {
for (size_t j = 0; j < l.Shape(2); ++j) {
@@ -29,7 +30,7 @@ void TestElementWiseKernel() {
}
}
t = l.View(0).Slice(linalg::All(), 1, linalg::All());
t = l.View(device).Slice(linalg::All(), 1, linalg::All());
ElementWiseKernelDevice(t, [] XGBOOST_DEVICE(size_t i, float v) { SPAN_CHECK(v == i); });
}
@@ -37,11 +38,11 @@ void TestElementWiseKernel() {
/**
* Contiguous
*/
auto t = l.View(0);
auto t = l.View(device);
ElementWiseTransformDevice(t, [] XGBOOST_DEVICE(size_t i, float) { return i; });
ASSERT_TRUE(t.CContiguous());
// CPU view
t = l.View(Context::kCpuId);
t = l.View(DeviceOrd::CPU());
size_t ind = 0;
for (size_t i = 0; i < l.Shape(0); ++i) {

View File

@@ -41,7 +41,7 @@ void TestCalcQueriesInvIDCG() {
p.UpdateAllowUnknown(Args{{"ndcg_exp_gain", "false"}});
cuda_impl::CalcQueriesInvIDCG(&ctx, linalg::MakeTensorView(&ctx, d_scores, d_scores.size()),
dh::ToSpan(group_ptr), inv_IDCG.View(ctx.gpu_id), p);
dh::ToSpan(group_ptr), inv_IDCG.View(ctx.Device()), p);
for (std::size_t i = 0; i < n_groups; ++i) {
double inv_idcg = inv_IDCG(i);
ASSERT_NEAR(inv_idcg, 0.00551782, kRtEps);

View File

@@ -47,7 +47,7 @@ class StatsGPU : public ::testing::Test {
data.insert(data.cend(), seg.begin(), seg.end());
data.insert(data.cend(), seg.begin(), seg.end());
linalg::Tensor<float, 1> arr{data.cbegin(), data.cend(), {data.size()}, 0};
auto d_arr = arr.View(0);
auto d_arr = arr.View(DeviceOrd::CUDA(0));
auto key_it = dh::MakeTransformIterator<std::size_t>(
thrust::make_counting_iterator(0ul),
@@ -71,8 +71,8 @@ class StatsGPU : public ::testing::Test {
}
void Weighted() {
auto d_arr = arr_.View(0);
auto d_key = indptr_.View(0);
auto d_arr = arr_.View(DeviceOrd::CUDA(0));
auto d_key = indptr_.View(DeviceOrd::CUDA(0));
auto key_it = dh::MakeTransformIterator<std::size_t>(
thrust::make_counting_iterator(0ul),
@@ -81,7 +81,7 @@ class StatsGPU : public ::testing::Test {
dh::MakeTransformIterator<float>(thrust::make_counting_iterator(0ul),
[=] XGBOOST_DEVICE(std::size_t i) { return d_arr(i); });
linalg::Tensor<float, 1> weights{{10}, 0};
linalg::ElementWiseTransformDevice(weights.View(0),
linalg::ElementWiseTransformDevice(weights.View(DeviceOrd::CUDA(0)),
[=] XGBOOST_DEVICE(std::size_t, float) { return 1.0; });
auto w_it = weights.Data()->ConstDevicePointer();
for (auto const& pair : TestSet{{0.0f, 1.0f}, {0.5f, 3.0f}, {1.0f, 5.0f}}) {
@@ -102,7 +102,7 @@ class StatsGPU : public ::testing::Test {
data.insert(data.cend(), seg.begin(), seg.end());
data.insert(data.cend(), seg.begin(), seg.end());
linalg::Tensor<float, 1> arr{data.cbegin(), data.cend(), {data.size()}, 0};
auto d_arr = arr.View(0);
auto d_arr = arr.View(DeviceOrd::CUDA(0));
auto key_it = dh::MakeTransformIterator<std::size_t>(
thrust::make_counting_iterator(0ul),
@@ -125,8 +125,8 @@ class StatsGPU : public ::testing::Test {
}
void NonWeighted() {
auto d_arr = arr_.View(0);
auto d_key = indptr_.View(0);
auto d_arr = arr_.View(DeviceOrd::CUDA(0));
auto d_key = indptr_.View(DeviceOrd::CUDA(0));
auto key_it = dh::MakeTransformIterator<std::size_t>(
thrust::make_counting_iterator(0ul), [=] __device__(std::size_t i) { return d_key(i); });

View File

@@ -22,7 +22,7 @@ TEST(ArrayInterface, Initialize) {
HostDeviceVector<size_t> u64_storage(storage.Size());
std::string u64_arr_str{ArrayInterfaceStr(linalg::TensorView<size_t const, 2>{
u64_storage.ConstHostSpan(), {kRows, kCols}, Context::kCpuId})};
u64_storage.ConstHostSpan(), {kRows, kCols}, DeviceOrd::CPU()})};
std::copy(storage.ConstHostVector().cbegin(), storage.ConstHostVector().cend(),
u64_storage.HostSpan().begin());
auto u64_arr = ArrayInterface<2>{u64_arr_str};

View File

@@ -129,8 +129,8 @@ TEST(MetaInfo, SaveLoadBinary) {
EXPECT_EQ(inforead.group_ptr_, info.group_ptr_);
EXPECT_EQ(inforead.weights_.HostVector(), info.weights_.HostVector());
auto orig_margin = info.base_margin_.View(xgboost::Context::kCpuId);
auto read_margin = inforead.base_margin_.View(xgboost::Context::kCpuId);
auto orig_margin = info.base_margin_.View(xgboost::DeviceOrd::CPU());
auto read_margin = inforead.base_margin_.View(xgboost::DeviceOrd::CPU());
EXPECT_TRUE(std::equal(orig_margin.Values().cbegin(), orig_margin.Values().cend(),
read_margin.Values().cbegin()));
@@ -267,8 +267,8 @@ TEST(MetaInfo, Validate) {
xgboost::HostDeviceVector<xgboost::bst_group_t> d_groups{groups};
d_groups.SetDevice(0);
d_groups.DevicePointer(); // pull to device
std::string arr_interface_str{ArrayInterfaceStr(
xgboost::linalg::MakeVec(d_groups.ConstDevicePointer(), d_groups.Size(), 0))};
std::string arr_interface_str{ArrayInterfaceStr(xgboost::linalg::MakeVec(
d_groups.ConstDevicePointer(), d_groups.Size(), xgboost::DeviceOrd::CUDA(0)))};
EXPECT_THROW(info.SetInfo(ctx, "group", xgboost::StringView{arr_interface_str}), dmlc::Error);
#endif // defined(XGBOOST_USE_CUDA)
}
@@ -307,5 +307,5 @@ TEST(MetaInfo, HostExtend) {
}
namespace xgboost {
TEST(MetaInfo, CPUStridedData) { TestMetaInfoStridedData(Context::kCpuId); }
TEST(MetaInfo, CPUStridedData) { TestMetaInfoStridedData(DeviceOrd::CPU()); }
} // namespace xgboost

View File

@@ -65,7 +65,7 @@ TEST(MetaInfo, FromInterface) {
}
info.SetInfo(ctx, "base_margin", str.c_str());
auto const h_base_margin = info.base_margin_.View(Context::kCpuId);
auto const h_base_margin = info.base_margin_.View(DeviceOrd::CPU());
ASSERT_EQ(h_base_margin.Size(), d_data.size());
for (size_t i = 0; i < d_data.size(); ++i) {
ASSERT_EQ(h_base_margin(i), d_data[i]);
@@ -83,7 +83,7 @@ TEST(MetaInfo, FromInterface) {
}
TEST(MetaInfo, GPUStridedData) {
TestMetaInfoStridedData(0);
TestMetaInfoStridedData(DeviceOrd::CUDA(0));
}
TEST(MetaInfo, Group) {

View File

@@ -14,10 +14,10 @@
#include "../../../src/data/array_interface.h"
namespace xgboost {
inline void TestMetaInfoStridedData(int32_t device) {
inline void TestMetaInfoStridedData(DeviceOrd device) {
MetaInfo info;
Context ctx;
ctx.UpdateAllowUnknown(Args{{"gpu_id", std::to_string(device)}});
ctx.UpdateAllowUnknown(Args{{"device", device.Name()}});
{
// labels
linalg::Tensor<float, 3> labels;
@@ -28,9 +28,9 @@ inline void TestMetaInfoStridedData(int32_t device) {
ASSERT_EQ(t_labels.Shape().size(), 2);
info.SetInfo(ctx, "label", StringView{ArrayInterfaceStr(t_labels)});
auto const& h_result = info.labels.View(-1);
auto const& h_result = info.labels.View(DeviceOrd::CPU());
ASSERT_EQ(h_result.Shape().size(), 2);
auto in_labels = labels.View(-1);
auto in_labels = labels.View(DeviceOrd::CPU());
linalg::ElementWiseKernelHost(h_result, omp_get_max_threads(), [&](size_t i, float& v_0) {
auto tup = linalg::UnravelIndex(i, h_result.Shape());
auto i0 = std::get<0>(tup);
@@ -62,9 +62,9 @@ inline void TestMetaInfoStridedData(int32_t device) {
ASSERT_EQ(t_margin.Shape().size(), 2);
info.SetInfo(ctx, "base_margin", StringView{ArrayInterfaceStr(t_margin)});
auto const& h_result = info.base_margin_.View(-1);
auto const& h_result = info.base_margin_.View(DeviceOrd::CPU());
ASSERT_EQ(h_result.Shape().size(), 2);
auto in_margin = base_margin.View(-1);
auto in_margin = base_margin.View(DeviceOrd::CPU());
linalg::ElementWiseKernelHost(h_result, omp_get_max_threads(), [&](size_t i, float v_0) {
auto tup = linalg::UnravelIndex(i, h_result.Shape());
auto i0 = std::get<0>(tup);

View File

@@ -298,8 +298,8 @@ TEST(SimpleDMatrix, Slice) {
ASSERT_EQ(p_m->Info().weights_.HostVector().at(ridx),
out->Info().weights_.HostVector().at(i));
auto out_margin = out->Info().base_margin_.View(Context::kCpuId);
auto in_margin = margin.View(Context::kCpuId);
auto out_margin = out->Info().base_margin_.View(DeviceOrd::CPU());
auto in_margin = margin.View(DeviceOrd::CPU());
for (size_t j = 0; j < kClasses; ++j) {
ASSERT_EQ(out_margin(i, j), in_margin(ridx, j));
}
@@ -372,8 +372,8 @@ TEST(SimpleDMatrix, SliceCol) {
out->Info().labels_upper_bound_.HostVector().at(i));
ASSERT_EQ(p_m->Info().weights_.HostVector().at(i), out->Info().weights_.HostVector().at(i));
auto out_margin = out->Info().base_margin_.View(Context::kCpuId);
auto in_margin = margin.View(Context::kCpuId);
auto out_margin = out->Info().base_margin_.View(DeviceOrd::CPU());
auto in_margin = margin.View(DeviceOrd::CPU());
for (size_t j = 0; j < kClasses; ++j) {
ASSERT_EQ(out_margin(i, j), in_margin(i, j));
}

View File

@@ -39,9 +39,9 @@ void TestGPUMakePair() {
auto make_args = [&](std::shared_ptr<ltr::RankingCache> p_cache, auto rank_idx,
common::Span<std::size_t const> y_sorted_idx) {
linalg::Vector<double> dummy;
auto d = dummy.View(ctx.gpu_id);
auto d = dummy.View(ctx.Device());
linalg::Vector<GradientPair> dgpair;
auto dg = dgpair.View(ctx.gpu_id);
auto dg = dgpair.View(ctx.Device());
cuda_impl::KernelInputs args{
d,
d,
@@ -50,9 +50,9 @@ void TestGPUMakePair() {
p_cache->DataGroupPtr(&ctx),
p_cache->CUDAThreadsGroupPtr(),
rank_idx,
info.labels.View(ctx.gpu_id),
info.labels.View(ctx.Device()),
predt.ConstDeviceSpan(),
linalg::MatrixView<GradientPair>{common::Span<GradientPair>{}, {0}, 0},
linalg::MatrixView<GradientPair>{common::Span<GradientPair>{}, {0}, DeviceOrd::CUDA(0)},
dg,
nullptr,
y_sorted_idx,

View File

@@ -226,7 +226,7 @@ TEST(GPUPredictor, ShapStump) {
auto dmat = RandomDataGenerator(3, 1, 0).GenerateDMatrix();
gpu_predictor->PredictContribution(dmat.get(), &predictions, model);
auto& phis = predictions.HostVector();
auto base_score = mparam.BaseScore(Context::kCpuId)(0);
auto base_score = mparam.BaseScore(DeviceOrd::CPU())(0);
EXPECT_EQ(phis[0], 0.0);
EXPECT_EQ(phis[1], base_score);
EXPECT_EQ(phis[2], 0.0);

View File

@@ -287,7 +287,7 @@ void TestCategoricalPrediction(Context const* ctx, bool is_column_split) {
predictor->InitOutPredictions(m->Info(), &out_predictions.predictions, model);
predictor->PredictBatch(m.get(), &out_predictions, model, 0);
auto score = mparam.BaseScore(Context::kCpuId)(0);
auto score = mparam.BaseScore(DeviceOrd::CPU())(0);
ASSERT_EQ(out_predictions.predictions.Size(), 1ul);
ASSERT_EQ(out_predictions.predictions.HostVector()[0],
right_weight + score); // go to right for matching cat