Skip to content

Update RMM memory resource APIs to ref-based equivalents#2920

Open
bdice wants to merge 4 commits intorapidsai:mainfrom
bdice:rmm-resource-ref-migration
Open

Update RMM memory resource APIs to ref-based equivalents#2920
bdice wants to merge 4 commits intorapidsai:mainfrom
bdice:rmm-resource-ref-migration

Conversation

@bdice
Copy link
Contributor

@bdice bdice commented Jan 13, 2026

Summary

  • Replace pointer-based RMM APIs (which will soon be deprecated) with ref-based alternatives as part of the RMM CCCL 3.2 memory resource migration
  • Update device_resources_manager.hpp: set_current_device_resource()set_current_device_resource_ref()
  • Update benchmark.hpp: Update using_pool_memory_res to use ref-based APIs
  • Update gather.cu: Use ref-based APIs for memory resource management
  • Update subsample.cu: Use ref-based APIs for memory resource management
  • Update buffer.cpp: set_current_device_resource()set_current_device_resource_ref()

Replace pointer-based RMM APIs (which will soon be deprecated) with
ref-based alternatives as part of the RMM CCCL 3.2 memory resource migration:

- device_resources_manager.hpp: set_current_device_resource() → set_current_device_resource_ref()
- benchmark.hpp: Update using_pool_memory_res to use ref-based APIs
- gather.cu: Update to use ref-based APIs for memory resource management
- subsample.cu: Update to use ref-based APIs for memory resource management
- buffer.cpp: set_current_device_resource() → set_current_device_resource_ref()
bdice added 2 commits January 13, 2026 22:04
Change mem_resource_ from device_async_resource_ref to
cuda::mr::any_resource to ensure the memory resource lifetime
is properly managed. This is part of the RMM CCCL 3.2 migration.
@csadorf
Copy link
Contributor

csadorf commented Jan 15, 2026

This seems like a change with potential for unexpected side-effects. I'd urge to merge this asap if it cannot be delayed to 26.04.

@dantegd
Copy link
Member

dantegd commented Jan 20, 2026

The changes look good to me, there is one CI failure that looks fairly minor:

[ RUN      ] DeviceResourcesManager.ObeysSetters
/tmp/conda-bld-output/bld/rattler-build_libraft-headers-only/work/cpp/tests/core/device_resources_manager.cpp:111: Failure
Expected: (mr) != (nullptr), actual: NULL vs (nullptr)

/tmp/conda-bld-output/bld/rattler-build_libraft-headers-only/work/cpp/tests/core/device_resources_manager.cpp:111: Failure
Expected: (mr) != (nullptr), actual: NULL vs (nullptr)

but I agree with @csadorf that I'd rather be conservative and aim this to main/26.04, @bdice would that be ok?

@bdice
Copy link
Contributor Author

bdice commented Jan 20, 2026

This should target 26.04, I’ll update you after I’ve had time to figure out the errors in this and other repositories.

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

Labels

None yet

Development

Successfully merging this pull request may close these issues.

3 participants