xgboost/tests/cpp/common/test_cuda_host_allocator.cu
Jiaming Yuan 55aef8f546
[EM] Avoid resizing host cache. (#10734)
* [EM] Avoid resizing host cache.

- Add SAM allocator and resource.
- Use page-based cache instead of stream-based cache.
2024-08-23 06:34:01 +08:00

37 lines
951 B
Plaintext

/**
* 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