Merge branch 'main' into fix/gb300-arena-test-timeout

Author: bdice

conda/environments/all_cuda-129_arch-aarch64.yaml

@@ -35,7 +35,7 @@ dependencies:
 - pytest-cov
 - python>=3.11
 - rapids-build-backend>=0.4.0,<0.5.0
-- rapids-logger==0.2.*,>=0.0.0a0
+- rapids-logger==0.3.*
 - scikit-build-core>=0.11.0
 - sphinx
 - sphinx-copybutton

conda/environments/all_cuda-129_arch-x86_64.yaml

@@ -35,7 +35,7 @@ dependencies:
 - pytest-cov
 - python>=3.11
 - rapids-build-backend>=0.4.0,<0.5.0
-- rapids-logger==0.2.*,>=0.0.0a0
+- rapids-logger==0.3.*
 - scikit-build-core>=0.11.0
 - sphinx
 - sphinx-copybutton

conda/environments/all_cuda-133_arch-aarch64.yaml

@@ -35,7 +35,7 @@ dependencies:
 - pytest-cov
 - python>=3.11
 - rapids-build-backend>=0.4.0,<0.5.0
-- rapids-logger==0.2.*,>=0.0.0a0
+- rapids-logger==0.3.*
 - scikit-build-core>=0.11.0
 - sphinx
 - sphinx-copybutton

conda/environments/all_cuda-133_arch-x86_64.yaml

@@ -35,7 +35,7 @@ dependencies:
 - pytest-cov
 - python>=3.11
 - rapids-build-backend>=0.4.0,<0.5.0
-- rapids-logger==0.2.*,>=0.0.0a0
+- rapids-logger==0.3.*
 - scikit-build-core>=0.11.0
 - sphinx
 - sphinx-copybutton

conda/recipes/librmm/recipe.yaml

@@ -73,7 +73,7 @@ cache:
       - cuda-version =${{ cuda_version }}
       - ${{ stdlib("c") }}
     host:
-      - rapids-logger =0.2
+      - rapids-logger =0.3
       - cuda-driver-dev
 
 outputs:
@@ -97,10 +97,10 @@ outputs:
       host:
         - cuda-version =${{ cuda_version }}
         - cuda-cudart-dev
-        - rapids-logger =0.2
+        - rapids-logger =0.3
       run:
         - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
-        - rapids-logger =0.2
+        - rapids-logger =0.3
       run_exports:
         - ${{ pin_subpackage("librmm", upper_bound="x.x") }}
       ignore_run_exports:
@@ -138,7 +138,7 @@ outputs:
       run:
         - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
         - ${{ pin_subpackage("librmm", exact=True) }}
-        - rapids-logger =0.2
+        - rapids-logger =0.3
         - cuda-cudart
       ignore_run_exports:
         from_package:

cpp/include/rmm/mr/detail/pool_memory_resource_impl.hpp

@@ -70,6 +70,7 @@ class pool_memory_resource_impl final
   block_type try_to_expand(std::size_t try_size, std::size_t min_size, cuda_stream_view stream);
   void initialize_pool(std::size_t initial_size, std::optional<std::size_t> maximum_size);
   block_type expand_pool(std::size_t size, free_list& blocks, cuda_stream_view stream);
+  void reclaim_free_blocks(std::size_t size, free_list& blocks, cuda_stream_view stream);
   [[nodiscard]] std::size_t size_to_grow(std::size_t size) const;
   block_type block_from_upstream(std::size_t size, cuda_stream_view stream);
   split_block allocate_from_block(block_type const& block, std::size_t size);
@@ -85,9 +86,9 @@ class pool_memory_resource_impl final
   cuda::mr::any_resource<cuda::mr::device_accessible> upstream_mr_;
   std::size_t current_pool_size_{};
   std::optional<std::size_t> maximum_pool_size_{};
-  std::set<block_type, compare_blocks<block_type>> upstream_blocks_;
+  std::set<block_type, compare_blocks<block_type>> upstream_blocks_;  ///< Upstream allocations.
 #ifdef RMM_POOL_TRACK_ALLOCATIONS
-  std::set<block_type, compare_blocks<block_type>> allocated_blocks_;
+  std::set<block_type, compare_blocks<block_type>> allocated_blocks_;  ///< Live suballocations.
 #endif
 };
 

cpp/src/mr/detail/pool_memory_resource_impl.cpp

@@ -106,9 +106,66 @@ void pool_memory_resource_impl::initialize_pool(std::size_t initial_size,
 }
 
 pool_memory_resource_impl::block_type pool_memory_resource_impl::expand_pool(
-  std::size_t size, [[maybe_unused]] free_list& blocks, cuda_stream_view stream)
+  std::size_t size, free_list& blocks, cuda_stream_view stream)
 {
-  return try_to_expand(size_to_grow(size), size, stream);
+  auto grow_size = size_to_grow(size);
+  // When the pool is capped and cannot grow enough to satisfy `size` in a single new upstream
+  // block, try to reclaim entirely-free upstream blocks (whose budget prevents growth) back to
+  // upstream, freeing headroom under `maximum_pool_size_` to grow a sufficiently large block.
+  if (grow_size < size && maximum_pool_size_.has_value()) {
+    reclaim_free_blocks(size, blocks, stream);
+    grow_size = size_to_grow(size);
+  }
+  return try_to_expand(grow_size, size, stream);
+}
+
+void pool_memory_resource_impl::reclaim_free_blocks(std::size_t size,
+                                                    free_list& blocks,
+                                                    cuda_stream_view stream)
+{
+  // Reclamation only frees headroom when the pool has a capped maximum size.
+  if (!maximum_pool_size_.has_value()) { return; }
+  auto const max_pool_size = maximum_pool_size_.value();
+
+  // If `size` can never fit under the cap even with every free block reclaimed, there is nothing to
+  // gain: avoid the synchronize and the destructive reclaim on a request that will fail anyway.
+  if (size > max_pool_size) { return; }
+
+  auto free_iter     = blocks.cbegin();
+  auto upstream_iter = upstream_blocks_.cbegin();
+  using compare_t    = decltype(upstream_blocks_)::key_compare;
+  auto const compare = compare_t{};
+
+  // This merge join requires `blocks` and `upstream_blocks_` to remain sorted by `compare_blocks`.
+  // coalescing_free_list maintains that order on insertion, and upstream_blocks_ uses the same
+  // comparator. Erasing matched entries below preserves the ordering of both collections.
+  while (free_iter != blocks.cend() && upstream_iter != upstream_blocks_.cend()) {
+    // `current_pool_size_ <= max_pool_size` is an invariant, so the subtraction cannot underflow.
+    if (max_pool_size - current_pool_size_ >= size) { return; }
+
+    if (compare(*free_iter, *upstream_iter)) {
+      ++free_iter;
+      continue;
+    }
+    if (compare(*upstream_iter, *free_iter)) {
+      ++upstream_iter;
+      continue;
+    }
+
+    auto const candidate = free_iter++;
+    auto const upstream  = upstream_iter++;
+    if (!candidate->is_head() || upstream->size() != candidate->size()) { continue; }
+
+    auto const blk = *candidate;
+    // The free lists were merged onto `stream`, which waits on their recorded events. Enqueueing
+    // the upstream deallocation on the same stream preserves those dependencies without blocking
+    // the host.
+    get_upstream_resource().deallocate(
+      stream, blk.pointer(), blk.size(), rmm::CUDA_ALLOCATION_ALIGNMENT);
+    blocks.erase(candidate);
+    upstream_blocks_.erase(upstream);
+    current_pool_size_ -= blk.size();
+  }
 }
 
 std::size_t pool_memory_resource_impl::size_to_grow(std::size_t size) const
@@ -152,6 +209,8 @@ pool_memory_resource_impl::block_type pool_memory_resource_impl::free_block(
   void* ptr, std::size_t size) noexcept
 {
 #ifdef RMM_POOL_TRACK_ALLOCATIONS
+  // Fetch the metadata recorded for this block's suballocation and validate
+  // the caller's provided size before returning the block to a free list.
   if (ptr == nullptr) return block_type{};
   auto const iter = allocated_blocks_.find(static_cast<char*>(ptr));
   RMM_LOGGING_ASSERT(iter != allocated_blocks_.end());
@@ -162,6 +221,9 @@ pool_memory_resource_impl::block_type pool_memory_resource_impl::free_block(
 
   return block;
 #else
+  // Reconstruct the block, trusting the validity of the caller's pointer and
+  // size. A pointer is a block head if and only if it is the start of an
+  // upstream allocation.
   auto const iter = upstream_blocks_.find(static_cast<char*>(ptr));
   return block_type{static_cast<char*>(ptr), size, (iter != upstream_blocks_.end())};
 #endif