Skip to content

Migrate RMM usage to CCCL MR design#2996

Draft
bdice wants to merge 1 commit intorapidsai:mainfrom
bdice:rmm-cccl-migration
Draft

Migrate RMM usage to CCCL MR design#2996
bdice wants to merge 1 commit intorapidsai:mainfrom
bdice:rmm-cccl-migration

Conversation

@bdice
Copy link
Copy Markdown
Contributor

@bdice bdice commented Apr 3, 2026

Summary

Changes

Core resource infrastructure

  • device_memory_resource.hpp: Remove any_resource_bridge (which inherited from rmm::mr::device_memory_resource), remove all shared_ptr<device_memory_resource> constructor overloads, consolidate to any_resource-only path
  • device_resources.hpp: Remove deprecated constructor taking shared_ptr<device_memory_resource>, update get_workspace_resource() return type (de-templated limiting_resource_adaptor)
  • device_resources_snmg.hpp: Remove stale include, de-template pool_memory_resource
  • handle.hpp: Remove deprecated constructors taking shared_ptr<device_memory_resource>
  • device_resources_manager.hpp: Remove workspace_mrs vector and set_workspace_memory_resource() API, de-template pool_mr_ to optional<pool_memory_resource>, remove dynamic_cast for upstream type detection, replace get/set_current_device_resource() with _ref variants

Memory tracking

  • memory_tracking_resources.hpp: Remove device_tracking_bridge (inherited from device_memory_resource), use set_current_device_resource_ref() directly

Call sites using get_workspace_resource()get_workspace_resource_ref()

  • select_k-inl.cuh, select_radix.cuh, select_warpsort.cuh, sparse/select_k-inl.cuh, bitmap_to_csr.cuh, bitset_to_csr.cuh

Benchmarks

  • benchmark.hpp: De-template pool_memory_resource, use any_resource for RAII restore
  • gather.cu, subsample.cu: Same pattern

Tests

  • handle.cpp: Dereference limiting_resource_adaptor* for device_buffer constructor
  • device_resources_manager.cpp: Remove workspace-related test code for removed APIs
  • mdarray.cu: Remove test_device_resource_bridge_unwrap (bridge no longer exists)
  • multi_variable_gaussian.cu: get_current_device_resource()get_current_device_resource_ref()

Remove device_memory_resource base class usage, de-template all resource
and adaptor types, replace pointer-based per-device resource APIs with
ref-based equivalents, and update all call sites for the new signatures.

Part of rapidsai/rmm#2011.
@copy-pr-bot
Copy link
Copy Markdown

copy-pr-bot bot commented Apr 3, 2026

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.

@achirkin
Copy link
Copy Markdown
Contributor

achirkin commented Apr 8, 2026

I assume this can only go in after downstream library PRs are merged (cuvs/cuml/cugraph), right?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

breaking Breaking change improvement Improvement / enhancement to an existing function

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

2 participants