From 4be94a64bc0b6bda0203bf01060054e807816dd9 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Mon, 1 Jun 2026 16:21:52 +0000 Subject: [PATCH 1/3] Remove Cython any resource helper --- python/rmm/rmm/librmm/memory_resource.pxd | 21 +-------------- python/rmm/rmm/pylibrmm/device_buffer.pyx | 6 ++--- .../memory_resource/_memory_resource.pyx | 27 ++++++++++--------- 3 files changed, 18 insertions(+), 36 deletions(-) diff --git a/python/rmm/rmm/librmm/memory_resource.pxd b/python/rmm/rmm/librmm/memory_resource.pxd index eb9c60031..61de3c897 100644 --- a/python/rmm/rmm/librmm/memory_resource.pxd +++ b/python/rmm/rmm/librmm/memory_resource.pxd @@ -79,26 +79,7 @@ cdef extern from "" namespace "cuda::mr" nogil: pass cdef cppclass any_resource[Properties]: any_resource() except + - - -# Inline C++ helper to convert device_async_resource_ref to -# any_resource for passing to constructor parameters. -# Needed because Cython wraps intermediate values in __Pyx_FakeReference, -# a proxy with operator T&(). CCCL's any_resource only has a constrained -# forwarding constructor that cannot see through this proxy. A free function -# performs the conversion in pure C++ where it works directly. -# TODO: Remove once CCCL merges NVIDIA/cccl#8320 and RMM upgrades. -cdef extern from *: - """ - #include - #include - inline cuda::mr::any_resource - make_any_device_resource(rmm::device_async_resource_ref ref) { - return cuda::mr::any_resource(ref); - } - """ - any_resource[device_accessible] make_any_device_resource( - device_async_resource_ref) nogil except + + any_resource(device_async_resource_ref) except + # Inline C++ helper to construct optional[device_async_resource_ref] from any diff --git a/python/rmm/rmm/pylibrmm/device_buffer.pyx b/python/rmm/rmm/pylibrmm/device_buffer.pyx index 19b039a57..dc6ecb739 100644 --- a/python/rmm/rmm/pylibrmm/device_buffer.pyx +++ b/python/rmm/rmm/pylibrmm/device_buffer.pyx @@ -28,7 +28,7 @@ from rmm.librmm.device_buffer cimport ( get_current_cuda_device, prefetch, ) -from rmm.librmm.memory_resource cimport make_any_device_resource +from rmm.librmm.memory_resource cimport any_resource, device_accessible from rmm.pylibrmm.memory_resource cimport ( DeviceMemoryResource, get_current_device_resource, @@ -90,12 +90,12 @@ cdef class DeviceBuffer: if c_ptr == NULL or size == 0: self.c_obj.reset(new device_buffer( size, stream.view(), - make_any_device_resource(self.mr.get_mr()) + any_resource[device_accessible](self.mr.get_mr()) )) else: self.c_obj.reset(new device_buffer( c_ptr, size, stream.view(), - make_any_device_resource(self.mr.get_mr()) + any_resource[device_accessible](self.mr.get_mr()) )) if stream.c_is_default(): diff --git a/python/rmm/rmm/pylibrmm/memory_resource/_memory_resource.pyx b/python/rmm/rmm/pylibrmm/memory_resource/_memory_resource.pyx index 35da9b710..bcc808634 100644 --- a/python/rmm/rmm/pylibrmm/memory_resource/_memory_resource.pyx +++ b/python/rmm/rmm/pylibrmm/memory_resource/_memory_resource.pyx @@ -40,6 +40,7 @@ from rmm.librmm.memory_resource cimport ( CppExcept, allocate_callback_t, allocation_handle_type, + any_resource, arena_memory_resource, available_device_memory as c_available_device_memory, binning_memory_resource, @@ -48,13 +49,13 @@ from rmm.librmm.memory_resource cimport ( cuda_async_view_memory_resource, cuda_memory_resource, deallocate_callback_t, + device_accessible, device_async_resource_ref, failure_callback_resource_adaptor_oom, failure_callback_t, fixed_size_memory_resource, limiting_resource_adaptor, logging_resource_adaptor, - make_any_device_resource, make_device_async_resource_ref, managed_memory_resource, percent_of_free_device_memory as c_percent_of_free_device_memory, @@ -364,7 +365,7 @@ cdef class PoolMemoryResource(UpstreamResourceAdaptor): else optional[size_t]( parse_bytes(maximum_pool_size)) ) self.c_obj.reset(new pool_memory_resource( - make_any_device_resource(upstream_mr.get_mr()), + any_resource[device_accessible](upstream_mr.get_mr()), c_initial_pool_size, c_maximum_pool_size )) @@ -408,7 +409,7 @@ cdef class ArenaMemoryResource(UpstreamResourceAdaptor): else optional[size_t]( parse_bytes(arena_size)) ) self.c_obj.reset(new arena_memory_resource( - make_any_device_resource(upstream_mr.get_mr()), + any_resource[device_accessible](upstream_mr.get_mr()), c_arena_size, dump_log_on_failure, )) @@ -445,7 +446,7 @@ cdef class FixedSizeMemoryResource(UpstreamResourceAdaptor): size_t blocks_to_preallocate=128 ): self.c_obj.reset(new fixed_size_memory_resource( - make_any_device_resource(upstream_mr.get_mr()), + any_resource[device_accessible](upstream_mr.get_mr()), block_size, blocks_to_preallocate )) @@ -490,11 +491,11 @@ cdef class BinningMemoryResource(UpstreamResourceAdaptor): if (min_size_exponent == -1 or max_size_exponent == -1): self.c_obj.reset(new binning_memory_resource( - make_any_device_resource(upstream_mr.get_mr()) + any_resource[device_accessible](upstream_mr.get_mr()) )) else: self.c_obj.reset(new binning_memory_resource( - make_any_device_resource(upstream_mr.get_mr()), + any_resource[device_accessible](upstream_mr.get_mr()), min_size_exponent, max_size_exponent )) @@ -684,7 +685,7 @@ cdef class LimitingResourceAdaptor(UpstreamResourceAdaptor): size_t allocation_limit ): self.c_obj.reset(new limiting_resource_adaptor( - make_any_device_resource(upstream_mr.get_mr()), + any_resource[device_accessible](upstream_mr.get_mr()), allocation_limit )) self.c_ref = make_device_async_resource_ref(deref(self.c_obj)) @@ -749,7 +750,7 @@ cdef class LoggingResourceAdaptor(UpstreamResourceAdaptor): self._log_file_name = log_file_name self.c_obj.reset(new logging_resource_adaptor( - make_any_device_resource(upstream_mr.get_mr()), + any_resource[device_accessible](upstream_mr.get_mr()), log_file_name.encode() )) self.c_ref = make_device_async_resource_ref(deref(self.c_obj)) @@ -785,7 +786,7 @@ cdef class StatisticsResourceAdaptor(UpstreamResourceAdaptor): DeviceMemoryResource upstream_mr ): self.c_obj.reset(new statistics_resource_adaptor( - make_any_device_resource(upstream_mr.get_mr()) + any_resource[device_accessible](upstream_mr.get_mr()) )) self.c_ref = make_device_async_resource_ref(deref(self.c_obj)) @@ -876,7 +877,7 @@ cdef class TrackingResourceAdaptor(UpstreamResourceAdaptor): bool capture_stacks=False ): self.c_obj.reset(new tracking_resource_adaptor( - make_any_device_resource(upstream_mr.get_mr()), + any_resource[device_accessible](upstream_mr.get_mr()), capture_stacks )) self.c_ref = make_device_async_resource_ref(deref(self.c_obj)) @@ -949,7 +950,7 @@ cdef class FailureCallbackResourceAdaptor(UpstreamResourceAdaptor): ): self._callback = callback self.c_obj.reset(new failure_callback_resource_adaptor_oom( - make_any_device_resource(upstream_mr.get_mr()), + any_resource[device_accessible](upstream_mr.get_mr()), (_oom_callback_function), (callback) )) @@ -979,7 +980,7 @@ cdef class PrefetchResourceAdaptor(UpstreamResourceAdaptor): DeviceMemoryResource upstream_mr ): self.c_obj.reset(new prefetch_resource_adaptor( - make_any_device_resource(upstream_mr.get_mr()) + any_resource[device_accessible](upstream_mr.get_mr()) )) self.c_ref = make_device_async_resource_ref(deref(self.c_obj)) @@ -1110,7 +1111,7 @@ cpdef set_per_device_resource(int device, DeviceMemoryResource mr): cpp_set_per_device_resource( deref(device_id), - make_any_device_resource(mr.get_mr()) + any_resource[device_accessible](mr.get_mr()) ) From f5703f185a259350ae3f1c03cf5aac712575dd49 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Wed, 3 Jun 2026 04:27:07 +0000 Subject: [PATCH 2/3] Keep Cython any resource helper --- python/rmm/rmm/librmm/memory_resource.pxd | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/python/rmm/rmm/librmm/memory_resource.pxd b/python/rmm/rmm/librmm/memory_resource.pxd index 61de3c897..350d4ec8a 100644 --- a/python/rmm/rmm/librmm/memory_resource.pxd +++ b/python/rmm/rmm/librmm/memory_resource.pxd @@ -82,6 +82,26 @@ cdef extern from "" namespace "cuda::mr" nogil: any_resource(device_async_resource_ref) except + +# Inline C++ helper to convert device_async_resource_ref to +# any_resource for passing to constructor parameters. +# Needed because Cython wraps intermediate values in __Pyx_FakeReference, +# a proxy with operator T&(). CCCL's any_resource only has a constrained +# forwarding constructor that cannot see through this proxy. A free function +# performs the conversion in pure C++ where it works directly. +# TODO: Remove once CCCL merges NVIDIA/cccl#8320 and RMM upgrades. +cdef extern from *: + """ + #include + #include + inline cuda::mr::any_resource + make_any_device_resource(rmm::device_async_resource_ref ref) { + return cuda::mr::any_resource(ref); + } + """ + any_resource[device_accessible] make_any_device_resource( + device_async_resource_ref) nogil except + + + # Inline C++ helper to construct optional[device_async_resource_ref] from any # concrete resource type. Returns optional so that Cython assignment # (self.c_ref = make_device_async_resource_ref(...)) uses optional's From 6bc91d604c229933ec07bb9bb75524c62c5f0674 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Wed, 3 Jun 2026 04:36:06 +0000 Subject: [PATCH 3/3] Remove Cython any resource helper --- python/rmm/rmm/librmm/memory_resource.pxd | 21 --------------------- 1 file changed, 21 deletions(-) diff --git a/python/rmm/rmm/librmm/memory_resource.pxd b/python/rmm/rmm/librmm/memory_resource.pxd index 350d4ec8a..80684a7c3 100644 --- a/python/rmm/rmm/librmm/memory_resource.pxd +++ b/python/rmm/rmm/librmm/memory_resource.pxd @@ -81,27 +81,6 @@ cdef extern from "" namespace "cuda::mr" nogil: any_resource() except + any_resource(device_async_resource_ref) except + - -# Inline C++ helper to convert device_async_resource_ref to -# any_resource for passing to constructor parameters. -# Needed because Cython wraps intermediate values in __Pyx_FakeReference, -# a proxy with operator T&(). CCCL's any_resource only has a constrained -# forwarding constructor that cannot see through this proxy. A free function -# performs the conversion in pure C++ where it works directly. -# TODO: Remove once CCCL merges NVIDIA/cccl#8320 and RMM upgrades. -cdef extern from *: - """ - #include - #include - inline cuda::mr::any_resource - make_any_device_resource(rmm::device_async_resource_ref ref) { - return cuda::mr::any_resource(ref); - } - """ - any_resource[device_accessible] make_any_device_resource( - device_async_resource_ref) nogil except + - - # Inline C++ helper to construct optional[device_async_resource_ref] from any # concrete resource type. Returns optional so that Cython assignment # (self.c_ref = make_device_async_resource_ref(...)) uses optional's