Dry Run Protocol - #2961
Merged
Merged
Dry Run Protocol#2961
Conversation
…mory Introduce a dry-run execution framework that replaces device and host memory resources with lightweight fake allocators to measure peak memory usage without holding real memory. New files: - dry_run_memory_resource.hpp: dry_run_allocator (lock-free bump allocator), dry_run_device_memory_resource, dry_run_host_memory_resource, dry_run_resource_manager (RAII), and dry_run_execute() helper. - dry_run_flag.hpp: boolean dry-run flag as a raft resource, allowing algorithms to skip kernel execution during profiling. - tests/util/dry_run_memory_resource.cpp: unit tests. The dry_run_allocator probes the upstream once to obtain a base address, then atomically bumps a pointer for each allocation — no mutex, no map, no real memory held after the initial probe.
…pinned_memory_resource Add pinned and managed resources to the raft::resources handle to make it possible to customize / temporarily replace these resources
…aking change due to transitive includes in downstream libraries
Merges Remove deprecated headers (NVIDIA#2939). Conflict resolutions: - rsvd.cuh: Use new mdspan-based raft::matrix::sqrt and reciprocal APIs (they have internal dry-run guards); kept cudaMemsetAsync guard - svd.cuh: Use raft::matrix::weighted_sqrt (has internal dry-run guard) - matrix.cuh: Accept deletion (deprecated, removed in main) Co-authored-by: Cursor <cursoragent@cursor.com>
achirkin
requested review from
cjnolet,
divyegala,
huuanhhuyn,
lowener and
tarang-jain
August 25, 2026 09:21
lowener
reviewed
Aug 28, 2026
Comment on lines
+43
to
+44
| int getUniquelabels( | ||
| bool dry_run, rmm::device_uvector<value_t>& unique, value_t* y, size_t n, cudaStream_t stream) |
Contributor
There was a problem hiding this comment.
Why add the dry_run as a first argument, and not at the end, defaulted to false?
This would spare having a second version of that function just below
Contributor
Author
There was a problem hiding this comment.
This is to be very explicit in the detail namespace. At some point we should gradually refactor raft to always pass raft::resources handle everywhere and remove the need to have this argument at all.
achirkin
dismissed
huuanhhuyn’s stale review
August 28, 2026 12:54
The review comments were addressed, but the reviewer is out of the office to approve the changes.
bdice
approved these changes
Aug 28, 2026
bdice
left a comment
bdice
left a comment
Contributor
There was a problem hiding this comment.
One request for CMake, no comments on CI. Approved.
Contributor
Author
|
/merge |
alexfallin
added a commit
to alexfallin/raft_mst_imp
that referenced
this pull request
Sep 10, 2026
… mst() overload - Merge upstream main (dry run protocol NVIDIA#2961) - mst_solve: fold the stream argument into a handle copy; launch on the resources; replace thrust::sequence/cudaMemsetAsync with raft::linalg::map_offset; resource::sync_stream; dry-run guards - Replace local atomic and word load/store helpers with raft/util/device_atomics.cuh primitives and new raft::ldg_ca/stg_wb in raft/util/device_loads_stores.cuh - Gate RAFT_MST_HAS_CAS128 on host __int128; validate offsets[v] == e; shared mst_row_of helper - New handle-only mst() overload; deprecate the stream-taking overload - Tests: DryRunCompliance (exact dry-run accounting), DeprecatedStreamOverload
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The dry run protocol defines a mechanism to simulate the execution of algorithms to get a precise estimate of the memory requirements for a real execution with the same parameters.
This PR:
raft::util::dry_run_execute,raft::dry_run_resources, andresource::get_dry_run_flagto let callers estimate peak memory usage of any RAFT algorithm without executing GPU work.Note for reviewers
The PR contains a lot of small tedious changes to cover all of raft library and the tests components.
Please start reading at
docs/sourceupdates to learn more about the topic and the principles guiding these changes.