Make the grow-only resource stream ordered. (#12185)

Author: trivialfis

src/common/device_vector.cu

@@ -1,5 +1,5 @@
 /**
- * Copyright 2017-2024, XGBoost contributors
+ * Copyright 2017-2026, XGBoost contributors
  */
 #include <numeric>  // for accumulate
 
@@ -59,7 +59,7 @@ void ThrowOOMError(std::string const &err, std::size_t bytes) {
   return std::accumulate(it, it + this->handles_.size(), static_cast<std::size_t>(0));
 }
 
-void GrowOnlyVirtualMemVec::Reserve(std::size_t new_size) {
+void GrowOnlyVirtualMemVec::Reserve(std::size_t new_size, xgboost::curt::StreamRef stream) {
   auto va_capacity = this->Capacity();
   if (new_size < va_capacity) {
     return;
@@ -83,6 +83,9 @@ void GrowOnlyVirtualMemVec::Reserve(std::size_t new_size) {
     // New allocation is successful. Map the pyhsical address to the virtual address.
     // First unmap the existing ptr.
     if (this->DevPtr() != 0) {
+      // Slow-path growth invalidates the existing virtual address. Wait for prior users on the
+      // caller's stream before changing the mapping.
+      stream.Sync();
       // Unmap the existing ptr.
       safe_cu(cu_.cuMemUnmap(this->DevPtr(), this->PhyCapacity()));
 

src/common/device_vector.cuh

@@ -204,20 +204,21 @@ class GrowOnlyVirtualMemVec {
     safe_cu(this->cu_.cuMemCreate(&alloc_handle, padded_size, &this->prop_, 0));
     return alloc_handle;
   }
-  void Reserve(std::size_t new_size);
+  void Reserve(std::size_t new_size, xgboost::curt::StreamRef stream);
 
  public:
   explicit GrowOnlyVirtualMemVec(CUmemLocationType type);
 
-  void GrowTo(std::size_t n_bytes) {
+  void GrowTo(std::size_t n_bytes,
+              xgboost::curt::StreamRef stream = xgboost::curt::DefaultStream()) {
     auto alloc_size = this->PhyCapacity();
     if (n_bytes <= alloc_size) {
       return;
     }
 
     std::size_t delta = n_bytes - alloc_size;
     auto const padded_delta = RoundUp(delta, this->granularity_);
-    this->Reserve(alloc_size + padded_delta);
+    this->Reserve(alloc_size + padded_delta, stream);
 
     this->handles_.emplace_back(
         std::unique_ptr<PhyAddrHandle, std::function<void(PhyAddrHandle *)>>{

src/common/resource.cuh

@@ -1,5 +1,5 @@
 /**
- * Copyright 2024-2025, XGBoost Contributors
+ * Copyright 2024-2026, XGBoost Contributors
  */
 #pragma once
 #include <cstddef>     // for size_t
@@ -48,8 +48,8 @@ class CudaGrowOnlyResource : public ResourceHandler {
       : ResourceHandler{kCudaGrowOnly}, alloc_{MakeNew()} {
     this->Resize(n_bytes);
   }
-  void Resize(std::size_t n_bytes) {
-    this->alloc_->GrowTo(n_bytes);
+  void Resize(std::size_t n_bytes, curt::StreamRef stream = curt::DefaultStream()) {
+    this->alloc_->GrowTo(n_bytes, stream);
     this->n_bytes_ = n_bytes;
   }
   void Clear() {

tests/cpp/common/test_ref_resource_view.cu

@@ -1,5 +1,5 @@
 /**
- * Copyright 2024-2025, XGBoost Contributors
+ * Copyright 2024-2026, XGBoost Contributors
  */
 #if defined(__linux__)
 
@@ -25,15 +25,17 @@ class TestCudaGrowOnly : public ::testing::TestWithParam<std::size_t> {
     thrust::sequence(ctx.CUDACtx()->CTP(), ref.begin(), ref.end(), 0.0);
     auto res = std::dynamic_pointer_cast<common::CudaGrowOnlyResource>(ref.Resource());
     CHECK(res);
-    res->Resize(n * sizeof(double));
+    res->Resize(n * sizeof(double), ctx.CUDACtx()->Stream());
 
     auto ref1 = RefResourceView<double>(res->DataAs<double>(), res->Size() / sizeof(double),
                                         ref.Resource());
     ASSERT_EQ(res->Size(), n * sizeof(double));
     ASSERT_EQ(ref1.size(), n);
     thrust::sequence(ctx.CUDACtx()->CTP(), ref1.begin(), ref1.end(), static_cast<double>(0.0));
     std::vector<double> h_vec(ref1.size());
-    dh::safe_cuda(cudaMemcpyAsync(h_vec.data(), ref1.data(), ref1.size_bytes(), cudaMemcpyDefault));
+    dh::safe_cuda(cudaMemcpyAsync(h_vec.data(), ref1.data(), ref1.size_bytes(), cudaMemcpyDefault,
+                                  ctx.CUDACtx()->Stream()));
+    ctx.CUDACtx()->Stream().Sync();
     for (std::size_t i = 0; i < h_vec.size(); ++i) {
       ASSERT_EQ(h_vec[i], i);
     }
@@ -64,8 +66,7 @@ TEST(HostPinnedMemPool, Alloc) {
 
     // Thread safety.
     auto n_threads = static_cast<std::int32_t>(std::thread::hardware_concurrency());
-    common::ThreadPool workers{"tmempool", n_threads, [] {
-                               }};
+    common::ThreadPool workers{"tmempool", n_threads, [] {}};
     std::vector<std::future<RefResourceView<double>>> alloc_futs;
     for (std::int32_t i = 0, n = n_threads * 4; i < n; ++i) {
       auto fut = workers.Submit([i, pool] {