[EM] Small fixes for hardware decompression. (#11512)

trivialfis · web-flow · commit f1871a40b8f0 · 2025-06-17T02:15:59.000+08:00
- Use a memory pool for result allocation.
- Copy the parameters.
diff --git a/src/common/cuda_rt_utils.cc b/src/common/cuda_rt_utils.cc
@@ -42,13 +42,13 @@ std::int32_t CurrentDevice(bool raise) {
 }
 
 // alternatively: `nvidia-smi -q | grep Addressing`
-bool SupportsPageableMem() {
+[[nodiscard]] bool SupportsPageableMem() {
   std::int32_t res{0};
   dh::safe_cuda(cudaDeviceGetAttribute(&res, cudaDevAttrPageableMemoryAccess, CurrentDevice()));
   return res == 1;
 }
 
-bool SupportsAts() {
+[[nodiscard]] bool SupportsAts() {
   std::int32_t res{0};
   dh::safe_cuda(cudaDeviceGetAttribute(&res, cudaDevAttrPageableMemoryAccessUsesHostPageTables,
                                        CurrentDevice()));
@@ -106,7 +106,7 @@ void DrVersion(std::int32_t* major, std::int32_t* minor) {
 
 [[nodiscard]] std::int32_t GetNumaId() {
   std::int32_t numa_id = -1;
-  dh::safe_cuda(cudaDeviceGetAttribute(&numa_id, cudaDevAttrNumaId, curt::CurrentDevice()));
+  dh::safe_cuda(cudaDeviceGetAttribute(&numa_id, cudaDevAttrHostNumaId, curt::CurrentDevice()));
   numa_id = std::max(numa_id, 0);
   return numa_id;
 }
diff --git a/src/common/cuda_rt_utils.h b/src/common/cuda_rt_utils.h
@@ -15,10 +15,10 @@ std::int32_t CurrentDevice(bool raise = true);
 
 // Whether the device supports coherently accessing pageable memory without calling
 // `cudaHostRegister` on it
-bool SupportsPageableMem();
+[[nodiscard]] bool SupportsPageableMem();
 
 // Address Translation Service (ATS)
-bool SupportsAts();
+[[nodiscard]] bool SupportsAts();
 
 void CheckComputeCapability();
 
diff --git a/src/common/device_compression.cu b/src/common/device_compression.cu
@@ -25,10 +25,9 @@
 #include <mutex>      // for once_flag, call_once
 #include <vector>     // for vector
 
-#include "compressed_iterator.h"    // for CompressedByteT
-#include "cuda_context.cuh"         // for CUDAContext
-#include "cuda_dr_utils.h"          // for GetGlobalCuDriverApi
-#include "cuda_pinned_allocator.h"  // for HostPinnedMemPool
+#include "compressed_iterator.h"  // for CompressedByteT
+#include "cuda_context.cuh"       // for CUDAContext
+#include "cuda_dr_utils.h"        // for GetGlobalCuDriverApi
 #include "device_compression.h"
 #include "device_vector.cuh"      // for DeviceUVector
 #include "nvtx_utils.h"           // for xgboost_NVTX_FN_RANGE
@@ -154,9 +153,11 @@ void SafeNvComp(nvcompStatus_t status) {
   return de;
 }
 
-SnappyDecomprMgrImpl::SnappyDecomprMgrImpl(
-    dh::CUDAStreamView s, std::shared_ptr<common::cuda_impl::HostPinnedMemPool> pool,
-    CuMemParams params, common::Span<std::uint8_t const> in_compressed_data) {
+SnappyDecomprMgrImpl::SnappyDecomprMgrImpl(dh::CUDAStreamView s,
+                                           std::shared_ptr<HostPinnedMemPool> pool,
+                                           CuMemParams params,
+                                           common::Span<std::uint8_t const> in_compressed_data)
+    : n_dst_bytes{params.TotalDstBytes()} {
   std::size_t n_chunks = params.size();
   if (n_chunks == 0) {
     return;
@@ -177,6 +178,7 @@ SnappyDecomprMgrImpl::SnappyDecomprMgrImpl(
     last_in += params[i].src_nbytes;
     last_out += params[i].dst_nbytes;
   }
+  CHECK_EQ(this->n_dst_bytes, last_out);
 
   // copy to d
   dh::CopyTo(in_chunk_ptrs, &this->d_in_chunk_ptrs, s);
@@ -210,6 +212,7 @@ SnappyDecomprMgrImpl::SnappyDecomprMgrImpl(
 
 common::Span<CUmemDecompressParams> SnappyDecomprMgrImpl::GetParams(
     common::Span<common::CompressedByteT> out) {
+  xgboost_NVTX_FN_RANGE_C(3, 252, 198);
   if (this->de_params.empty()) {
     return {};
   }
@@ -218,19 +221,34 @@ common::Span<CUmemDecompressParams> SnappyDecomprMgrImpl::GetParams(
   // Set the output buffers.
   std::size_t last_out = 0;
   for (std::size_t i = 0; i < n_chunks; ++i) {
+    this->de_params_copy[i] = this->de_params[i];
     this->de_params_copy[i].dst = out.subspan(last_out, de_params[i].dstNumBytes).data();
     last_out += de_params[i].dstNumBytes;
   }
 
   return this->de_params_copy.ToSpan();
 }
 
+[[nodiscard]] bool SnappyDecomprMgrImpl::Empty() const {
+#if defined(CUDA_HW_DECOM_AVAILABLE)
+  return this->de_params.empty();
+#else
+  return true;
+#endif
+}
+
 SnappyDecomprMgr::SnappyDecomprMgr() : pimpl_{std::make_unique<SnappyDecomprMgrImpl>()} {}
 SnappyDecomprMgr::SnappyDecomprMgr(SnappyDecomprMgr&& that) = default;
 SnappyDecomprMgr& SnappyDecomprMgr::operator=(SnappyDecomprMgr&& that) = default;
 
 SnappyDecomprMgr::~SnappyDecomprMgr() = default;
 
+[[nodiscard]] bool SnappyDecomprMgr::Empty() const { return this->Impl()->Empty(); }
+
+[[nodiscard]] std::size_t SnappyDecomprMgr::DecompressedBytes() const {
+  return this->Impl()->n_dst_bytes;
+}
+
 SnappyDecomprMgrImpl* SnappyDecomprMgr::Impl() const { return this->pimpl_.get(); }
 
 void DecompressSnappy(dh::CUDAStreamView stream, SnappyDecomprMgr const& mgr,
@@ -272,7 +290,7 @@ 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(),
@@ -387,8 +405,9 @@ void DecompressSnappy(dh::CUDAStreamView stream, SnappyDecomprMgr const& mgr,
 }
 
 [[nodiscard]] common::RefResourceView<std::uint8_t> CoalesceCompressedBuffersToHost(
-    dh::CUDAStreamView stream, CuMemParams const& in_params,
-    dh::DeviceUVector<std::uint8_t> const& in_buf, CuMemParams* p_out) {
+    dh::CUDAStreamView stream, std::shared_ptr<HostPinnedMemPool> pool,
+    CuMemParams const& in_params, dh::DeviceUVector<std::uint8_t> const& in_buf,
+    CuMemParams* p_out) {
   std::size_t n_total_act_bytes = in_params.TotalSrcActBytes();
   std::size_t n_total_bytes = in_params.TotalSrcBytes();
   if (n_total_bytes == 0) {
@@ -399,8 +418,8 @@ void DecompressSnappy(dh::CUDAStreamView stream, SnappyDecomprMgr const& mgr,
   // copy from device buffer to the host cache.
   CHECK_EQ(n_total_bytes, in_buf.size());
   auto c_page =
-      common::MakeFixedVecWithPinnedMalloc<std::remove_reference_t<decltype(in_buf)>::value_type>(
-          n_total_act_bytes);
+      common::MakeFixedVecWithPinnedMemPool<std::remove_reference_t<decltype(in_buf)>::value_type>(
+          pool, n_total_act_bytes, stream);
   std::vector<std::uint8_t const*> srcs(in_params.size());
   std::vector<std::uint8_t*> dsts(in_params.size());
   std::vector<std::size_t> sizes(in_params.size());
@@ -450,6 +469,12 @@ SnappyDecomprMgr& SnappyDecomprMgr::operator=(SnappyDecomprMgr&& that) = default
 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;
+}
+
 // Round-trip compression
 void DecompressSnappy(dh::CUDAStreamView, SnappyDecomprMgr const&,
                       common::Span<common::CompressedByteT>, bool) {
@@ -464,7 +489,8 @@ void DecompressSnappy(dh::CUDAStreamView, SnappyDecomprMgr const&,
 }
 
 [[nodiscard]] common::RefResourceView<std::uint8_t> CoalesceCompressedBuffersToHost(
-    dh::CUDAStreamView, CuMemParams const&, dh::DeviceUVector<std::uint8_t> const&, CuMemParams*) {
+    dh::CUDAStreamView, std::shared_ptr<HostPinnedMemPool>, CuMemParams const&,
+    dh::DeviceUVector<std::uint8_t> const&, CuMemParams*) {
   common::AssertNvCompSupport();
   return {};
 }
diff --git a/src/common/device_compression.cuh b/src/common/device_compression.cuh
@@ -15,6 +15,9 @@
 #include "xgboost/span.h"           // for Span
 
 namespace xgboost::dc {
+
+using HostPinnedMemPool = common::cuda_impl::HostPinnedMemPool;
+
 /**
  * @brief Use nvcomp to compress the data.
  *
@@ -44,19 +47,22 @@ void DecompressSnappy(dh::CUDAStreamView stream, SnappyDecomprMgr const& mgr,
  * @brief Coalesce the compressed chunks into a contiguous host pinned buffer.
  *
  * @param stream CUDA stream.
+ * @param pool Pinned memory pool for storing the results.
  * @param in_params Params from @ref CompressSnappy, specifies the chunks.
  * @param in_buf The buffer storing compressed chunks.
  * @param p_out Re-newed parameters to keep track of the buffers.
  */
 [[nodiscard]] common::RefResourceView<std::uint8_t> CoalesceCompressedBuffersToHost(
-    dh::CUDAStreamView stream, CuMemParams const& in_params,
-    dh::DeviceUVector<std::uint8_t> const& in_buf, CuMemParams* p_out);
+    dh::CUDAStreamView stream, std::shared_ptr<HostPinnedMemPool> pool,
+    CuMemParams const& in_params, dh::DeviceUVector<std::uint8_t> const& in_buf,
+    CuMemParams* p_out);
 
 // We store decompression parameters in struct of vectors. This is due to nvcomp works
 // with this format. But the CUDA driver works with vector of structs. We can optimize
 // toward the driver decompression function if the overhead is significant (too many
 // chunks).
 struct SnappyDecomprMgrImpl {
+  std::size_t n_dst_bytes{0};
   // src of the CUmemDecompressParams
   dh::device_vector<void const*> d_in_chunk_ptrs;
   // srcNumBytes of the CUmemDecompressParams
@@ -81,8 +87,7 @@ struct SnappyDecomprMgrImpl {
 #endif  // defined(CUDA_HW_DECOM_AVAILABLE)
   }
 
-  SnappyDecomprMgrImpl(dh::CUDAStreamView s,
-                       std::shared_ptr<common::cuda_impl::HostPinnedMemPool> pool,
+  SnappyDecomprMgrImpl(dh::CUDAStreamView s, std::shared_ptr<HostPinnedMemPool> pool,
                        CuMemParams params, common::Span<std::uint8_t const> in_compressed_data);
 
 #if defined(CUDA_HW_DECOM_AVAILABLE) && defined(XGBOOST_USE_NVCOMP)
@@ -95,10 +100,11 @@ struct SnappyDecomprMgrImpl {
   SnappyDecomprMgrImpl(SnappyDecomprMgrImpl&& that) = default;
   SnappyDecomprMgrImpl& operator=(SnappyDecomprMgrImpl const&) = delete;
   SnappyDecomprMgrImpl& operator=(SnappyDecomprMgrImpl&&) = default;
+
+  [[nodiscard]] bool Empty() const;
 };
 
-inline auto MakeSnappyDecomprMgr(dh::CUDAStreamView s,
-                                 std::shared_ptr<common::cuda_impl::HostPinnedMemPool> pool,
+inline auto MakeSnappyDecomprMgr(dh::CUDAStreamView s, std::shared_ptr<HostPinnedMemPool> pool,
                                  CuMemParams params,
                                  common::Span<std::uint8_t const> in_compressed_data) {
   SnappyDecomprMgr mgr;
diff --git a/src/common/device_compression.h b/src/common/device_compression.h
@@ -112,6 +112,12 @@ class SnappyDecomprMgr {
 
   SnappyDecomprMgrImpl* Impl() const;
 
+  [[nodiscard]] bool Empty() const;
+  /**
+   * @brief The number of bytes of the uncompressed data.
+   */
+  [[nodiscard]] std::size_t DecompressedBytes() const;
+
  private:
   // Hide the CUDA API calls.
 #if defined(XGBOOST_USE_NVCOMP)
diff --git a/src/data/sparse_page_source.h b/src/data/sparse_page_source.h
@@ -362,7 +362,7 @@ class SparsePageSourceImpl : public BatchIteratorImpl<S>, public FormatStreamPol
 
     timer.Stop();
     if (bytes != InvalidPageSize()) {
-      // Not entirely accurate, the kernels doesn't have to flush the data.
+      // Not entirely accurate, the kernel doesn't have to flush the data.
       LOG(INFO) << common::HumanMemUnit(bytes) << " written in " << timer.ElapsedSeconds()
                 << " seconds.";
       cache_info_->Push(bytes);
diff --git a/tests/cpp/common/test_device_compression.cu b/tests/cpp/common/test_device_compression.cu
@@ -71,11 +71,12 @@ class TestNvComp : public ::testing::TestWithParam<std::tuple<std::size_t, std::
       ASSERT_GE(params.size(), n_bytes / n_chunk_bytes);
     }
 
+    auto pool = std::make_shared<common::cuda_impl::HostPinnedMemPool>();
+
     CuMemParams out_params;
-    auto page = CoalesceCompressedBuffersToHost(cuctx->Stream(), params, compr, &out_params);
+    auto page = CoalesceCompressedBuffersToHost(cuctx->Stream(), pool, params, compr, &out_params);
 
     dh::device_vector<common::CompressedByteT> dout(in.size(), 0);
-    auto pool = std::make_shared<common::cuda_impl::HostPinnedMemPool>();
     auto mgr = MakeSnappyDecomprMgr(cuctx->Stream(), pool, out_params, page.ToSpan());
     DecompressSnappy(cuctx->Stream(), mgr, dh::ToSpan(dout), true);
 
