Migrate RMM usage to CCCL MR design by bdice · Pull Request #1990 · rapidsai/cuvs

bdice · 2026-04-03T19:01:33Z

Summary

Migrate all RMM usage to the new CCCL memory resource design (de-templated resources, device_async_resource_ref instead of device_memory_resource*, value semantics)
Replace get_workspace_resource() / get_large_workspace_resource() with _ref() variants across 65 call sites
Rewrite cuda_huge_page_resource to satisfy CCCL resource concept directly
Remove owning_wrapper / dynamic_cast patterns in C API and benchmarks

Depends on rapidsai/rmm#2361.
Depends on rapidsai/ucxx#636.
Depends on rapidsai/raft#2996.

Changes

33 files changed (~208 insertions, ~221 deletions)
device_memory_resource* params → device_async_resource_ref (ivf_common, ivf_pq, naive_knn)
get_current_device_resource() → get_current_device_resource_ref()
set_current_device_resource() → set_current_device_resource_ref()
De-templated pool_memory_resource, failure_callback_resource_adaptor in bench utils
Removed &resource pointer patterns (resources are now copyable value types)
Removed spurious mr arg from select_k calls (previously compiled due to implicit pointer→bool conversion)
C API pool resource management rewritten without owning_wrapper

Adapt cuVS to RMM breaking changes: removal of device_memory_resource base class, de-templated resource/adaptor types, new per-device resource ref APIs, and CCCL resource concept requirements. Key changes: - get_workspace_resource() -> get_workspace_resource_ref() (44 sites) - get_large_workspace_resource() -> get_large_workspace_resource_ref() (21 sites) - get_current_device_resource() -> get_current_device_resource_ref() - device_memory_resource* params -> device_async_resource_ref - Remove &resource pointer patterns (resources are now value types) - Migrate cuda_huge_page_resource to CCCL concept - De-template pool_memory_resource, failure_callback_resource_adaptor - Rewrite C API pool resource management without owning_wrapper - Remove deleted rmm/mr/device_memory_resource.hpp includes

copy-pr-bot · 2026-04-03T19:01:38Z

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.

…evice_resource

# Conflicts: # cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh

bdice · 2026-04-19T07:03:41Z

/ok to test

bdice · 2026-04-19T07:16:52Z

@@ -17,37 +19,25 @@

 namespace raft::mr {


Why is cuVS defining classes in the raft:: namespace? (This is not in scope for this PR.)

bdice · 2026-04-19T07:23:38Z

                                          indices_topk + offset * k,
-                                          type != cuvs::distance::DistanceType::InnerProduct,
-                                          mr);
+                                          type != cuvs::distance::DistanceType::InnerProduct);


Need to check if this change is intentional or a mistake.

Seems like this was changed in rapidsai/raft#1786?

Maybe this was a dead code path. I can't see how it would've been building successfully prior to this migration otherwise.

bdice · 2026-04-19T07:26:50Z

                                                indices_topk + offset * k,
-                                                cuvs::distance::is_min_close(type),
-                                                mr);
+                                                cuvs::distance::is_min_close(type));


Need to check if this API change was intentional.

Seems like this was changed in rapidsai/raft#1786?

Maybe this was a dead code path. I can't see how it would've been building successfully prior to this migration otherwise.

…e owns the resource

…its upstream

…e_resource owns it

Add explicit operator!= (required when _CCCL_HAS_CONCEPTS is 0, e.g. nvcc C++17 mode, since operator!= is not auto-synthesized from operator==). Use constexpr friend for get_property to match the RMM resource pattern. Add synchronous_resource and resource static_asserts for finer diagnostics.

bdice · 2026-04-21T18:10:00Z

Requesting admin-merge. Most CI is passing, the one failure seems like a fluke. I am not going to wait for the slow L4 C++ job since all others have passed and this is blocking builds for cuML + cuGraph.

## Summary - Migrate all raw RMM `allocate`/`deallocate` calls to the new CCCL 3-argument API that requires explicit alignment - Replace removed `rmm.librmm.per_device_resource` Cython import with `rmm.pylibrmm.memory_resource` and use `make_any_device_resource` to obtain the resource for `device_buffer` construction Depends on rapidsai/rmm#2361. Depends on rapidsai/ucxx#636. Depends on rapidsai/raft#2996. Depends on rapidsai/cuvs#1990. ## Changes - **`cpp/src/genetic/genetic.cu`**: Add explicit `alignof(node)` / `alignof(program)` to all `allocate` and `deallocate` calls in `parallel_evolve` and `symFit`; fix deallocation bug in `parallel_evolve` where `h_nextprogs[i].len` was incorrectly used instead of `tmp.len` to compute the buffer size being freed - **`cpp/examples/symreg/symreg_example.cpp`**: Use `params.population_size * sizeof(cg::program)` and `alignof(cg::program)` for `allocate`/`deallocate` calls, fixing incorrect byte-size computation; remove unused `<rmm/aligned.hpp>` include - **`cpp/tests/sg/genetic/evolution_test.cu`**: Add alignment arguments to allocate/deallocate in `SymReg` test - **`cpp/tests/sg/genetic/program_test.cu`**: Add alignment arguments to `SetUp`/`TearDown` allocate/deallocate calls - **`python/cuml/cuml/manifold/umap/umap.pyx`**: Replace `get_current_device_resource()` with `make_any_device_resource(get_current_device_resource().get_mr())` for `device_buffer` construction Authors: - Bradley Dice (https://github.com/bdice) Approvers: - Simon Adorf (https://github.com/csadorf) - Divye Gala (https://github.com/divyegala) - Victor Lafargue (https://github.com/viclafargue) URL: #7951

## Summary - Replace removed `rmm::mr::device_memory_resource` base class, `owning_wrapper`, `shared_ptr`-based resource management, and deprecated per-device resource APIs with CCCL-native memory resource types - Use `cuda::mr::any_resource<cuda::mr::device_accessible>` for owning type-erased storage, `rmm::device_async_resource_ref` for non-owning references, and value-typed resources (`cuda_memory_resource`, `pinned_host_memory_resource`) - Pass the memory resource to `raft::handle_t` as the `workspace_resource` (3rd) constructor argument, matching the new raft API (`stream_view`, `stream_pool`, `std::optional<raft::mr::device_resource>`) Depends on rapidsai/rmm#2361. Depends on rapidsai/ucxx#636. Depends on rapidsai/raft#2996. Depends on rapidsai/cuvs#1990. ## Files changed **Headers:** - `algorithms.hpp`, `dendrogram.hpp`, `legacy/graph.hpp`, `legacy/functions.hpp`: `get_current_device_resource()` → `get_current_device_resource_ref()` in default argument expressions - `host_staging_buffer_manager.hpp`: Remove `owning_wrapper`, store `pool_memory_resource` by value in a `std::optional`, accept `pinned_host_memory_resource` by value in `init()` - `large_buffer_manager.hpp`: Store `pinned_host_memory_resource` by value (not `shared_ptr`), return `device_async_resource_ref` from `get()`, `std::move` the resource into storage - `mtmg/resource_manager.hpp`: Use `cuda::mr::any_resource<device_accessible>` instead of `shared_ptr<device_memory_resource>` for `per_device_rmm_resources_`, use non-deprecated `set_per_device_resource`, pass resource as `workspace_resource` to `raft::handle_t` **Tests:** - `base_fixture.hpp`: Return `any_resource<device_accessible>` from `create_memory_resource()`, use value-typed MR factory helpers (`make_cuda`, `make_managed`, `make_pool`, `make_binning`), switch to non-deprecated `set_current_device_resource` / `get_current_device_resource_ref` - `multi_node_threaded_test.cpp`: Switch to non-deprecated `set_current_device_resource(resource)` - `mg_graph500_bfs_test.cu`, `mg_graph500_sssp_test.cu`: Store `pinned_mr_` as `optional<pinned_host_memory_resource>` by value, prefer `.value()` over `operator*` for optional access **Examples:** - All 4 example files (`sg_graph_algorithms.cpp`, `mg_graph_algorithms.cpp`, `vertex_and_edge_partition.cu`, `graph_operations.cu`): Use value-typed `cuda_memory_resource`, non-deprecated `set_current_device_resource`, pass the resource to `raft::handle_t` as the `workspace_resource` (3rd positional arg, with `nullptr` for the unused `stream_pool`) Authors: - Bradley Dice (https://github.com/bdice) - Chuck Hastings (https://github.com/ChuckHastings) Approvers: - Chuck Hastings (https://github.com/ChuckHastings) - Vyas Ramasubramani (https://github.com/vyasr) URL: #5483

github-project-automation Bot added this to Unstructured Data Processing Apr 3, 2026

bdice mentioned this pull request Apr 3, 2026

[FEA] Support memory resources from CCCL 3.2 rapidsai/rmm#2011

Open

50 tasks

aamijar assigned bdice Apr 6, 2026

bdice added 4 commits April 14, 2026 19:17

Merge remote-tracking branch 'upstream/main' into rmm-cccl-migration

1e4936f

Fix device_uvector construction for new RMM resource model

c7f2a8b

Replace deprecated set_current_device_resource_ref with set_current_d…

ae767d4

…evice_resource

Merge remote-tracking branch 'upstream/main' into rmm-cccl-migration

7d00c6f

# Conflicts: # cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh

bdice force-pushed the rmm-cccl-migration branch 4 times, most recently from 942f4e8 to 4c59082 Compare April 17, 2026 07:09

bdice added breaking Introduces a breaking change improvement Improves an existing functionality labels Apr 17, 2026

bdice force-pushed the rmm-cccl-migration branch from 4c59082 to ae4f344 Compare April 17, 2026 07:29

bdice changed the title ~~Migrate RMM usage to CCCL memory resource design~~ Migrate RMM usage to CCCL MR design Apr 17, 2026

bdice force-pushed the rmm-cccl-migration branch 5 times, most recently from 4df8847 to e45a83e Compare April 18, 2026 15:52

bdice commented Apr 19, 2026

View reviewed changes

bdice added 3 commits April 19, 2026 10:53

Migrate examples to RMM 26.06 CCCL memory resource API

f20ea87

Update docs RMM example for 26.06 API

3dd1b15

Add FAISS patch for RMM 26.06 device_memory_resource removal

358b72e

bdice force-pushed the rmm-cccl-migration branch from 1428e7c to 423fa1c Compare April 19, 2026 15:54

bdice added 2 commits April 19, 2026 11:26

Remove thread_local pool state from C API; set_current_device_resourc…

0db0b3f

…e owns the resource

Add host_accessible property to cuda_huge_page_resource

b88d0ba

bdice added 3 commits April 19, 2026 11:26

Remove pool_resource_ member; failure_callback_resource_adaptor owns …

4ab4ef4

…its upstream

Remove resource_ member from shared_raft_resources; set_current_devic…

17ce188

…e_resource owns it

Style

029cf1f

bdice force-pushed the rmm-cccl-migration branch from 423fa1c to 8f50dc0 Compare April 19, 2026 16:27

bdice marked this pull request as ready for review April 19, 2026 16:47

bdice requested review from a team as code owners April 19, 2026 16:47

bdice requested a review from AyodeAwe April 19, 2026 16:47

bdice force-pushed the rmm-cccl-migration branch from 8f50dc0 to 06e0ae5 Compare April 19, 2026 19:15

This was referenced Apr 20, 2026

Migrate RMM usage to CCCL MR design rapidsai/cuml#7951

Merged

Migrate RMM usage to CCCL MR design rapidsai/cugraph#5483

Merged

divyegala approved these changes Apr 20, 2026

View reviewed changes

bdice mentioned this pull request Apr 21, 2026

CCCL Memory Resource Migration — Merge Train rapidsai/rmm#2364

Closed

bdice force-pushed the rmm-cccl-migration branch from 06e0ae5 to fe7ca1b Compare April 21, 2026 11:57

Fix clang-format

f1230a4

bdice removed the request for review from a team April 21, 2026 14:27

gforsyth merged commit c999208 into rapidsai:main Apr 21, 2026
77 of 78 checks passed

github-project-automation Bot moved this to Done in Unstructured Data Processing Apr 21, 2026

achirkin mentioned this pull request Apr 22, 2026

[REVIEW] Generalize and improve cagra::optimize #1830

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Migrate RMM usage to CCCL MR design#1990

Migrate RMM usage to CCCL MR design#1990
gforsyth merged 15 commits intorapidsai:mainfrom
bdice:rmm-cccl-migration

bdice commented Apr 3, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Apr 3, 2026

Uh oh!

bdice commented Apr 19, 2026

Uh oh!

Uh oh!

bdice Apr 19, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

bdice Apr 19, 2026

Uh oh!

bdice Apr 19, 2026

Uh oh!

bdice Apr 21, 2026

Uh oh!

bdice Apr 19, 2026

Uh oh!

bdice Apr 19, 2026

Uh oh!

bdice Apr 21, 2026

Uh oh!

bdice commented Apr 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

bdice commented Apr 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Uh oh!

copy-pr-bot Bot commented Apr 3, 2026

Uh oh!

bdice commented Apr 19, 2026

Uh oh!

Uh oh!

bdice Apr 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

bdice Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

bdice Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

bdice Apr 21, 2026

Choose a reason for hiding this comment

Uh oh!

bdice Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

bdice Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

bdice Apr 21, 2026

Choose a reason for hiding this comment

Uh oh!

bdice commented Apr 21, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

bdice commented Apr 3, 2026 •

edited

Loading

bdice Apr 19, 2026 •

edited

Loading