1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
|
/**
* Copyright 2024, XGBoost Contributors
*/
#include <gtest/gtest.h>
#include <xgboost/context.h> // for Context
#include <vector>
#include "../../../src/common/cuda_pinned_allocator.h"
#include "../../../src/common/device_helpers.cuh" // for DefaultStream
#include "../../../src/common/numeric.h" // for Iota
namespace xgboost {
TEST(CudaHostMalloc, Pinned) {
std::vector<float, common::cuda_impl::PinnedAllocator<float>> vec;
vec.resize(10);
ASSERT_EQ(vec.size(), 10);
Context ctx;
common::Iota(&ctx, vec.begin(), vec.end(), 0);
float k = 0;
for (auto v : vec) {
ASSERT_EQ(v, k);
++k;
}
}
TEST(CudaHostMalloc, Managed) {
std::vector<float, common::cuda_impl::ManagedAllocator<float>> vec;
vec.resize(10);
#if defined(__linux__)
dh::safe_cuda(
cudaMemPrefetchAsync(vec.data(), vec.size() * sizeof(float), 0, dh::DefaultStream()));
#endif
dh::DefaultStream().Sync();
}
} // namespace xgboost
|
