Fix Buffer.from_handle(mr=...) not calling mr.deallocate()

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -565,6 +565,42 @@ DevicePtrHandle deviceptr_create_with_owner(CUdeviceptr ptr, PyObject* owner) {
  return DevicePtrHandle(box, &box->resource);
 }
 
+// ============================================================================
+// MemoryResource-owned Device Pointer Handles
+// ============================================================================
+
+static MRDeallocCallback mr_dealloc_cb = nullptr;
+
+void register_mr_dealloc_callback(MRDeallocCallback cb) {
+ mr_dealloc_cb = cb;
+}
+
+DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* mr) {
+ if (!mr) {
+ return deviceptr_create_ref(ptr);
+ }
+ // GIL required when mr is provided
+ GILAcquireGuard gil;
+ if (!gil.acquired()) {
+ return deviceptr_create_ref(ptr);
+ }
+ Py_INCREF(mr);
+ auto box = std::shared_ptr<DevicePtrBox>(
+ new DevicePtrBox{ptr, StreamHandle{}},
+ [mr, size](DevicePtrBox* b) {
+ GILAcquireGuard gil;
+ if (gil.acquired()) {
+ if (mr_dealloc_cb) {
+ mr_dealloc_cb(mr, b->resource, size, b->h_stream);
+ }
+ Py_DECREF(mr);
+ }
+ delete b;
+ }
+ );
+ return DevicePtrHandle(box, &box->resource);
+}
+
 // ============================================================================
 // IPC Pointer Cache
 // ============================================================================

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -239,6 +239,23 @@ DevicePtrHandle deviceptr_create_ref(CUdeviceptr ptr);
 // If owner is nullptr, equivalent to deviceptr_create_ref.
 DevicePtrHandle deviceptr_create_with_owner(CUdeviceptr ptr, PyObject* owner);
 
+// Callback type for MemoryResource deallocation.
+// Called from the shared_ptr deleter when a handle created via
+// deviceptr_create_with_mr is destroyed. The implementation is responsible
+// for converting raw C types to Python objects and calling
+// mr.deallocate(ptr, size, stream).
+using MRDeallocCallback = void (*)(PyObject* mr, CUdeviceptr ptr,
+ size_t size, const StreamHandle& stream);
+
+// Register the MR deallocation callback.
+void register_mr_dealloc_callback(MRDeallocCallback cb);
+
+// Create a device pointer handle whose destructor calls mr.deallocate()
+// via the registered callback. The mr's refcount is incremented and
+// decremented when the handle is released.
+// If mr is nullptr, equivalent to deviceptr_create_ref.
+DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* mr);
+
 // Import a device pointer from IPC via cuMemPoolImportPointer.
 // When the last reference is released, cuMemFreeAsync is called on the stored stream.
 // Note: Does not yet implement reference counting for nvbug 5570902.

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -17,12 +17,14 @@ from cuda.core._resource_handles cimport (
  DevicePtrHandle,
  StreamHandle,
  deviceptr_create_with_owner,
+ deviceptr_create_with_mr,
+ register_mr_dealloc_callback,
  as_intptr,
  as_cu,
  set_deallocation_stream,
 )
 
-from cuda.core._stream cimport Stream_accept, Stream
+from cuda.core._stream cimport Stream, Stream_accept
 from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
 
 import sys
@@ -37,6 +39,30 @@ from cuda.core._dlpack import DLDeviceType, make_py_capsule
 from cuda.core._utils.cuda_utils import driver
 from cuda.core._device import Device
 
+
+# =============================================================================
+# MR deallocation callback (invoked from C++ shared_ptr deleter)
+# =============================================================================
+
+cdef void _mr_dealloc_callback(
+ object mr,
+ cydriver.CUdeviceptr ptr,
+ size_t size,
+ const StreamHandle& h_stream,
+) noexcept:
+ """Called by the C++ deleter to deallocate via MemoryResource.deallocate."""
+ try:
+ stream = None
+ if h_stream:
+ stream = Stream._from_handle(Stream, h_stream)
+ mr.deallocate(int(ptr), size, stream)
+ except Exception as exc:
+ print(f"Warning: mr.deallocate() failed during Buffer destruction: {exc}",
+ file=sys.stderr)
+
+register_mr_dealloc_callback(_mr_dealloc_callback)
+
+
 __all__ = ['Buffer', 'MemoryResource']
 
 
@@ -73,27 +99,30 @@ cdef class Buffer:
  @classmethod
  def _init(
  cls, ptr: DevicePointerT, size_t size, mr: MemoryResource | None = None,
-stream: Stream | None = None, ipc_descriptor: IPCBufferDescriptor | None = None,
+ ipc_descriptor: IPCBufferDescriptor | None = None,
  owner : object | None = None
  ):
- """Legacy init for compatibility - creates a non-owning ref handle.
+ """Create a Buffer from a raw pointer.
 
- Note: The stream parameter is accepted for API compatibility but is
- ignored since non-owning refs are never freed by the handle.
+ When ``mr`` is provided, the buffer takes ownership: ``mr.deallocate()``
+ is called when the buffer is closed or garbage collected. When ``owner``
+ is provided, the owner is kept alive but no deallocation is performed.
  """
  if mr is not None and owner is not None:
  raise ValueError("owner and memory resource cannot be both specified together")
  cdef Buffer self = Buffer.__new__(cls)
- self._h_ptr = deviceptr_create_with_owner(<uintptr_t>(int(ptr)), owner)
+ cdef uintptr_t c_ptr = <uintptr_t>(int(ptr))
+ if mr is not None:
+ self._h_ptr = deviceptr_create_with_mr(c_ptr, size, mr)
+ else:
+ self._h_ptr = deviceptr_create_with_owner(c_ptr, owner)
  self._size = size
  self._memory_resource = mr
  self._ipc_data = IPCDataForBuffer(ipc_descriptor, True) if ipc_descriptor is not None else None
  self._owner = owner
  self._mem_attrs_inited = False
  return self
 
- # No __dealloc__ needed - RAII handles cleanup via _h_ptr destructor
-
  def __reduce__(self):
  # Must not serialize the parent's stream!
  return Buffer.from_ipc_descriptor, (self.memory_resource, self.get_ipc_descriptor())
@@ -112,16 +141,19 @@ cdef class Buffer:
  size : int
  Memory size of the buffer
  mr : :obj:`~_memory.MemoryResource`, optional
- Memory resource associated with the buffer
+ Memory resource associated with the buffer. When provided,
+ :meth:`MemoryResource.deallocate` is called when the buffer is
+ closed or garbage collected.
  owner : object, optional
  An object holding external allocation that the ``ptr`` points to.
  The reference is kept as long as the buffer is alive.
  The ``owner`` and ``mr`` cannot be specified together.
 
  Note
  ----
- This creates a non-owning reference. The pointer will NOT be freed
- when the Buffer is closed or garbage collected.
+ When neither ``mr`` nor ``owner`` is specified, this creates a
+ non-owning reference. The pointer will NOT be freed when the
+ :class:`Buffer` is closed or garbage collected.
  """
  return Buffer._init(ptr, size, mr=mr, owner=owner)
 

cuda_core/cuda/core/_memory/_legacy.py

@@ -51,7 +51,7 @@ def allocate(self, size, stream=None) -> Buffer:
  raise_if_driver_error(err)
  else:
  ptr = 0
- return Buffer._init(ptr, size, self, stream)
+ return Buffer._init(ptr, size, self)
 
  def deallocate(self, ptr: DevicePointerT, size, stream):
  """Deallocate a buffer previously allocated by this resource.
@@ -106,7 +106,7 @@ def allocate(self, size, stream=None) -> Buffer:
  raise_if_driver_error(err)
  else:
  ptr = 0
- return Buffer._init(ptr, size, self, stream)
+ return Buffer._init(ptr, size, self)
 
  def deallocate(self, ptr, size, stream):
  if stream is not None:

cuda_core/cuda/core/_resource_handles.pxd

@@ -108,6 +108,15 @@ cdef DevicePtrHandle deviceptr_alloc(size_t size) except+ nogil
 cdef DevicePtrHandle deviceptr_alloc_host(size_t size) except+ nogil
 cdef DevicePtrHandle deviceptr_create_ref(cydriver.CUdeviceptr ptr) except+ nogil
 cdef DevicePtrHandle deviceptr_create_with_owner(cydriver.CUdeviceptr ptr, object owner) except+ nogil
+cdef DevicePtrHandle deviceptr_create_with_mr(
+ cydriver.CUdeviceptr ptr, size_t size, object mr) except+ nogil
+
+# MR deallocation callback type and registration
+ctypedef void (*MRDeallocCallback)(
+ object mr, cydriver.CUdeviceptr ptr, size_t size,
+ const StreamHandle& stream) noexcept
+cdef void register_mr_dealloc_callback(MRDeallocCallback cb) noexcept
+
 cdef DevicePtrHandle deviceptr_import_ipc(
  const MemoryPoolHandle& h_pool, const void* export_data, const StreamHandle& h_stream) except+ nogil
 cdef StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept nogil

cuda_core/cuda/core/_resource_handles.pyx

@@ -92,6 +92,16 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
  cydriver.CUdeviceptr ptr) except+ nogil
  DevicePtrHandle deviceptr_create_with_owner "cuda_core::deviceptr_create_with_owner" (
  cydriver.CUdeviceptr ptr, object owner) except+ nogil
+
+ # MR deallocation callback
+ ctypedef void (*MRDeallocCallback)(
+ object mr, cydriver.CUdeviceptr ptr, size_t size,
+ const StreamHandle& stream) noexcept
+ void register_mr_dealloc_callback "cuda_core::register_mr_dealloc_callback" (
+ MRDeallocCallback cb) noexcept
+ DevicePtrHandle deviceptr_create_with_mr "cuda_core::deviceptr_create_with_mr" (
+ cydriver.CUdeviceptr ptr, size_t size, object mr) except+ nogil
+
  DevicePtrHandle deviceptr_import_ipc "cuda_core::deviceptr_import_ipc" (
  const MemoryPoolHandle& h_pool, const void* export_data, const StreamHandle& h_stream) except+ nogil
  StreamHandle deallocation_stream "cuda_core::deallocation_stream" (

cuda_core/tests/helpers/buffers.py

@@ -14,6 +14,7 @@
  "DummyUnifiedMemoryResource",
  "make_scratch_buffer",
  "PatternGen",
+ "TrackingMR",
 ]
 
 
@@ -41,6 +42,37 @@ def device_id(self) -> int:
  return self.device
 
 
+class TrackingMR(MemoryResource):
+ """A MemoryResource that tracks active allocations via a dict.
+
+ Useful for verifying that deallocate is called at the expected time.
+ """
+
+ def __init__(self):
+ self.active = {}
+
+ def allocate(self, size, stream=None):
+ ptr = handle_return(driver.cuMemAlloc(size))
+ self.active[int(ptr)] = size
+ return Buffer.from_handle(ptr=ptr, size=size, mr=self)
+
+ def deallocate(self, ptr, size, stream=None):
+ handle_return(driver.cuMemFree(ptr))
+ del self.active[int(ptr)]
+
+ @property
+ def is_device_accessible(self):
+ return True
+
+ @property
+ def is_host_accessible(self):
+ return False
+
+ @property
+ def device_id(self):
+ return 0
+
+
 class PatternGen:
  """
  Provides methods to fill a target buffer with known test patterns and

cuda_core/tests/test_memory.py

@@ -39,7 +39,7 @@
 from cuda.core._utils.cuda_utils import CUDAError, handle_return
 from cuda.core.utils import StridedMemoryView
 from helpers import IS_WINDOWS, supports_ipc_mempool
-from helpers.buffers import DummyUnifiedMemoryResource
+from helpers.buffers import DummyUnifiedMemoryResource, TrackingMR
 
 from conftest import (
  create_managed_memory_resource_or_skip,
@@ -471,6 +471,49 @@ def test_buffer_external_managed(change_device):
  handle_return(driver.cuMemFree(ptr))
 
 
+def test_mr_deallocate_called_on_close():
+ """Buffer.from_handle(mr=mr) calls mr.deallocate() on close (issue #1619)."""
+ device = Device()
+ device.set_current()
+ mr = TrackingMR()
+ buf = mr.allocate(1024)
+ assert len(mr.active) == 1
+ buf.close()
+ assert len(mr.active) == 0
+
+
+def test_mr_deallocate_called_on_gc():
+ """Buffer.from_handle(mr=mr) calls mr.deallocate() on GC (issue #1619)."""
+ import gc
+
+ device = Device()
+ device.set_current()
+ mr = TrackingMR()
+ buf = mr.allocate(1024)
+ assert len(mr.active) == 1
+ del buf
+ gc.collect()
+ assert len(mr.active) == 0
+
+
+def test_mr_deallocate_receives_stream():
+ """Buffer.close(stream) forwards the stream to mr.deallocate() (issue #1619)."""
+ device = Device()
+ device.set_current()
+ stream = device.create_stream()
+ received = {}
+
+ class StreamCaptureMR(TrackingMR):
+ def deallocate(self, ptr, size, stream=None):
+ received["stream"] = stream
+ super().deallocate(ptr, size, stream)
+
+ mr = StreamCaptureMR()
+ buf = mr.allocate(1024)
+ buf.close(stream)
+ assert received["stream"].handle == stream.handle
+
+
 def test_memory_resource_and_owner_disallowed():
  with pytest.raises(ValueError, match="cannot be both specified together"):
  a = (ctypes.c_byte * 20)()