Key arena per-stream state by CUDA stream ID - #2471
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR keys per-stream arena state by normalized CUDA stream IDs from ChangesArena stream ID keying
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related issues
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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: 1
🤖 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/src/mr/detail/arena_memory_resource_impl.cpp`:
- Around line 18-29: The get_stream_id helper currently uses
RMM_ASSERT_CUDA_SUCCESS around cudaStreamGetId, which can be compiled out in
release builds and silently leave stream_id at 0. Update get_stream_id in
arena_memory_resource_impl.cpp to use RMM_CUDA_TRY for the cudaStreamGetId call
so failures throw immediately, while keeping the existing stream_to_store
selection and stream_id_type handling unchanged.
🪄 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: 7abeee32-a042-4e92-9329-6234930bfe09
📒 Files selected for processing (3)
cpp/include/rmm/mr/detail/arena_memory_resource_impl.hppcpp/src/mr/detail/arena_memory_resource_impl.cppcpp/tests/mr/arena_mr_tests.cpp
fd6a2ed to
a215f77
Compare
bdice
left a comment
bdice
left a comment
There was a problem hiding this comment.
Thanks for this and the other PRs you filed! Seems mostly good, I have a couple questions.
| * `noexcept`-safe requirements of the `deallocate`/`deallocate_sync` path. | ||
| * @return The CUDA stream ID. | ||
| */ | ||
| stream_id_type get_stream_id(cuda_stream_view stream, bool may_throw); |
There was a problem hiding this comment.
Hmm. Do we need something like may_throw in the corresponding code for stream_ordered_memory_resource?
Is it possible to catch and swallow the exception in the deallocation path and still count that as noexcept?
There was a problem hiding this comment.
nit: Let's make this a method on cuda_stream_view, I think.
There was a problem hiding this comment.
@wence- We'll be migrating from rmm::cuda_stream_view to cuda::stream_ref in the relatively near future. I'd like to keep this a free function for now, though we can file an upstream feature request with CCCL for that if you like.
There was a problem hiding this comment.
ah ok, I withdraw my suggestion
|
I thought about this some more. A cleaner design would be to split this into For I think any |
| * `noexcept`-safe requirements of the `deallocate`/`deallocate_sync` path. | ||
| * @return The CUDA stream ID. | ||
| */ | ||
| stream_id_type get_stream_id(cuda_stream_view stream, bool may_throw); |
There was a problem hiding this comment.
nit: Let's make this a method on cuda_stream_view, I think.
| // Safe to throw: only called from the non-noexcept allocate() path. | ||
| RMM_CUDA_TRY(cudaStreamGetId(stream_to_store, &stream_id)); | ||
| } else { | ||
| // Called from the noexcept deallocate()/deallocate_sync() path: must not throw. A failure | ||
| // here is only surfaced via assertion in debug builds. | ||
| RMM_ASSERT_CUDA_SUCCESS(cudaStreamGetId(stream_to_store, &stream_id)); |
There was a problem hiding this comment.
If this were a method on cuda_stream_view we could plausibly query the ID on creation (if it is fast) and stash it on the stream view. That's a ctor that may throw, so we're all good. And then get_stream_id could itself be noexcept
There was a problem hiding this comment.
I like the idea of querying/caching on construction, we'll need to gather some benchmarks to understand the cost.
|
Thanks @bdice, the split makes sense to me. A throwing |
The arena memory resource keyed per-stream arenas by the raw cudaStream_t handle. CUDA can hand a destroyed stream's handle value to a new stream (ABA reuse), so a new logical stream could inherit stale arena state. Key stream_arenas_ by cudaStreamGetId instead, mirroring the pattern already used by stream_ordered_memory_resource. Closes rapidsai#2394 Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
a215f77 to
293338d
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/include/rmm/mr/detail/arena_memory_resource_impl.hpp`:
- Around line 126-130: Update the stream_arenas_ lifecycle in the arena memory
resource implementation to retire and erase arenas associated with destroyed or
otherwise inactive CUDA streams, while preserving safety for arenas still in use
and avoiding stream-ID reuse hazards. Ensure eviction releases arena-owned
resources and does not leave permanently retained entries for every stream ever
seen.
In `@cpp/tests/mr/arena_mr_tests.cpp`:
- Line 657: Replace the direct cudaStreamGetId assertion in the affected test
with RMM_CUDA_TRY, preserving the existing stream.view().value() argument and
expected_id output parameter.
🪄 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: 3c5143cc-9f5a-4b62-a206-65bcf0fd1b61
📒 Files selected for processing (3)
cpp/include/rmm/mr/detail/arena_memory_resource_impl.hppcpp/src/mr/detail/arena_memory_resource_impl.cppcpp/tests/mr/arena_mr_tests.cpp
Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
293338d to
afbf0a3
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (2)
cpp/tests/mr/arena_mr_tests.cpp (2)
658-660: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExercise the
try_get_stream_idfailure path.This test only covers successful lookup. Add a controlled invalid or destroyed-stream case and assert that
try_get_stream_idreturnsstd::nulloptwithout throwing, since deallocation relies on that contract.🤖 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/arena_mr_tests.cpp` around lines 658 - 660, Add a failure-path case alongside the existing successful `try_get_stream_id` test, using a controlled invalid or destroyed stream. Assert that the call returns std::nullopt and does not throw, while preserving the current expected-ID assertion for valid streams.
664-680: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert default-stream normalization in PTDS builds too.
The
#ifndef CUDA_API_PER_THREAD_DEFAULT_STREAMexcludes the default-versus-legacy assertion in PTDS builds, although the helper contract says normalization applies in both modes. Keep this assertion unconditional.Proposed adjustment
-#ifndef CUDA_API_PER_THREAD_DEFAULT_STREAM - // Non-PTDS: the default stream normalizes to cudaStreamLegacy, so it keys to the legacy id. + // The default stream normalizes to cudaStreamLegacy in both default-stream modes. EXPECT_EQ(rmm::mr::detail::get_stream_id(rmm::cuda_stream_view{}), rmm::mr::detail::get_stream_id(rmm::cuda_stream_legacy)); -#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/tests/mr/arena_mr_tests.cpp` around lines 664 - 680, Make the default-versus-legacy stream ID assertion in ArenaStreamIdTest unconditional by removing the CUDA_API_PER_THREAD_DEFAULT_STREAM guard, while preserving the existing real-stream collision check and other test assertions.
🤖 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.
Nitpick comments:
In `@cpp/tests/mr/arena_mr_tests.cpp`:
- Around line 658-660: Add a failure-path case alongside the existing successful
`try_get_stream_id` test, using a controlled invalid or destroyed stream. Assert
that the call returns std::nullopt and does not throw, while preserving the
current expected-ID assertion for valid streams.
- Around line 664-680: Make the default-versus-legacy stream ID assertion in
ArenaStreamIdTest unconditional by removing the
CUDA_API_PER_THREAD_DEFAULT_STREAM guard, while preserving the existing
real-stream collision check and other test assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 329e0086-c249-41a0-8848-9fabb7ccabe8
📒 Files selected for processing (4)
cpp/include/rmm/mr/arena_memory_resource.hppcpp/include/rmm/mr/detail/arena_memory_resource_impl.hppcpp/src/mr/detail/arena_memory_resource_impl.cppcpp/tests/mr/arena_mr_tests.cpp
🚧 Files skipped from review as they are similar to previous changes (2)
- cpp/include/rmm/mr/detail/arena_memory_resource_impl.hpp
- cpp/src/mr/detail/arena_memory_resource_impl.cpp
Switch the raw ASSERT_EQ(..., cudaSuccess) check to RMM_CUDA_TRY, the repo's CUDA error-checking macro, and include <rmm/detail/error.hpp> directly where the macro is defined rather than relying on a transitive include. Signed-off-by: nethum529 <nethumweerasinghe.nw@gmail.com>
|
Two notes on the latest CodeRabbit nitpicks that were in the review body rather than inline threads:
Default-stream normalization assertion under PTDS (removing the |
|
@nethum529 I am working through a few items for our 26.08 release and I'll keep working through reviewing and merging the PRs you filed in the next few days. I really appreciate the PRs you've filed, they are helpful contributions! |
Of course not a problem Bradley! This is cool stuff so I don't mind 😁 |
4 participants
Description
arena_memory_resource_implstored per-stream arena state in astd::mapkeyed by the rawcudaStream_thandle, looked up withstream.value().CUDA can reassign a destroyed stream's raw handle value to a newly created stream (ABA reuse), so a new logical stream could inherit stale arena state left behind by an earlier, unrelated stream that happened to share the same handle value.
This mirrors the fix already applied to
stream_ordered_memory_resourcein #2252, which keys stream state bycudaStreamGetIdfor exactly this reason.Changes:
rmm::mr::detail::get_stream_idhelper that normalizes the default stream tocudaStreamLegacyand returns the ID fromcudaStreamGetId.stream_arenas_bystream_id_type(the stream ID) instead of the rawcudaStream_thandle.get_stream_arenanow looks up and inserts by the stream ID.closes #2394
Before / After
stream_arenas_cudaStream_thandle (stream.value())cudaStreamGetId(stream)Tests
Added three tests to
cpp/tests/mr/arena_mr_tests.cpp:ARENA_MR_TESTandARENA_MR_PTDS_TESTboth pass 44/44.Checklist