Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

11
tests/cpp/common/test_linalg.cu
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) {