Skip to content

Comments

Refactor pool_memory_resource to shared CCCL MR design#2258

Draft
bdice wants to merge 2 commits intorapidsai:stagingfrom
bdice:refactor-pool-memory-resource
Draft

Refactor pool_memory_resource to shared CCCL MR design#2258
bdice wants to merge 2 commits intorapidsai:stagingfrom
bdice:refactor-pool-memory-resource

Conversation

@bdice
Copy link
Collaborator

@bdice bdice commented Feb 18, 2026

⚠️ Note to reviewers: I plan to self-review this PR before I ask reviewers to look at it. Just wanted to get this PR up before ending for the day.

Summary

Convert pool_memory_resource from a templated, header-only class to a non-templated class using cuda::mr::shared_resource for reference-counted ownership. Implementation is now compiled into librmm.so.

This follows the same pattern established in #2246. Part of #2011.

Key changes

  • Remove template parameter (was pool_memory_resource<Upstream>)
  • Add pool_memory_resource_impl in rmm::mr::detail namespace
  • Remove Upstream* constructor overloads (use device_async_resource_ref)
  • Store upstream as any_resource for owning semantics
  • Use identity-based equality for the impl (distinct pools are never interchangeable)
  • Class is now copyable with shared ownership semantics
  • Temporarily keep device_memory_resource inheritance for Python compatibility
  • Remove POOL_MR_REF_TEST and PINNED_POOL_MR_REF_TEST (replaced by dedicated fixture-based tests)
  • Add dedicated PoolMemoryResourceTest fixture with tests for GetUpstreamResource, AllocateDeallocate, SharedOwnership, Equality, PoolSize

Breaking API changes

  1. Template removal: pool_memory_resource<Upstream>pool_memory_resource
  2. Constructor changes: Removed Upstream* overloads, use device_async_resource_ref
  3. Equality semantics: Now identity-based (distinct pools are never equal, copies are equal)

Future changes

Once we migrate everything from using device_memory_resource to the CCCL MR design, we can drop the pieces of code labeled as legacy device_memory_resource compatibility layer.

@copy-pr-bot
Copy link

copy-pr-bot bot commented Feb 18, 2026

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@bdice bdice marked this pull request as draft February 18, 2026 05:27
Extract pool_memory_resource_impl into detail header and source files,
rewrite pool_memory_resource as a non-template outer class with
private shared_resource inheritance. Update Cython bindings, tests,
and doc comment examples. Remove pool/pinned_pool mr_ref tests
(replaced by dedicated fixture-based tests in pool_mr_tests.cpp).

81/81 C++ tests pass, 1163/1163 Python tests pass.
@bdice bdice force-pushed the refactor-pool-memory-resource branch from 7f067a6 to 64208e0 Compare February 20, 2026 21:10
@copy-pr-bot
Copy link

copy-pr-bot bot commented Feb 20, 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