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.

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

Labels

None yet

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

1 participant