Caching allocator - #2480
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
📝 WalkthroughSummary by CodeRabbit
WalkthroughChangesCaching memory resource
Estimated code review effort: 4 (Complex) | ~60 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 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: 3
🧹 Nitpick comments (2)
cpp/src/mr/detail/caching_memory_resource_impl.cpp (1)
32-36: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
release()triggersmerge_all_free_blocks()twice.
release()callsrelease_pool(true)thenrelease_pool(false), and each independently callsmerge_all_free_blocks()(seerelease_pool, lines 176-178). For a fullrelease(), the second merge is redundant since the first already merged/synced everything into a single free list.🤖 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/caching_memory_resource_impl.cpp` around lines 32 - 36, The caching_memory_resource_impl::release() path is doing redundant work because release_pool(true) and release_pool(false) each invoke merge_all_free_blocks(), so the second merge is unnecessary after a full release. Update release() (and/or release_pool) so the full-release path only merges free blocks once while still releasing both pools, and keep the behavior localized around caching_memory_resource_impl::release(), release_pool(), and merge_all_free_blocks().cpp/include/rmm/mr/caching_memory_resource.hpp (1)
73-94: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument blocking/synchronization behavior of
release*APIs.
release(),release_small_blocks(), andrelease_large_blocks()are documented only in terms of bytes freed, but the implementation (seecaching_memory_resource_impl.cpp) synchronizes CUDA events across streams before releasing segments, which can block the calling thread. Worth calling out in the Doxygen comments given the class docstring already commits to documenting thread-safety/stream semantics.As per path instructions, "For public C++ API headers, additionally check: 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/caching_memory_resource.hpp` around lines 73 - 94, Update the Doxygen comments for caching_memory_resource::release(), release_small_blocks(), and release_large_blocks() to explicitly mention their blocking/synchronization behavior; these APIs do more than return bytes to the upstream resource because they wait on CUDA event/stream synchronization before freeing cached segments. Add this note in the public header near the existing declarations so the documented thread/stream semantics match the implementation in caching_memory_resource_impl.cpp.Source: Path instructions
🤖 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 `@caching_allocator.md`:
- Around line 147-157: The Stream-Ordered Reuse section in caching_allocator.md
conflicts with the existing test coverage for caching_memory_resource. Update
this paragraph to match the actual coverage by referencing the cross-stream
scenarios already exercised in caching_mr_tests.cpp, including
ReleaseMergesCrossStreamFreeBlocks, ReusesOtherStreamAllocation, and
ReusesDeletedStreamAllocation, and remove or soften the claim that no dedicated
cross-stream correctness test exists.
In `@cpp/src/mr/detail/caching_memory_resource_impl.cpp`:
- Around line 111-134: free_block() currently fabricates a block for pointers
not found in allocated_blocks_, which lets foreign or double-freed pointers
enter the stream free list. Update caching_memory_resource_impl::free_block to
treat any ptr that is not an exact match in allocated_blocks_ and not provably
valid from upstream_blocks_ as an error path, and avoid returning a reusable
block for it. Make stream_ordered_memory_resource::deallocate rely on this
validation so invalid pointers are rejected with a hard failure or log instead
of being inserted into the free list.
In `@cpp/tests/mr/caching_mr_tests.cpp`:
- Around line 1-4: The SPDX header in this C++ test file does not match the
required copyright template. Update the existing header comment to use the
mandated C++/CUDA form with the 2025-2026 year range and the “NVIDIA CORPORATION
& AFFILIATES. All rights reserved.” suffix, matching the project’s standard
SPDX-FileCopyrightText header style.
---
Nitpick comments:
In `@cpp/include/rmm/mr/caching_memory_resource.hpp`:
- Around line 73-94: Update the Doxygen comments for
caching_memory_resource::release(), release_small_blocks(), and
release_large_blocks() to explicitly mention their blocking/synchronization
behavior; these APIs do more than return bytes to the upstream resource because
they wait on CUDA event/stream synchronization before freeing cached segments.
Add this note in the public header near the existing declarations so the
documented thread/stream semantics match the implementation in
caching_memory_resource_impl.cpp.
In `@cpp/src/mr/detail/caching_memory_resource_impl.cpp`:
- Around line 32-36: The caching_memory_resource_impl::release() path is doing
redundant work because release_pool(true) and release_pool(false) each invoke
merge_all_free_blocks(), so the second merge is unnecessary after a full
release. Update release() (and/or release_pool) so the full-release path only
merges free blocks once while still releasing both pools, and keep the behavior
localized around caching_memory_resource_impl::release(), release_pool(), and
merge_all_free_blocks().
🪄 Autofix (Beta)
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: 5b302f4c-c802-445e-9e33-1332f62cf739
📒 Files selected for processing (12)
caching_allocator.mdcpp/CMakeLists.txtcpp/benchmarks/multi_stream_allocations/multi_stream_allocations_bench.cucpp/benchmarks/random_allocations/random_allocations.cppcpp/include/rmm/mr/caching_memory_resource.hppcpp/include/rmm/mr/detail/caching_memory_resource_impl.hppcpp/include/rmm/mr/detail/coalescing_free_list.hppcpp/include/rmm/mr/detail/stream_ordered_memory_resource.hppcpp/src/mr/caching_memory_resource.cppcpp/src/mr/detail/caching_memory_resource_impl.cppcpp/tests/CMakeLists.txtcpp/tests/mr/caching_mr_tests.cpp
| ### Stream-Ordered Reuse | ||
|
|
||
| The resource inherits the stream ordering model from `stream_ordered_memory_resource`. | ||
|
|
||
| This means allocation uses a stream-associated event and free-list. Same-stream blocks can be reused | ||
| without waiting. Blocks freed on other streams can be reused after the allocating stream waits on the | ||
| other stream's event. | ||
|
|
||
| This has not yet received a dedicated cross-stream correctness test for `caching_memory_resource`. | ||
| It is inherited from existing RMM infrastructure, but because this resource has new block sizing and | ||
| release behavior, it still needs direct tests. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Contradicts test coverage claimed later in the same doc.
This section says caching_memory_resource "has not yet received a dedicated cross-stream correctness test," but the "Covered scenarios" list (lines 227-228) and caching_mr_tests.cpp both include cross-stream free-list-merge and cross-stream-reuse tests (ReleaseMergesCrossStreamFreeBlocks, ReusesOtherStreamAllocation, ReusesDeletedStreamAllocation). Recommend reconciling this section with the actual test coverage.
🤖 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 `@caching_allocator.md` around lines 147 - 157, The Stream-Ordered Reuse
section in caching_allocator.md conflicts with the existing test coverage for
caching_memory_resource. Update this paragraph to match the actual coverage by
referencing the cross-stream scenarios already exercised in
caching_mr_tests.cpp, including ReleaseMergesCrossStreamFreeBlocks,
ReusesOtherStreamAllocation, and ReusesDeletedStreamAllocation, and remove or
soften the claim that no dedicated cross-stream correctness test exists.
| caching_memory_resource_impl::block_type caching_memory_resource_impl::free_block( | ||
| void* ptr, std::size_t size) noexcept | ||
| { | ||
| auto const iter = allocated_blocks_.find(static_cast<char*>(ptr)); | ||
| if (iter != allocated_blocks_.end()) { | ||
| auto block = *iter; | ||
| allocated_blocks_.erase(iter); | ||
| return block; | ||
| } | ||
|
|
||
| auto const aligned_size = rmm::align_up(size, rmm::CUDA_ALLOCATION_ALIGNMENT); | ||
| auto const upstream_iter = | ||
| upstream_blocks_.upper_bound(segment{static_cast<char*>(ptr), 0, false}); | ||
| if (upstream_iter != upstream_blocks_.begin()) { | ||
| auto const candidate = std::prev(upstream_iter); | ||
| auto const segment_end = candidate->pointer + candidate->size; | ||
| if (candidate->pointer <= static_cast<char*>(ptr) && static_cast<char*>(ptr) < segment_end) { | ||
| return block_type{ | ||
| static_cast<char*>(ptr), aligned_size, candidate->pointer == static_cast<char*>(ptr)}; | ||
| } | ||
| } | ||
|
|
||
| return block_type{static_cast<char*>(ptr), aligned_size, false}; | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant file and surrounding symbols.
ast-grep outline cpp/src/mr/detail/caching_memory_resource_impl.cpp --view expanded || true
# Read the relevant region with line numbers.
sed -n '1,240p' cpp/src/mr/detail/caching_memory_resource_impl.cpp | cat -n
# Find the caller and related deallocation path.
rg -n "free_block\\(|deallocate\\(|allocated_blocks_|upstream_blocks_|stream_ordered_memory_resource" cpp/src/mr/detail/caching_memory_resource_impl.cpp cpp/include cpp/src -g '!**/build/**'Repository: rapidsai/rmm
Length of output: 30452
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the deallocation path that consumes free_block().
sed -n '120,240p' cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp | cat -n
# Inspect the coalescing free-list insertion/merge behavior.
rg -n "class coalescing_free_list|insert\\(|get_block\\(|merge|remove_free_block|add_free_block" cpp/include/rmm/mr/detail cpp/src/mr/detail -g '!**/build/**'
# Read the relevant free-list implementation if present.
fd -a "coalescing_free_list*" cpp/include cpp/srcRepository: rapidsai/rmm
Length of output: 15679
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the coalescing free-list implementation.
sed -n '1,320p' cpp/include/rmm/mr/detail/coalescing_free_list.hpp | cat -n
# Read the block/segment comparison helpers if they are in the same area.
sed -n '1,240p' cpp/include/rmm/mr/detail/free_list.hpp | cat -nRepository: rapidsai/rmm
Length of output: 18720
Reject untracked pointers in free_block()
free_block() should not synthesize a plausible block when allocated_blocks_ has no match. stream_ordered_memory_resource::deallocate() inserts whatever comes back into the stream free list, so a double-free or foreign pointer can corrupt coalescing and later hand out overlapping live allocations. Add a hard failure/logging path and keep invalid blocks out of the free list.
🤖 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/caching_memory_resource_impl.cpp` around lines 111 - 134,
free_block() currently fabricates a block for pointers not found in
allocated_blocks_, which lets foreign or double-freed pointers enter the stream
free list. Update caching_memory_resource_impl::free_block to treat any ptr that
is not an exact match in allocated_blocks_ and not provably valid from
upstream_blocks_ as an error path, and avoid returning a reusable block for it.
Make stream_ordered_memory_resource::deallocate rely on this validation so
invalid pointers are rejected with a hard failure or log instead of being
inserted into the free list.
| /* | ||
| * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
SPDX header doesn't match the required template.
Header omits the "2025-" start year and "& AFFILIATES. All rights reserved." suffix mandated for C++/CUDA files.
As per coding guidelines: "SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved."
📝 Proposed fix
/*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*/📝 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.
| /* | |
| * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. | |
| * SPDX-License-Identifier: Apache-2.0 | |
| */ | |
| /* | |
| * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| * SPDX-License-Identifier: Apache-2.0 | |
| */ |
🤖 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/caching_mr_tests.cpp` around lines 1 - 4, The SPDX header in
this C++ test file does not match the required copyright template. Update the
existing header comment to use the mandated C++/CUDA form with the 2025-2026
year range and the “NVIDIA CORPORATION & AFFILIATES. All rights reserved.”
suffix, matching the project’s standard SPDX-FileCopyrightText header style.
Source: Coding guidelines
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/src/mr/detail/caching_memory_resource_impl.cpp (2)
57-65: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winHIGH: Synchronize cache-accounting reads.
These accessors read non-atomic counters while allocation and release paths mutate them under the resource mutex. Concurrent callers can cause a C++ data race. Lock these reads or use atomic accounting.
As per coding guidelines, “Memory resources must be thread-safe by default with proper synchronization.”
🤖 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/caching_memory_resource_impl.cpp` around lines 57 - 65, Synchronize the cached_small_bytes() and cached_large_bytes() accessors with the resource mutex before reading their counters, matching the locking used by allocation and release paths. Preserve their const noexcept behavior and return values while ensuring concurrent accounting updates cannot race.Source: Coding guidelines
138-142: 🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy liftCRITICAL: Make allocator metadata updates exception-safe.
If either tracking-container insertion throws after ownership has moved, the block is no longer recoverable:
allocated_blocks_loses a removed free block, andupstream_blocks_leaks an upstream allocation. Roll back the free-list transition or deallocate upstream via a scope guard before rethrowing.As per coding guidelines, “Device memory allocated must be properly returned to upstream on all error paths to prevent GPU memory leaks.”
Also applies to: 285-295
🤖 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/caching_memory_resource_impl.cpp` around lines 138 - 142, Make the metadata updates in the allocation path exception-safe around the visible allocated_blocks_.insert(alloc) and the corresponding upstream_blocks_ insertion near the additionally referenced code. Add rollback or scope-guard handling so any container insertion failure restores the removed free-list block or returns the upstream allocation before rethrowing, while preserving the normal split and return behavior.Source: Coding guidelines
🧹 Nitpick comments (1)
cpp/include/rmm/mr/caching_memory_resource.hpp (1)
64-75: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the observable semantics of the new policies.
These public
@paramdescriptions are too generic: explainrelease_allversusrelease_oversized_then_all,separateversusunified, and thatcross_streamenables event-waited reuse whilesame_streamdoes not. Also reflect these guarantees in the class-level stream-semantics documentation.As per path instructions, public C++ API changes must be documented and public memory resources must document stream semantics and thread-safety guarantees.
Suggested documentation update
- * `@param` oom_fallback_policy Policy used after an upstream allocation failure. - * `@param` pool_policy Policy controlling cross-pool cached block reuse. - * `@param` stream_reuse_policy Policy controlling cross-stream cached block reuse. + * `@param` oom_fallback_policy Controls which cached segments are released before retrying + * an upstream allocation. + * `@param` pool_policy Controls whether small and large cached segments may satisfy each + * other's requests. + * `@param` stream_reuse_policy Controls whether cached blocks may be reused across streams; + * cross-stream reuse inserts the required CUDA event waits.🤖 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/caching_memory_resource.hpp` around lines 64 - 75, Update the public documentation for the policy parameters in caching_memory_resource: describe the observable behavior of release_all versus release_oversized_then_all, separate versus unified pool reuse, and cross_stream versus same_stream, including that cross_stream uses event-waited reuse while same_stream does not. Extend the class-level stream-semantics documentation to reflect these guarantees and document the resource’s thread-safety guarantees as required for public memory resources.Source: Path instructions
🤖 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/caching_allocator_reports.py`:
- Around line 141-157: Update the default_build path used by the argument parser
to the repository-supported standard build location, cpp/build, so the
--random-bench and --multi-stream-bench defaults resolve to benchmarks produced
by the documented CMake workflow. Keep the existing argument names and benchmark
executable paths unchanged.
In `@cpp/src/mr/detail/caching_memory_resource_impl.cpp`:
- Around line 80-82: Update
caching_memory_resource_impl::get_maximum_allocation_size() to return the
largest size that can be safely rounded up to CUDA_ALLOCATION_ALIGNMENT, rather
than std::numeric_limits<std::size_t>::max(). Ensure near-limit allocation
requests are capped or rejected before alignment can overflow.
---
Outside diff comments:
In `@cpp/src/mr/detail/caching_memory_resource_impl.cpp`:
- Around line 57-65: Synchronize the cached_small_bytes() and
cached_large_bytes() accessors with the resource mutex before reading their
counters, matching the locking used by allocation and release paths. Preserve
their const noexcept behavior and return values while ensuring concurrent
accounting updates cannot race.
- Around line 138-142: Make the metadata updates in the allocation path
exception-safe around the visible allocated_blocks_.insert(alloc) and the
corresponding upstream_blocks_ insertion near the additionally referenced code.
Add rollback or scope-guard handling so any container insertion failure restores
the removed free-list block or returns the upstream allocation before
rethrowing, while preserving the normal split and return behavior.
---
Nitpick comments:
In `@cpp/include/rmm/mr/caching_memory_resource.hpp`:
- Around line 64-75: Update the public documentation for the policy parameters
in caching_memory_resource: describe the observable behavior of release_all
versus release_oversized_then_all, separate versus unified pool reuse, and
cross_stream versus same_stream, including that cross_stream uses event-waited
reuse while same_stream does not. Extend the class-level stream-semantics
documentation to reflect these guarantees and document the resource’s
thread-safety guarantees as required for public memory resources.
🪄 Autofix (Beta)
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: 7a301b87-ff4c-4759-84ca-1b5e7694806d
📒 Files selected for processing (11)
.gitignorecaching_allocator_plan.mdcpp/benchmarks/caching_allocator_reports.pycpp/benchmarks/multi_stream_allocations/multi_stream_allocations_bench.cucpp/benchmarks/random_allocations/random_allocations.cppcpp/include/rmm/mr/caching_memory_resource.hppcpp/include/rmm/mr/detail/caching_memory_resource_impl.hppcpp/include/rmm/mr/detail/stream_ordered_memory_resource.hppcpp/src/mr/caching_memory_resource.cppcpp/src/mr/detail/caching_memory_resource_impl.cppcpp/tests/mr/caching_mr_tests.cpp
🚧 Files skipped from review as they are similar to previous changes (4)
- cpp/benchmarks/random_allocations/random_allocations.cpp
- cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp
- cpp/benchmarks/multi_stream_allocations/multi_stream_allocations_bench.cu
- cpp/tests/mr/caching_mr_tests.cpp
| default_build = root / "cpp" / "build-bench-opencode" / "gbenchmarks" | ||
| parser = argparse.ArgumentParser( | ||
| description="Generate caching allocator benchmark reports." | ||
| ) | ||
| parser.add_argument( | ||
| "--random-bench", | ||
| type=Path, | ||
| default=default_build / "RANDOM_ALLOCATIONS_BENCH", | ||
| ) | ||
| parser.add_argument( | ||
| "--multi-stream-bench", | ||
| type=Path, | ||
| default=default_build / "MULTI_STREAM_ALLOCATIONS_BENCH", | ||
| ) | ||
| parser.add_argument( | ||
| "--output-dir", type=Path, default=root / "caching-allocator-reports" | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
HIGH: Use a repository-supported default benchmark path.
cpp/build-bench-opencode is not produced by the documented build workflows, so the script fails with default arguments on normal builds. Default to the standard cpp/build benchmark output location or discover it explicitly.
As per coding guidelines, “cmake -S cpp -B cpp/build -DBUILD_TESTS=ON -DCMAKE_BUILD_TYPE=Release”.
🤖 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/caching_allocator_reports.py` around lines 141 - 157, Update
the default_build path used by the argument parser to the repository-supported
standard build location, cpp/build, so the --random-bench and
--multi-stream-bench defaults resolve to benchmarks produced by the documented
CMake workflow. Keep the existing argument names and benchmark executable paths
unchanged.
Source: Coding guidelines
| std::size_t caching_memory_resource_impl::get_maximum_allocation_size() const | ||
| { | ||
| return std::numeric_limits<std::size_t>::max(); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant implementation and allocation paths.
git ls-files 'cpp/src/mr/detail/*' 'cpp/include/rmm/mr/*' | sed -n '1,200p'
echo '---'
rg -n "get_maximum_allocation_size|align_up\(|maximum_allocation_size|stream_ordered_memory_resource::allocate|allocate\(" cpp/src cpp/include/rmm/mr -g '!**/build/**' -g '!**/dist/**'Repository: rapidsai/rmm
Length of output: 27173
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the relevant file sections with line numbers.
FILE=cpp/src/mr/detail/caching_memory_resource_impl.cpp
wc -l "$FILE"
echo '---'
cat -n "$FILE" | sed -n '1,220p'
echo '---'
# Inspect the stream-ordered allocation implementation if present.
FILE2=$(rg -l "stream_ordered_memory_resource" cpp/include/rmm/mr cpp/src | head -n 1 || true)
if [[ -n "${FILE2:-}" ]]; then
echo "FILE2=$FILE2"
wc -l "$FILE2"
echo '---'
cat -n "$FILE2" | sed -n '1,260p'
fiRepository: rapidsai/rmm
Length of output: 16286
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Probe the exact rounding behavior and compare against size limits.
python3 - <<'PY'
import sys
from math import isfinite
SIZE_MAX = (1 << (8 * (sys.maxsize.bit_length() + 1))) - 1 if False else (2**64 - 1)
def align_up(n, a):
return (n + a - 1) // a * a
for align in [2, 8, 16, 64, 256, 4096, 1 << 20]:
safe_max = SIZE_MAX - (align - 1)
# Check whether any n <= SIZE_MAX can round to > SIZE_MAX.
bad = []
for n in [SIZE_MAX, SIZE_MAX-1, safe_max, safe_max+1 if safe_max < SIZE_MAX else safe_max]:
try:
rounded = align_up(n, align)
if rounded > SIZE_MAX:
bad.append((n, rounded))
except Exception as e:
bad.append((n, repr(e)))
print(f"align={align} safe_max={safe_max} bad={bad[:3]}")
PYRepository: rapidsai/rmm
Length of output: 1384
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the exact implementation around the reported lines and the allocation path it feeds.
for f in \
cpp/src/mr/detail/caching_memory_resource_impl.cpp \
cpp/include/rmm/mr/stream_ordered_memory_resource.hpp \
cpp/src/mr/stream_ordered_memory_resource.cpp \
cpp/include/rmm/mr/detail/align.hpp \
cpp/include/rmm/detail/align.hpp
do
if [[ -f "$f" ]]; then
echo "=== $f ==="
wc -l "$f"
cat -n "$f" | sed -n '1,260p'
echo
fi
doneRepository: rapidsai/rmm
Length of output: 10789
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the alignment helper and the stream-ordered allocation check sequence.
rg -n "std::size_t align_up|align_up\(" cpp/src cpp/include/rmm -g '!**/build/**'
echo '---'
FILE=cpp/src/aligned.cpp
if [[ -f "$FILE" ]]; then
cat -n "$FILE" | sed -n '1,120p'
fi
echo '---'
FILE=cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp
cat -n "$FILE" | sed -n '80,130p'Repository: rapidsai/rmm
Length of output: 6214
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Compare similar maximum-allocation contracts in other memory resources.
rg -n "get_maximum_allocation_size\(\) const" cpp/src cpp/include/rmm/mr/detail -g '!**/build/**' -A4 -B4Repository: rapidsai/rmm
Length of output: 7467
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find the exact align_up implementation used by the MR codepaths.
FILE=cpp/include/rmm/aligned.hpp
if [[ -f "$FILE" ]]; then
cat -n "$FILE" | sed -n '1,220p'
fi
echo '---'
FILE=cpp/src/aligned.cpp
if [[ -f "$FILE" ]]; then
cat -n "$FILE" | sed -n '1,120p'
fiRepository: rapidsai/rmm
Length of output: 5117
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the alignment helper and the stream-ordered allocation path.
FILE=cpp/include/rmm/aligned.hpp
if [[ -f "$FILE" ]]; then
cat -n "$FILE" | sed -n '1,220p'
fi
echo '---'
FILE=cpp/src/aligned.cpp
if [[ -f "$FILE" ]]; then
cat -n "$FILE" | sed -n '1,120p'
fi
echo '---'
FILE=cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp
cat -n "$FILE" | sed -n '88,116p'Repository: rapidsai/rmm
Length of output: 6463
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the get_block logic to understand the effect of a wrapped-to-zero size.
FILE=cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp
cat -n "$FILE" | sed -n '130,230p'Repository: rapidsai/rmm
Length of output: 4766
Cap requests before alignment. stream_ordered_memory_resource::allocate() rounds before checking get_maximum_allocation_size(), so std::numeric_limits<std::size_t>::max() can wrap rmm::align_up() for near-limit requests and let an oversized request slip through. Return the largest safely roundable size for CUDA_ALLOCATION_ALIGNMENT, or reject oversized requests before rounding.
🤖 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/caching_memory_resource_impl.cpp` around lines 80 - 82,
Update caching_memory_resource_impl::get_maximum_allocation_size() to return the
largest size that can be safely rounded up to CUDA_ALLOCATION_ALIGNMENT, rather
than std::numeric_limits<std::size_t>::max(). Ensure near-limit allocation
requests are capped or rejected before alignment can overflow.
1 participant
Description
This is a draft, not ready for review.
I am experimenting with a memory allocator inspired by PyTorch's caching allocator. I plan to compare its performance over a range of workloads, fragmentation behavior, and OOM resilience against the existing RMM pool and arena memory resources.
Checklist