Add indexed stream-ordered pooled resource - #2507
Conversation
Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis change adds indexed free lists, a stream-ordered indexed memory resource, and a public indexed pool resource. It includes upstream pool growth and reclamation, cross-stream recovery, extensive tests, build integration, and benchmarks for allocation and fragmentation behavior. ChangesIndexed pool resource
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (14)
cpp/benchmarks/pool_fragmentation/pool_fragmentation_bench.cpp (1)
29-32: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse SCREAMING_SNAKE_CASE for constants.
Rename the new
constexprconstants. For example, renamealignmenttoALIGNMENTandfirst_lookup_iterationstoFIRST_LOOKUP_ITERATIONS.As per coding guidelines, “C++ constants and macros should use SCREAMING_SNAKE_CASE naming convention.”
Also applies to: 35-42, 48-50, 102-104, 138-140, 151-151, 176-176, 218-221, 282-282
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/benchmarks/pool_fragmentation/pool_fragmentation_bench.cpp` around lines 29 - 32, Rename all newly introduced C++ constexpr constants in fragment_size and the other referenced benchmark sections to SCREAMING_SNAKE_CASE, including alignment to ALIGNMENT and first_lookup_iterations to FIRST_LOOKUP_ITERATIONS; update every use consistently without changing behavior.Source: Coding guidelines
cpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hpp (2)
80-95: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse the existing
crtphelper and fix the broken link.
indexed_crtpduplicatescrtpfrom the existing stream-ordered resource header, and the rename corrupted the reference URL.https://www.fluentcpp.com/2017/05/19/indexed_crtp-helper/does not exist; the article iscrtp-helper. Include the existing helper instead of copying it, or at minimum correct the URL.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hpp` around lines 80 - 95, Remove the duplicated indexed_crtp definition and reuse the existing crtp helper from the stream-ordered resource header, updating references to the existing symbol as needed. If the helper remains locally, correct its documentation URL to https://www.fluentcpp.com/2017/05/19/crtp-helper/.
613-626: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winReplace the donor-lookup asserts with a runtime fallback.
Lines 617 and 621 assume the shared maximum index agrees with the donor free list. With
NDEBUGa disagreement is not caught:selection.blockequalsowner_blocks.end(), andallocate_from_selectiondereferences it at line 463. The donor's own cached failed lookup makes this state reachable if either index ever drifts. Skip the candidate instead.🛡️ Proposed fallback
- auto const owner = candidate->second; - auto owner_iter = stream_free_blocks_.find(owner); - assert(owner_iter != stream_free_blocks_.end()); - - auto& owner_blocks = owner_iter->second; - auto const selection = owner_blocks.find_block(size); - assert(selection.block != owner_blocks.end()); + auto const owner = candidate->second; + auto owner_iter = stream_free_blocks_.find(owner); + if (owner_iter == stream_free_blocks_.end()) { return {}; } + + auto& owner_blocks = owner_iter->second; + auto const selection = owner_blocks.find_block(size); + if (selection.block == owner_blocks.end()) { return {}; }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hpp` around lines 613 - 626, Replace the assertions in the candidate-selection path around stream_free_blocks_ and owner_blocks.find_block with runtime validation: if the donor lookup is missing or selection.block equals owner_blocks.end(), skip this candidate and continue searching. Ensure allocate_from_selection is called only with a valid owner iterator and block selection, while preserving the existing empty-result behavior when no candidates remain.cpp/include/rmm/mr/detail/indexed_coalescing_free_list.hpp (3)
334-347: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd Doxygen comments to the remaining public members.
diagnostics_index_active,diagnostics_indexes_consistent, andclearhave no Doxygen block.insert(block_type const&)at line 172 documents@parambut not itsstd::size_treturn value, which callers use as the inserted (possibly coalesced) block size for the maximum index.📝 Proposed documentation
+ /** + * `@brief` Reports whether the size and address indexes are active. + * + * `@return` true if both indexes are maintained for this free list. + */ [[nodiscard]] bool diagnostics_index_active() const noexcept { return index_active_; } + /** + * `@brief` Checks that both indexes agree with the block list. + * + * `@return` true if the indexes are consistent with the list contents. + */ [[nodiscard]] bool diagnostics_indexes_consistent() const noexcept+ /** + * `@brief` Erases all blocks, both indexes, and the cached failed lookup. + */ void clear() noexceptAs per coding guidelines: "All public functions in rmm:: must have Doxygen documentation".
Also applies to: 365-372
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/rmm/mr/detail/indexed_coalescing_free_list.hpp` around lines 334 - 347, Add Doxygen documentation to the public members diagnostics_index_active(), diagnostics_indexes_consistent(), and clear(), describing their behavior and return values where applicable. Update insert(block_type const&) documentation to include its std::size_t return value as the inserted or coalesced block size used by the maximum index.Source: Coding guidelines
24-26: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winHide the inherited public
eraseto protect the indexes.
free_list<block>::erase(const_iterator)is public and non-virtual. Any caller can erase a list node directly. Bothblocks_by_size_andblocks_by_address_then hold a danglingstd::listiterator, and the next lookup or rekey dereferences it.clear()is already shadowed for the same reason, so makeeraseunreachable as well.🛡️ Proposed guard
struct indexed_coalescing_free_list : free_list<block> { private: using base_type = free_list<block>; + + // Erasing through the base class would leave both indexes holding dangling list iterators. + using base_type::erase;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/rmm/mr/detail/indexed_coalescing_free_list.hpp` around lines 24 - 26, Hide the inherited free_list<block>::erase(const_iterator) in indexed_coalescing_free_list so callers cannot remove nodes without updating blocks_by_size_ and blocks_by_address_. Add the appropriate private guard alongside the existing clear() shadowing, preserving the indexed cleanup path for supported removals.
47-50: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueActivation-threshold constants do not follow the required naming convention. Both new
static constexprthresholds use lower_snake_case, while the guidelines require SCREAMING_SNAKE_CASE for constants.
cpp/include/rmm/mr/detail/indexed_coalescing_free_list.hpp#L47-L50: renameenable_index_thresholdtoENABLE_INDEX_THRESHOLDand update its uses inprepare_for_spliceandinsert_uncoalesced.cpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hpp#L838-L842: renamemaximum_index_activation_owner_counttoMAXIMUM_INDEX_ACTIVATION_OWNER_COUNTand update its use inactivate_maximum_index_if_needed.As per coding guidelines: "C++ constants and macros should use SCREAMING_SNAKE_CASE naming convention".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/rmm/mr/detail/indexed_coalescing_free_list.hpp` around lines 47 - 50, Rename the constant in cpp/include/rmm/mr/detail/indexed_coalescing_free_list.hpp:47-50 from enable_index_threshold to ENABLE_INDEX_THRESHOLD and update its uses in prepare_for_splice and insert_uncoalesced. Also rename maximum_index_activation_owner_count to MAXIMUM_INDEX_ACTIVATION_OWNER_COUNT in cpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hpp:838-842 and update its use in activate_maximum_index_if_needed, preserving behavior.Source: Coding guidelines
cpp/include/rmm/mr/detail/free_list.hpp (1)
150-157: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the parameters of the new
spliceoverload.The neighboring
splice(const_iterator, free_list&&)documents every parameter. This overload documents none. Add@paramentries. State thatListTypemust supply the three-argument node-transfersplice, because thenoexceptguarantee depends on a non-allocating transfer.📝 Proposed documentation
/** * `@brief` Transfers one existing node from `other` before `pos` without allocation. * * Iterators and references to the transferred element remain valid. + * + * `@note` Requires `ListType` to provide a non-allocating three-argument `splice`, which is what + * makes this operation `noexcept`. + * + * `@param` pos Iterator before which the element is inserted. `pos` may be the `end()` iterator. + * `@param` other The free list that currently owns the element. + * `@param` iter Iterator to the element in `other` to transfer. */ void splice(const_iterator pos, free_list& other, const_iterator iter) noexcept { blocks.splice(pos, other.blocks, iter); }As per path instructions for
cpp/include/rmm/**/*: "Doxygen documentation for all public functions/classes".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/rmm/mr/detail/free_list.hpp` around lines 150 - 157, Update the public free_list::splice(const_iterator pos, free_list& other, const_iterator iter) documentation with `@param` entries for pos, other, and iter, and state that ListType must provide a three-argument, non-allocating node-transfer splice so the noexcept guarantee remains valid.Source: Path instructions
cpp/src/mr/detail/indexed_pool_memory_resource_impl.cpp (3)
46-46: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winPrevent exceptions from escaping the destructor.
Line 46 calls
release()from the destructor.release()callsget_upstream_resource().deallocate_sync(...), which can throw. An exception that escapes a destructor callsstd::terminate.Make the destructor swallow errors, or add a
noexceptrelease path that usesRMM_CUDA_TRY_NOEXCEPTsemantics.As per coding guidelines: "Use RMM_CUDA_TRY_NOEXCEPT in destructors and noexcept functions for CUDA error checking".
🛡️ Proposed destructor guard
-indexed_pool_memory_resource_impl::~indexed_pool_memory_resource_impl() { release(); } +indexed_pool_memory_resource_impl::~indexed_pool_memory_resource_impl() +{ + try { + release(); + } catch (...) { + // Never propagate exceptions out of the destructor. + } +}Also applies to: 247-256
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/mr/detail/indexed_pool_memory_resource_impl.cpp` at line 46, Update indexed_pool_memory_resource_impl::~indexed_pool_memory_resource_impl() so cleanup cannot let exceptions escape; use a noexcept release path or catch failures around release(), applying RMM_CUDA_TRY_NOEXCEPT semantics to upstream deallocation. Ensure the corresponding release cleanup path is also safe when invoked from noexcept contexts.Source: Coding guidelines
170-179: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueGuard the capped-size subtraction against underflow.
Line 173 computes
maximum_pool_size_.value() - pool_size(). Ifcurrent_pool_size_ever exceeds the maximum, the unsigned subtraction wraps to a very large value, andsize_to_growthen reports large available headroom. The pool would grow past its cap.The current call paths appear to keep
current_pool_size_at or below the maximum. Add a defensive clamp so future changes to the growth or reclamation paths cannot break the cap silently.♻️ Proposed clamp
if (maximum_pool_size_.has_value()) { - auto const unaligned_remaining = maximum_pool_size_.value() - pool_size(); + auto const max_size = maximum_pool_size_.value(); + auto const unaligned_remaining = (max_size > pool_size()) ? (max_size - pool_size()) : 0; auto const remaining = rmm::align_up(unaligned_remaining, rmm::CUDA_ALLOCATION_ALIGNMENT);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/mr/detail/indexed_pool_memory_resource_impl.cpp` around lines 170 - 179, Update indexed_pool_memory_resource_impl::size_to_grow so the remaining capacity calculation clamps pool_size() to maximum_pool_size_.value() before subtraction, preventing unsigned underflow when the current pool exceeds the cap. Preserve the existing alignment and growth behavior when the pool is within its configured maximum.
101-104: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winMake initial pool construction exception-safe.
block_from_upstreamallocates beforeupstream_blocks_.emplace. Ifemplacethrows,try_to_expandretries without returning the allocation. Ifinsert_blockthrows, the derived destructor does not run, sorelease()does not return tracked blocks. Use RAII or explicit guards on both paths.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/src/mr/detail/indexed_pool_memory_resource_impl.cpp` around lines 101 - 104, Make the initial construction path around try_to_expand, upstream_blocks_.emplace, and insert_block exception-safe: ensure any allocation obtained by try_to_expand is reclaimed if tracking or insertion throws, and ensure already-tracked blocks are released if initial construction fails before the derived destructor runs. Use RAII or explicit scope guards while preserving normal successful initialization behavior.Source: Coding guidelines
cpp/include/rmm/mr/indexed_pool_memory_resource.hpp (1)
24-33: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider documenting stream semantics and pool growth limits in the class documentation.
The class documentation states thread safety. It does not state the stream-ordering semantics of
allocate/deallocate, nor the behavior whenmaximum_pool_sizeis reached. The implementation reclaims entirely free upstream blocks and can wait on another owner's event before it grows the pool. Users need this in the public documentation.Add a short paragraph that documents:
- Memory returned on a stream is reusable on that stream immediately, and on other streams after cross-stream synchronization performed by the resource.
- Behavior when the pool reaches
maximum_pool_size, includingrmm::out_of_memoryon failure to grow.As per coding guidelines: "Document thread-safety guarantees and stream semantics in class documentation for all public memory resources".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/include/rmm/mr/indexed_pool_memory_resource.hpp` around lines 24 - 33, Expand the class documentation for the indexed pool memory resource to describe stream ordering: memory is immediately reusable on the same stream and becomes reusable on other streams after the resource performs cross-stream synchronization. Also document behavior at maximum_pool_size, including reclamation or growth attempts and propagation of rmm::out_of_memory when growth fails.Source: Coding guidelines
cpp/tests/mr/pool_mr_tests.cpp (3)
381-412: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd zero-size and alignment edge-case tests for
indexed_pool_memory_resource.The new
IndexedPoolTestcases cover maximum-size allocation, partial-block retention, and cross-stream reclamation. They do not cover a zero-byte allocation or a non-default alignment. The indexed resource has its own free-list and index bookkeeping, so the existingpool_mrcoverage does not exercise these paths.Add cases for
allocate_sync(0)/deallocate_sync(ptr, 0)and for allocations with an alignment larger thanrmm::CUDA_ALLOCATION_ALIGNMENT.As per coding guidelines: "Edge case testing must cover zero-size allocations, alignment edge cases, and stream synchronization scenarios".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/mr/pool_mr_tests.cpp` around lines 381 - 412, Add IndexedPoolTest cases covering allocate_sync(0) followed by deallocate_sync(ptr, 0), and an allocation/deallocation using an alignment greater than rmm::CUDA_ALLOCATION_ALIGNMENT. Exercise these through indexed_pool_memory_resource so its free-list and index bookkeeping handle both edge cases.Source: Coding guidelines
1292-1374: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd a contended concurrent allocation and deallocation test.
PerThreadDefaultThreeThreadTransitivePublicationuses three threads, but thestagecounter serializes them. No test drives concurrentallocateanddeallocatecalls onindexed_pool_memory_resourcefrom several threads at the same time. The indexed free lists, the shared owner-maximum index, and the failure cache are all shared mutable state, so contention is the case that can expose missing synchronization.Add a stress test that runs several threads which allocate and free on their own streams in a loop, then verify that all pointers are unique while held and that the pool releases cleanly.
As per coding guidelines: "Concurrent allocation/deallocation should be tested with multiple threads to verify thread safety".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/mr/pool_mr_tests.cpp` around lines 1292 - 1374, Add a dedicated multithreaded stress test near PerThreadDefaultThreeThreadTransitivePublication that launches several threads, gives each its own CUDA stream, and repeatedly allocates and deallocates blocks concurrently through indexed_pool_memory_resource. Track currently held pointers with synchronization to assert no duplicate allocation, release each block before the next iteration, and join all workers while propagating failures; finally verify the pool can be cleaned up without outstanding allocations.Source: Coding guidelines
861-913: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRestore recovery hooks with RAII in
run_failure🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/mr/pool_mr_tests.cpp` around lines 861 - 913, Update run_failure to install recovery hooks through an RAII guard that captures their original values and restores them on scope exit, including early returns and exceptions. Ensure each test’s injected hooks are isolated and no hook state leaks into subsequent tests.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/benchmarks/pool_fragmentation/pool_fragmentation_bench.cpp`:
- Around line 262-266: Remove the post-loop PauseTiming()/ResumeTiming() pairs
in cpp/benchmarks/pool_fragmentation/pool_fragmentation_bench.cpp at lines
262-266 and 319-324, while preserving the timing pair inside the loop. Keep the
existing synchronization and cudaEventDestroy operations unchanged.
In `@cpp/include/rmm/mr/detail/indexed_coalescing_free_list.hpp`:
- Around line 222-236: Update free_list::insert(free_list&& other) to move each
block into the destination by passing it as an rvalue to the insertion path, and
ensure the source free list’s indexed metadata is emptied consistently after
transfer. Preserve largest_inserted tracking while preventing nodes or memory
from remaining observable in both lists.
In `@cpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hpp`:
- Around line 40-78: Invert the preprocessor branches around
indexed_recovery_test_hooks so RMM_INDEXED_RECOVERY_DISABLE_TEST_HOOKS selects
the existing testable implementation, while the default production path uses
direct cudaStreamWaitEvent/cudaEventRecord calls and a no-op
metadata_checkpoint(). Update the tests/build configuration to define the macro
when fault-injection hooks are required, preserving the existing hook APIs for
those tests.
- Around line 184-215: Make deallocate exception-safe by ensuring metadata
operations after underlying().free_block cannot leave allocation tracking,
published free blocks, or total_free_bytes_ inconsistent when they throw. Update
deallocate to stage required owner-block/event metadata before freeing, or
provide rollback covering get_or_create_owner_blocks, blocks.insert, and related
bookkeeping while preserving its noexcept contract.
In `@cpp/tests/mr/pool_mr_tests.cpp`:
- Around line 1644-1654: Move the release_guard declaration in
ExpansionPublishesRemainderReadiness to after the upstream
delayed_async_memory_resource and indexed_pool_memory_resource pool are
constructed, matching the ordering used by
CrossStreamStealAfterMergeWaitsForDonorStream. Preserve the guard’s existing
flag and release behavior so it is destroyed first during teardown.
---
Nitpick comments:
In `@cpp/benchmarks/pool_fragmentation/pool_fragmentation_bench.cpp`:
- Around line 29-32: Rename all newly introduced C++ constexpr constants in
fragment_size and the other referenced benchmark sections to
SCREAMING_SNAKE_CASE, including alignment to ALIGNMENT and
first_lookup_iterations to FIRST_LOOKUP_ITERATIONS; update every use
consistently without changing behavior.
In `@cpp/include/rmm/mr/detail/free_list.hpp`:
- Around line 150-157: Update the public free_list::splice(const_iterator pos,
free_list& other, const_iterator iter) documentation with `@param` entries for
pos, other, and iter, and state that ListType must provide a three-argument,
non-allocating node-transfer splice so the noexcept guarantee remains valid.
In `@cpp/include/rmm/mr/detail/indexed_coalescing_free_list.hpp`:
- Around line 334-347: Add Doxygen documentation to the public members
diagnostics_index_active(), diagnostics_indexes_consistent(), and clear(),
describing their behavior and return values where applicable. Update
insert(block_type const&) documentation to include its std::size_t return value
as the inserted or coalesced block size used by the maximum index.
- Around line 24-26: Hide the inherited free_list<block>::erase(const_iterator)
in indexed_coalescing_free_list so callers cannot remove nodes without updating
blocks_by_size_ and blocks_by_address_. Add the appropriate private guard
alongside the existing clear() shadowing, preserving the indexed cleanup path
for supported removals.
- Around line 47-50: Rename the constant in
cpp/include/rmm/mr/detail/indexed_coalescing_free_list.hpp:47-50 from
enable_index_threshold to ENABLE_INDEX_THRESHOLD and update its uses in
prepare_for_splice and insert_uncoalesced. Also rename
maximum_index_activation_owner_count to MAXIMUM_INDEX_ACTIVATION_OWNER_COUNT in
cpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hpp:838-842 and
update its use in activate_maximum_index_if_needed, preserving behavior.
In `@cpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hpp`:
- Around line 80-95: Remove the duplicated indexed_crtp definition and reuse the
existing crtp helper from the stream-ordered resource header, updating
references to the existing symbol as needed. If the helper remains locally,
correct its documentation URL to
https://www.fluentcpp.com/2017/05/19/crtp-helper/.
- Around line 613-626: Replace the assertions in the candidate-selection path
around stream_free_blocks_ and owner_blocks.find_block with runtime validation:
if the donor lookup is missing or selection.block equals owner_blocks.end(),
skip this candidate and continue searching. Ensure allocate_from_selection is
called only with a valid owner iterator and block selection, while preserving
the existing empty-result behavior when no candidates remain.
In `@cpp/include/rmm/mr/indexed_pool_memory_resource.hpp`:
- Around line 24-33: Expand the class documentation for the indexed pool memory
resource to describe stream ordering: memory is immediately reusable on the same
stream and becomes reusable on other streams after the resource performs
cross-stream synchronization. Also document behavior at maximum_pool_size,
including reclamation or growth attempts and propagation of rmm::out_of_memory
when growth fails.
In `@cpp/src/mr/detail/indexed_pool_memory_resource_impl.cpp`:
- Line 46: Update
indexed_pool_memory_resource_impl::~indexed_pool_memory_resource_impl() so
cleanup cannot let exceptions escape; use a noexcept release path or catch
failures around release(), applying RMM_CUDA_TRY_NOEXCEPT semantics to upstream
deallocation. Ensure the corresponding release cleanup path is also safe when
invoked from noexcept contexts.
- Around line 170-179: Update indexed_pool_memory_resource_impl::size_to_grow so
the remaining capacity calculation clamps pool_size() to
maximum_pool_size_.value() before subtraction, preventing unsigned underflow
when the current pool exceeds the cap. Preserve the existing alignment and
growth behavior when the pool is within its configured maximum.
- Around line 101-104: Make the initial construction path around try_to_expand,
upstream_blocks_.emplace, and insert_block exception-safe: ensure any allocation
obtained by try_to_expand is reclaimed if tracking or insertion throws, and
ensure already-tracked blocks are released if initial construction fails before
the derived destructor runs. Use RAII or explicit scope guards while preserving
normal successful initialization behavior.
In `@cpp/tests/mr/pool_mr_tests.cpp`:
- Around line 381-412: Add IndexedPoolTest cases covering allocate_sync(0)
followed by deallocate_sync(ptr, 0), and an allocation/deallocation using an
alignment greater than rmm::CUDA_ALLOCATION_ALIGNMENT. Exercise these through
indexed_pool_memory_resource so its free-list and index bookkeeping handle both
edge cases.
- Around line 1292-1374: Add a dedicated multithreaded stress test near
PerThreadDefaultThreeThreadTransitivePublication that launches several threads,
gives each its own CUDA stream, and repeatedly allocates and deallocates blocks
concurrently through indexed_pool_memory_resource. Track currently held pointers
with synchronization to assert no duplicate allocation, release each block
before the next iteration, and join all workers while propagating failures;
finally verify the pool can be cleaned up without outstanding allocations.
- Around line 861-913: Update run_failure to install recovery hooks through an
RAII guard that captures their original values and restores them on scope exit,
including early returns and exceptions. Ensure each test’s injected hooks are
isolated and no hook state leaks into subsequent tests.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c3eddf66-fa1f-487d-a17a-f0e851768028
📒 Files selected for processing (15)
cpp/CMakeLists.txtcpp/benchmarks/CMakeLists.txtcpp/benchmarks/device_uvector/device_uvector_bench.cucpp/benchmarks/multi_stream_allocations/multi_stream_allocations_bench.cucpp/benchmarks/pool_fragmentation/pool_fragmentation_bench.cppcpp/benchmarks/random_allocations/random_allocations.cppcpp/benchmarks/replay/replay.cppcpp/include/rmm/mr/detail/free_list.hppcpp/include/rmm/mr/detail/indexed_coalescing_free_list.hppcpp/include/rmm/mr/detail/indexed_pool_memory_resource_impl.hppcpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hppcpp/include/rmm/mr/indexed_pool_memory_resource.hppcpp/src/mr/detail/indexed_pool_memory_resource_impl.cppcpp/src/mr/indexed_pool_memory_resource.cppcpp/tests/mr/pool_mr_tests.cpp
| #ifdef RMM_INDEXED_RECOVERY_DISABLE_TEST_HOOKS | ||
| struct indexed_recovery_test_hooks { | ||
| static cudaError_t wait(cudaStream_t stream, cudaEvent_t event, unsigned int flags) | ||
| { return cudaStreamWaitEvent(stream, event, flags); } | ||
|
|
||
| static cudaError_t record(cudaEvent_t event, cudaStream_t stream) | ||
| { return cudaEventRecord(event, stream); } | ||
|
|
||
| static void metadata_checkpoint() noexcept {} | ||
| }; | ||
| #else | ||
| struct indexed_recovery_test_hooks { | ||
| using wait_function = cudaError_t (*)(cudaStream_t, cudaEvent_t, unsigned int); | ||
| using record_function = cudaError_t (*)(cudaEvent_t, cudaStream_t); | ||
|
|
||
| static cudaError_t default_wait(cudaStream_t stream, cudaEvent_t event, unsigned int flags) | ||
| { return cudaStreamWaitEvent(stream, event, flags); } | ||
|
|
||
| static cudaError_t default_record(cudaEvent_t event, cudaStream_t stream) | ||
| { return cudaEventRecord(event, stream); } | ||
|
|
||
| static void metadata_checkpoint() | ||
| { | ||
| if (metadata_fail_after == 0) { throw std::bad_alloc{}; } | ||
| if (metadata_fail_after > 0) { --metadata_fail_after; } | ||
| } | ||
|
|
||
| static void reset() noexcept | ||
| { | ||
| wait = default_wait; | ||
| record = default_record; | ||
| metadata_fail_after = -1; | ||
| } | ||
|
|
||
| inline static thread_local wait_function wait{default_wait}; | ||
| inline static thread_local record_function record{default_record}; | ||
| inline static thread_local int metadata_fail_after{-1}; | ||
| }; | ||
| #endif |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Invert the test-hook macro so production builds exclude the hooks.
RMM_INDEXED_RECOVERY_DISABLE_TEST_HOOKS is an opt-out. No build defines it, so the fault-injection version is the default. Two consequences follow in normal builds:
cudaStreamWaitEventandcudaEventRecordare reached through mutablethread_localfunction pointers on the allocation path, which blocks inlining in a latency-sensitive allocator.metadata_checkpoint()becomes a throwing call that consumer code can arm by writingmetadata_fail_after, through an installed public header.
Make the hooks opt-in and let the tests define the macro.
🔒️ Proposed change
-#ifdef RMM_INDEXED_RECOVERY_DISABLE_TEST_HOOKS
+#ifndef RMM_INDEXED_RECOVERY_ENABLE_TEST_HOOKS
struct indexed_recovery_test_hooks {📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| #ifdef RMM_INDEXED_RECOVERY_DISABLE_TEST_HOOKS | |
| struct indexed_recovery_test_hooks { | |
| static cudaError_t wait(cudaStream_t stream, cudaEvent_t event, unsigned int flags) | |
| { return cudaStreamWaitEvent(stream, event, flags); } | |
| static cudaError_t record(cudaEvent_t event, cudaStream_t stream) | |
| { return cudaEventRecord(event, stream); } | |
| static void metadata_checkpoint() noexcept {} | |
| }; | |
| #else | |
| struct indexed_recovery_test_hooks { | |
| using wait_function = cudaError_t (*)(cudaStream_t, cudaEvent_t, unsigned int); | |
| using record_function = cudaError_t (*)(cudaEvent_t, cudaStream_t); | |
| static cudaError_t default_wait(cudaStream_t stream, cudaEvent_t event, unsigned int flags) | |
| { return cudaStreamWaitEvent(stream, event, flags); } | |
| static cudaError_t default_record(cudaEvent_t event, cudaStream_t stream) | |
| { return cudaEventRecord(event, stream); } | |
| static void metadata_checkpoint() | |
| { | |
| if (metadata_fail_after == 0) { throw std::bad_alloc{}; } | |
| if (metadata_fail_after > 0) { --metadata_fail_after; } | |
| } | |
| static void reset() noexcept | |
| { | |
| wait = default_wait; | |
| record = default_record; | |
| metadata_fail_after = -1; | |
| } | |
| inline static thread_local wait_function wait{default_wait}; | |
| inline static thread_local record_function record{default_record}; | |
| inline static thread_local int metadata_fail_after{-1}; | |
| }; | |
| #endif | |
| #ifndef RMM_INDEXED_RECOVERY_ENABLE_TEST_HOOKS | |
| struct indexed_recovery_test_hooks { | |
| static cudaError_t wait(cudaStream_t stream, cudaEvent_t event, unsigned int flags) | |
| { return cudaStreamWaitEvent(stream, event, flags); } | |
| static cudaError_t record(cudaEvent_t event, cudaStream_t stream) | |
| { return cudaEventRecord(event, stream); } | |
| static void metadata_checkpoint() noexcept {} | |
| }; | |
| #else | |
| struct indexed_recovery_test_hooks { | |
| using wait_function = cudaError_t (*)(cudaStream_t, cudaEvent_t, unsigned int); | |
| using record_function = cudaError_t (*)(cudaEvent_t, cudaStream_t); | |
| static cudaError_t default_wait(cudaStream_t stream, cudaEvent_t event, unsigned int flags) | |
| { return cudaStreamWaitEvent(stream, event, flags); } | |
| static cudaError_t default_record(cudaEvent_t event, cudaStream_t stream) | |
| { return cudaEventRecord(event, stream); } | |
| static void metadata_checkpoint() | |
| { | |
| if (metadata_fail_after == 0) { throw std::bad_alloc{}; } | |
| if (metadata_fail_after > 0) { --metadata_fail_after; } | |
| } | |
| static void reset() noexcept | |
| { | |
| wait = default_wait; | |
| record = default_record; | |
| metadata_fail_after = -1; | |
| } | |
| inline static thread_local wait_function wait{default_wait}; | |
| inline static thread_local record_function record{default_record}; | |
| inline static thread_local int metadata_fail_after{-1}; | |
| }; | |
| #endif |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hpp` around
lines 40 - 78, Invert the preprocessor branches around
indexed_recovery_test_hooks so RMM_INDEXED_RECOVERY_DISABLE_TEST_HOOKS selects
the existing testable implementation, while the default production path uses
direct cudaStreamWaitEvent/cudaEventRecord calls and a no-op
metadata_checkpoint(). Update the tests/build configuration to define the macro
when fault-injection hooks are required, preserving the existing hook APIs for
those tests.
There was a problem hiding this comment.
We are intentionally deferring changes to the fault-injection hook architecture pending maintainer direction. A macro-only switch in this installed inline template header could give test and library translation units different definitions and create an ODR risk. If maintainers want production hooks removed, we propose using an explicit policy or test-only instantiation instead. The current tradeoff is documented in the PR.
There was a problem hiding this comment.
@revans2, agreed. A per-translation-unit macro in this installed template header can produce different inline definitions and create an ODR risk.
The production-path overhead remains a valid concern, but the proposed macro-only change is not safe without a consistent build-wide definition. I will leave this thread open pending maintainer direction on an explicit policy or test-only instantiation design.
Would you like to open a follow-up GitHub issue for the hook architecture?
✏️ Learnings added
Learnt from: revans2
Repo: rapidsai/rmm PR: 2507
File: cpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hpp:40-78
Timestamp: 2026-08-13T13:43:45.314Z
Learning: In `cpp/include/rmm/mr/detail/indexed_stream_ordered_memory_resource.hpp`, do not use a macro-only switch for `indexed_recovery_test_hooks` unless every library and consumer translation unit receives the same definition. This installed inline template header can otherwise create ODR risk. Changes to remove production fault-injection hooks are deferred pending maintainer direction toward an explicit policy or test-only instantiation.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
You are interacting with an AI system.
1 participant
Description
To be 100% clear this is AI slop for the most part. I tried hard to make sure that it fixes the issues that I am seeing, and not just passes the tests, but does the right thing, but I am no expert on memory allocateors. So, any comments you have I am happy to work with/around.
I am very happy to have some or all of these changes become a part of the pooled allocator itself, but if people want to keep them separate I totally understand. In Spark we have two separate pools for host memory, currently. One is for pinned memory, and the other is for pageable memory. If the data is going to go to/from the GPU we obviously want pinned memory for it, but we also do not want to fail a job if pinned memory is not available so we fall back to using pagable memory. This means the time to fail an allocation is also very important to us. Knowing quickly that it is not going to work so we can try with pageable memory is important.
This adds
indexed_pool_memory_resource, a separate stream-ordered pooled resource for workloadswhere fragmented best-fit lookup and cross-owner recovery dominate allocation latency.
There are a few things that I am still working on with the AI, like making sure that the overhead of the extra indexing does not cause problems if we have mixed workloads, where there may be periods with lots of small allocations followed by periods with larger allocations. Right now once it identifies a job as needs the extra indexing it just stays there. Depending on what we want to do with this I am also happy to update the documentation.
Why
pool_memory_resourcefinds best-fit blocks by walking an owner-local coalescing free list. Whenthat owner cannot satisfy a request, recovery may wait on donor owners and merge complete donor
trees before searching again. The cost grows with fragmentation depth and retained owner state.
The new resource preserves the stock pool's stream ordering, shared capacity, address coalescing,
best-fit allocation, cross-owner reuse, and capped-pool reclamation while changing how free blocks
are found and recovered.
Changes
capped-pool reclamation.
This is a new type; it does not change the ABI or behavior of
pool_memory_resource.Performance versus
pool_memory_resourceThe following Release SM86 microbenchmarks used six isolated processes per resource with
interleaved ordering. Change is
(indexed / stock - 1) × 100; negative values mean lower latency.pool_memory_resourceindexed_pool_memory_resourceValidation
git diff --checkpasses.Current tradeoffs
not measure a long-lived pool after a transient large query activates the indexes; that
persistent-pool behavior needs explicit measurement.
pool_memory_resource, which avoids changingits ABI or behavior but temporarily duplicates some infrastructure pending maintainer direction.
feedback.
Checklist