[EM] Multi-level quantile sketching for GPU. (#10813)

This commit is contained in:
Jiaming Yuan
2024-09-10 13:08:34 +08:00
committed by GitHub
parent 3ef8383d93
commit ed5f33df16
6 changed files with 111 additions and 34 deletions

View File

@@ -17,6 +17,7 @@
#include "../../../include/xgboost/logging.h"
#include "../../../src/common/cuda_context.cuh"
#include "../../../src/common/cuda_rt_utils.h" // for SetDevice
#include "../../../src/common/device_helpers.cuh"
#include "../../../src/common/hist_util.cuh"
#include "../../../src/common/hist_util.h"
@@ -59,8 +60,7 @@ TEST(HistUtil, SketchBatchNumElements) {
GTEST_SKIP_("Test not runnable with RMM enabled.");
#endif // defined(XGBOOST_USE_RMM) && XGBOOST_USE_RMM == 1
size_t constexpr kCols = 10000;
int device;
dh::safe_cuda(cudaGetDevice(&device));
std::int32_t device = dh::CurrentDevice();
auto avail = static_cast<size_t>(dh::AvailableMemory(device) * 0.8);
auto per_elem = detail::BytesPerElement(false);
auto avail_elem = avail / per_elem;
@@ -576,7 +576,7 @@ TEST(HistUtil, AdapterDeviceSketchBatches) {
namespace {
auto MakeData(Context const* ctx, std::size_t n_samples, bst_feature_t n_features) {
dh::safe_cuda(cudaSetDevice(ctx->Ordinal()));
common::SetDevice(ctx->Ordinal());
auto n = n_samples * n_features;
std::vector<float> x;
x.resize(n);

View File

@@ -13,7 +13,7 @@
namespace xgboost::data {
void TestEquivalent(float sparsity) {
Context ctx{MakeCUDACtx(0)};
auto ctx = MakeCUDACtx(0);
CudaArrayIterForTest iter{sparsity};
IterativeDMatrix m(&iter, iter.Proxy(), nullptr, Reset, Next,