[EM] Initial support for using snappy with sparse data. (#11520)

Author: trivialfis

include/xgboost/data.h

@@ -291,7 +291,6 @@ struct BatchParam {
    * @brief The number of batches to pre-fetch for external memory.
    */
   std::int32_t n_prefetch_batches{3};
-
   /**
    * @brief Exact or others that don't need histogram.
    */
@@ -542,6 +541,10 @@ struct ExtMemConfig {
   float missing;
   // The number of CPU threads.
   std::int32_t n_threads{0};
+  // The ratio of the cache that can be compressed. Used for testing.
+  float hw_decomp_ratio{std::numeric_limits<float>::quiet_NaN()};
+  // Fallback to using nvcomp. Used for testing.
+  bool allow_decomp_fallback{false};
 
   ExtMemConfig() = delete;
   ExtMemConfig(std::string cache, bool on_host, float h_ratio, std::int64_t min_cache,
@@ -552,6 +555,12 @@ struct ExtMemConfig {
         min_cache_page_bytes{min_cache},
         missing{missing},
         n_threads{n_threads} {}
+
+  ExtMemConfig& SetParamsForTest(float _hw_decomp_ratio, bool _allow_decomp_fallback) {
+    this->hw_decomp_ratio = _hw_decomp_ratio;
+    this->allow_decomp_fallback = _allow_decomp_fallback;
+    return *this;
+  }
 };
 
 /**

src/c_api/c_api.cu

@@ -1,11 +1,10 @@
 /**
- * Copyright 2019-2024, XGBoost Contributors
+ * Copyright 2019-2025, XGBoost Contributors
  */
 #include <thrust/transform.h>  // for transform
 
 #include "../common/api_entry.h"       // for XGBAPIThreadLocalEntry
 #include "../common/cuda_context.cuh"  // for CUDAContext
-#include "../common/threading_utils.h"
 #include "../data/array_interface.h"  // for DispatchDType, ArrayInterface
 #include "../data/device_adapter.cuh"
 #include "../data/proxy_dmatrix.h"
@@ -18,6 +17,9 @@
 #if defined(XGBOOST_USE_NCCL)
 #include <nccl.h>
 #endif
+#if defined(XGBOOST_USE_NVCOMP)
+#include <nvcomp/version.h>
+#endif  // defined(XGBOOST_USE_NVCOMP)
 
 namespace xgboost {
 void XGBBuildInfoDevice(Json *p_info) {
@@ -56,6 +58,15 @@ void XGBBuildInfoDevice(Json *p_info) {
 #else
   info["USE_RMM"] = Boolean{false};
 #endif
+
+#if defined(XGBOOST_USE_NVCOMP)
+  info["USE_NVCOMP"] = Boolean{true};
+  v = {Json{Integer{NVCOMP_VER_MAJOR}}, Json{Integer{NVCOMP_VER_MINOR}},
+       Json{Integer{NVCOMP_VER_PATCH}}};
+  info["NVCOMP_VERSION"] = v;
+#else
+  info["USE_NVCOMP"] = Boolean{false};
+#endif
 }
 
 void XGBoostAPIGuard::SetGPUAttribute() {

src/common/cuda_dr_utils.h

@@ -16,7 +16,7 @@
 
 #include "xgboost/string_view.h"  // for StringView
 
-#if CUDART_VERSION >= 12080
+#if CUDART_VERSION >= 12080 && defined(__linux__)
 #define CUDA_HW_DECOM_AVAILABLE 1
 #endif
 

src/common/cuda_pinned_allocator.cu

@@ -63,6 +63,7 @@ namespace xgboost::common::cuda_impl {
       [](cudaMemPool_t* mem_pool) {
         if (mem_pool) {
           dh::safe_cuda(cudaMemPoolDestroy(*mem_pool));
+          delete mem_pool;
         }
       }};
   return mem_pool;

src/common/device_compression.cu

@@ -25,9 +25,11 @@
 #include <mutex>      // for once_flag, call_once
 #include <vector>     // for vector
 
+#include "common.h"               // for HumanMemUnit
 #include "compressed_iterator.h"  // for CompressedByteT
 #include "cuda_context.cuh"       // for CUDAContext
 #include "cuda_dr_utils.h"        // for GetGlobalCuDriverApi
+#include "cuda_rt_utils.h"        // for CurrentDevice
 #include "device_compression.h"
 #include "device_vector.cuh"      // for DeviceUVector
 #include "nvtx_utils.h"           // for xgboost_NVTX_FN_RANGE
@@ -290,7 +292,6 @@ void DecompressSnappy(dh::CUDAStreamView stream, SnappyDecomprMgr const& mgr,
     dh::device_vector<void*> d_out_ptrs(n_chunks);
     dh::safe_cuda(cudaMemcpyAsync(d_out_ptrs.data().get(), h_out_ptrs.data(),
                                   dh::ToSpan(d_out_ptrs).size_bytes(), cudaMemcpyDefault, stream));
-    CHECK(curt::SupportsPageableMem() || curt::SupportsAts());
     // Run nvcomp
     SafeNvComp(nvcompBatchedSnappyDecompressAsync(
         mgr_impl->d_in_chunk_ptrs.data().get(), mgr_impl->d_in_chunk_sizes.data().get(),
@@ -383,8 +384,11 @@ void DecompressSnappy(dh::CUDAStreamView stream, SnappyDecomprMgr const& mgr,
   auto n_bytes = thrust::reduce(cuctx->CTP(), out_sizes.cbegin(), out_sizes.cend());
   auto n_total_bytes = p_out->size();
   auto ratio = static_cast<double>(n_total_bytes) / in.size_bytes();
-  LOG(DEBUG) << "[snappy] Input: " << in.size_bytes() << ", need:" << n_bytes
-             << " allocated:" << n_total_bytes << " ratio:" << ratio;
+  auto ratio_act = static_cast<double>(n_bytes) / in.size_bytes();
+  LOG(DEBUG) << "[snappy] Input: " << common::HumanMemUnit(in.size_bytes())
+             << ", need:" << common::HumanMemUnit(n_bytes)
+             << ", allocated:" << common::HumanMemUnit(n_total_bytes) << ", ratio:" << ratio
+             << ", actual ratio:" << ratio_act;
 
   /**
    * Meta
@@ -470,10 +474,7 @@ SnappyDecomprMgr::~SnappyDecomprMgr() = default;
 SnappyDecomprMgrImpl* SnappyDecomprMgr::Impl() const { return nullptr; }
 
 [[nodiscard]] bool SnappyDecomprMgr::Empty() const { return true; }
-[[nodiscard]] std::size_t SnappyDecomprMgr::DecompressedBytes() const {
-  common::AssertNvCompSupport();
-  return 0;
-}
+[[nodiscard]] std::size_t SnappyDecomprMgr::DecompressedBytes() const { return 0; }
 
 // Round-trip compression
 void DecompressSnappy(dh::CUDAStreamView, SnappyDecomprMgr const&,
@@ -482,15 +483,22 @@ void DecompressSnappy(dh::CUDAStreamView, SnappyDecomprMgr const&,
 }
 
 [[nodiscard]] CuMemParams CompressSnappy(Context const*,
-                                         common::Span<common::CompressedByteT const>,
+                                         common::Span<common::CompressedByteT const> in,
                                          dh::DeviceUVector<std::uint8_t>*, std::size_t) {
+  if (in.empty()) {
+    return {};
+  }
   common::AssertNvCompSupport();
   return {};
 }
 
 [[nodiscard]] common::RefResourceView<std::uint8_t> CoalesceCompressedBuffersToHost(
-    dh::CUDAStreamView, std::shared_ptr<HostPinnedMemPool>, CuMemParams const&,
+    dh::CUDAStreamView, std::shared_ptr<HostPinnedMemPool>, CuMemParams const& in_params,
     dh::DeviceUVector<std::uint8_t> const&, CuMemParams*) {
+  std::size_t n_total_bytes = in_params.TotalSrcBytes();
+  if (n_total_bytes == 0) {
+    return {};
+  }
   common::AssertNvCompSupport();
   return {};
 }

src/common/device_compression.cuh

@@ -104,11 +104,20 @@ struct SnappyDecomprMgrImpl {
   [[nodiscard]] bool Empty() const;
 };
 
-inline auto MakeSnappyDecomprMgr(dh::CUDAStreamView s, std::shared_ptr<HostPinnedMemPool> pool,
-                                 CuMemParams params,
-                                 common::Span<std::uint8_t const> in_compressed_data) {
+#if defined(XGBOOST_USE_NVCOMP)
+[[nodiscard]] inline auto MakeSnappyDecomprMgr(
+    dh::CUDAStreamView s, std::shared_ptr<HostPinnedMemPool> pool, CuMemParams params,
+    common::Span<std::uint8_t const> in_compressed_data) {
   SnappyDecomprMgr mgr;
   *mgr.Impl() = SnappyDecomprMgrImpl{s, std::move(pool), std::move(params), in_compressed_data};
   return mgr;
 }
+#else
+[[nodiscard]] inline auto MakeSnappyDecomprMgr(dh::CUDAStreamView,
+                                               std::shared_ptr<HostPinnedMemPool>, CuMemParams,
+                                               common::Span<std::uint8_t const>) {
+  SnappyDecomprMgr mgr;
+  return mgr;
+}
+#endif  // defined(XGBOOST_USE_NVCOMP)
 }  // namespace xgboost::dc

src/common/io.cc

@@ -42,6 +42,10 @@
 #include <limits>  // for numeric_limits
 #endif
 
+#if defined(__linux__)
+#include <sys/sysinfo.h>
+#endif
+
 namespace xgboost::common {
 size_t PeekableInStream::Read(void* dptr, size_t size) {
   size_t nbuffer = buffer_.length() - buffer_ptr_;
@@ -350,4 +354,19 @@ AlignedMemWriteStream::~AlignedMemWriteStream() = default;
   }
   return result;
 }
+
+[[nodiscard]] std::size_t TotalMemory() {
+#if defined(__linux__)
+  struct sysinfo info;
+  CHECK_EQ(sysinfo(&info), 0) << SystemErrorMsg();
+  return info.totalram * info.mem_unit;
+#elif defined(xgboost_IS_WIN)
+  MEMORYSTATUSEX status;
+  status.dwLength = sizeof(status);
+  CHECK(GlobalMemoryStatusEx(&status)) << SystemErrorMsg();
+  return static_cast<std::size_t>(status.ullTotalPhys);
+#else
+  LOG(FATAL) << "Not implemented";
+#endif  // defined(__linux__)
+}
 }  // namespace xgboost::common

src/common/io.h

@@ -633,5 +633,7 @@ class AlignedMemWriteStream : public AlignedFileWriteStream {
 
 // Run a system command, get its stdout.
 [[nodiscard]] std::string CmdOutput(StringView cmd);
+
+[[nodiscard]] std::size_t TotalMemory();
 }  // namespace xgboost::common
 #endif  // XGBOOST_COMMON_IO_H_

src/data/ellpack_page_raw_format.cu

@@ -102,8 +102,10 @@ template <typename T>
   auto* impl = page->Impl();
   CHECK(this->cuts_->cut_values_.DeviceCanRead());
 
+  auto ctx = Context{}.MakeCUDA(curt::CurrentDevice());
+
   auto dispatch = [&] {
-    fi->Read(page, this->param_.prefetch_copy || !this->has_hmm_ats_);
+    fi->Read(&ctx, page, this->param_.prefetch_copy || !this->has_hmm_ats_);
     impl->SetCuts(this->cuts_);
   };