diff --git a/cuda_core/cuda/core/_memory/_buffer.pxd b/cuda_core/cuda/core/_memory/_buffer.pxd index a9fa0d7e99c..e7a900da561 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pxd +++ b/cuda_core/cuda/core/_memory/_buffer.pxd @@ -2,10 +2,14 @@ # # SPDX-License-Identifier: Apache-2.0 +from libc.stddef cimport size_t + from libcpp cimport bool as cpp_bool from libcpp.atomic cimport atomic as std_atomic +from cuda.bindings cimport cydriver from cuda.core._resource_handles cimport DevicePtrHandle +from cuda.core._stream cimport Stream cdef struct _MemAttrs: @@ -46,6 +50,15 @@ cdef Buffer Buffer_from_deviceptr_handle( ) +# Wrap a raw device pointer with MR-owned teardown and record the stream. +cdef DevicePtrHandle deviceptr_create_owned_by_mr( + cydriver.CUdeviceptr ptr, + size_t size, + object mr, + Stream stream, +) except * + + # Shared argument coercion for the batched free functions (copy_batch, # prefetch_batch, discard_batch, discard_prefetch_batch). `single_hint` # names the per-buffer API to use instead when a bare Buffer is passed. diff --git a/cuda_core/cuda/core/_memory/_buffer.pyi b/cuda_core/cuda/core/_memory/_buffer.pyi index 5bf6511fa79..be21ff36dc8 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyi +++ b/cuda_core/cuda/core/_memory/_buffer.pyi @@ -297,6 +297,15 @@ class MemoryResource: returned by the :meth:`allocate` method would hold a reference to self, the buffer properties are retrieved simply by looking up the underlying memory resource's respective property.) + + Notes + ----- + Python subclasses of pool-backed memory resources + (:class:`~_memory.DeviceMemoryResource`, :class:`~_memory.PinnedMemoryResource`, + :class:`~_memory.ManagedMemoryResource`) and :class:`~_memory.GraphMemoryResource` + route buffer teardown through :meth:`deallocate`, including for buffers + returned by :meth:`allocate`. Built-in resource types use a direct C++ + deallocation path that bypasses Python during interpreter shutdown. """ def allocate(self, size: int, *, stream: Stream | GraphBuilder) -> Buffer: @@ -336,6 +345,11 @@ class MemoryResource: Keyword-only. The stream on which to perform the deallocation asynchronously. Must be passed explicitly; pass ``device.default_stream`` to use the default stream. + + Notes + ----- + For memory resources that own buffer pointers, this method is also + invoked when a :class:`Buffer` is closed or garbage-collected. """ @property diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index 552bb0bcc8d..57e60b7c4f5 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -92,6 +92,22 @@ cdef inline void _apply_deallocation_stream( HANDLE_RETURN(status) +cdef DevicePtrHandle deviceptr_create_owned_by_mr( + cydriver.CUdeviceptr ptr, + size_t size, + object mr, + Stream stream, +) except *: + """Create an MR-owned device pointer handle with a recorded deallocation stream.""" + cdef DevicePtrHandle h_ptr = deviceptr_create_with_mr(ptr, size, mr) + try: + _apply_deallocation_stream(h_ptr, stream._h_stream) + except BaseException: + h_ptr.reset() + raise + return h_ptr + + __all__ = ['Buffer', 'MemoryResource'] @@ -627,6 +643,15 @@ cdef class MemoryResource: returned by the :meth:`allocate` method would hold a reference to self, the buffer properties are retrieved simply by looking up the underlying memory resource's respective property.) + + Notes + ----- + Python subclasses of pool-backed memory resources + (:class:`~_memory.DeviceMemoryResource`, :class:`~_memory.PinnedMemoryResource`, + :class:`~_memory.ManagedMemoryResource`) and :class:`~_memory.GraphMemoryResource` + route buffer teardown through :meth:`deallocate`, including for buffers + returned by :meth:`allocate`. Built-in resource types use a direct C++ + deallocation path that bypasses Python during interpreter shutdown. """ def allocate(self, size_t size, *, stream: Stream | GraphBuilder) -> Buffer: @@ -673,6 +698,11 @@ cdef class MemoryResource: Keyword-only. The stream on which to perform the deallocation asynchronously. Must be passed explicitly; pass ``device.default_stream`` to use the default stream. + + Notes + ----- + For memory resources that own buffer pointers, this method is also + invoked when a :class:`Buffer` is closed or garbage-collected. """ raise TypeError("MemoryResource.deallocate must be implemented by subclasses.") diff --git a/cuda_core/cuda/core/_memory/_graph_memory_resource.pyx b/cuda_core/cuda/core/_memory/_graph_memory_resource.pyx index 67ecf97f58c..a930b2f7a24 100644 --- a/cuda_core/cuda/core/_memory/_graph_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_graph_memory_resource.pyx @@ -7,7 +7,12 @@ from __future__ import annotations from libc.stdint cimport intptr_t from cuda.bindings cimport cydriver -from cuda.core._memory._buffer cimport Buffer, Buffer_from_deviceptr_handle, MemoryResource +from cuda.core._memory._buffer cimport ( + Buffer, + Buffer_from_deviceptr_handle, + MemoryResource, + deviceptr_create_owned_by_mr, +) from cuda.core._resource_handles cimport ( DevicePtrHandle, deviceptr_alloc_async, @@ -208,20 +213,35 @@ cdef inline int check_capturing(cydriver.CUstream s) except?-1 nogil: "a non-capturing stream.") +cdef inline bint _GMR_is_builtin_type(cyGraphMemoryResource self): + """Return True for the built-in GraphMemoryResource (not Python subclasses).""" + cdef str name = self.__class__.__name__ + cdef str mod = self.__class__.__module__ + return name == "GraphMemoryResource" and mod == "cuda.core._memory._graph_memory_resource" + + cdef inline Buffer GMR_allocate(cyGraphMemoryResource self, size_t size, Stream stream): cdef cydriver.CUstream s = as_cu(stream._h_stream) + cdef cydriver.CUdeviceptr ptr cdef DevicePtrHandle h_ptr + cdef bint use_direct_deleter = _GMR_is_builtin_type(self) with nogil: check_capturing(s) - h_ptr = deviceptr_alloc_async(size, stream._h_stream) - if not h_ptr: - HANDLE_RETURN(get_last_error()) - raise RuntimeError( - f"Failed to allocate {size} bytes from GraphMemoryResource: " - "cuda-core returned an empty allocation handle without recording a CUDA error. " - "This is an internal cuda-core error; please report it with your CUDA driver, " - "CUDA Toolkit, and cuda-python versions." - ) + if use_direct_deleter: + h_ptr = deviceptr_alloc_async(size, stream._h_stream) + else: + HANDLE_RETURN(cydriver.cuMemAllocAsync(&ptr, size, s)) + if use_direct_deleter: + if not h_ptr: + HANDLE_RETURN(get_last_error()) + raise RuntimeError( + f"Failed to allocate {size} bytes from GraphMemoryResource: " + "cuda-core returned an empty allocation handle without recording a CUDA error. " + "This is an internal cuda-core error; please report it with your CUDA driver, " + "CUDA Toolkit, and cuda-python versions." + ) + return Buffer_from_deviceptr_handle(h_ptr, size, self, None) + h_ptr = deviceptr_create_owned_by_mr(ptr, size, self, stream) return Buffer_from_deviceptr_handle(h_ptr, size, self, None) diff --git a/cuda_core/cuda/core/_memory/_memory_pool.pyx b/cuda_core/cuda/core/_memory/_memory_pool.pyx index cccc95a01a2..fad08b02886 100644 --- a/cuda_core/cuda/core/_memory/_memory_pool.pyx +++ b/cuda_core/cuda/core/_memory/_memory_pool.pyx @@ -10,7 +10,12 @@ from libc.stdint cimport uintptr_t from libc.string cimport memset from cuda.bindings cimport cydriver -from cuda.core._memory._buffer cimport Buffer, Buffer_from_deviceptr_handle, MemoryResource +from cuda.core._memory._buffer cimport ( + Buffer, + Buffer_from_deviceptr_handle, + MemoryResource, + deviceptr_create_owned_by_mr, +) from cuda.core._memory cimport _ipc from cuda.core._stream cimport Stream_accept, Stream from cuda.core._resource_handles cimport ( @@ -328,20 +333,39 @@ cdef inline int check_not_capturing(cydriver.CUstream s) except?-1 nogil: "a capturing stream (consider using GraphMemoryResource).") +cdef inline bint _MP_is_builtin_type(_MemPool self): + """Return True for built-in pool MR types (not Python subclasses).""" + cdef str name = self.__class__.__name__ + cdef str mod = self.__class__.__module__ + return mod.startswith("cuda.core._memory._") and name in ( + "DeviceMemoryResource", + "PinnedMemoryResource", + "ManagedMemoryResource", + ) + + cdef Buffer _MP_allocate(_MemPool self, size_t size, Stream stream, type cls = Buffer): cdef cydriver.CUstream s = as_cu(stream._h_stream) + cdef cydriver.CUdeviceptr ptr cdef DevicePtrHandle h_ptr + cdef bint use_direct_deleter = _MP_is_builtin_type(self) with nogil: check_not_capturing(s) - h_ptr = deviceptr_alloc_from_pool(size, self._h_pool, stream._h_stream) - if not h_ptr: - HANDLE_RETURN(get_last_error()) - raise RuntimeError( - f"Failed to allocate {size} bytes from {self.__class__.__name__}: " - "cuda-core returned an empty allocation handle without recording a CUDA error. " - "This is an internal cuda-core error; please report it with your CUDA driver, " - "CUDA Toolkit, and cuda-python versions." - ) + if use_direct_deleter: + h_ptr = deviceptr_alloc_from_pool(size, self._h_pool, stream._h_stream) + else: + HANDLE_RETURN(cydriver.cuMemAllocFromPoolAsync(&ptr, size, as_cu(self._h_pool), s)) + if use_direct_deleter: + if not h_ptr: + HANDLE_RETURN(get_last_error()) + raise RuntimeError( + f"Failed to allocate {size} bytes from {self.__class__.__name__}: " + "cuda-core returned an empty allocation handle without recording a CUDA error. " + "This is an internal cuda-core error; please report it with your CUDA driver, " + "CUDA Toolkit, and cuda-python versions." + ) + return Buffer_from_deviceptr_handle(h_ptr, size, self, None, cls) + h_ptr = deviceptr_create_owned_by_mr(ptr, size, self, stream) return Buffer_from_deviceptr_handle(h_ptr, size, self, None, cls) diff --git a/cuda_core/docs/source/release/1.2.0-notes.rst b/cuda_core/docs/source/release/1.2.0-notes.rst index 4e6f29ba9a6..ea0fbb1d71c 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -44,6 +44,13 @@ Fixes and enhancements free operation. Previously, these errors could be suppressed. Automatic buffer cleanup remains non-raising and reports failures as warnings. +- Python subclasses of pool-backed memory resources and + :class:`GraphMemoryResource` now route buffer teardown from + :meth:`MemoryResource.allocate` through :meth:`MemoryResource.deallocate`, + matching :meth:`Buffer.from_handle` behavior. Built-in resource types keep + the direct C++ deallocation path for interpreter-shutdown safety. + (`#2615 `__) + - Graph node resources are now retained independently across graph clones, executable graphs, updates, node deletion, and in-flight launches. Previously, modifying a graph definition could release resources still used by an diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 3fdbe98885f..7a7a79322ee 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -572,6 +572,94 @@ def test_set_deallocation_stream_rejects_none_and_closed_buffer(): buf.set_deallocation_stream(stream) +class _RecordingPoolMR(DeviceMemoryResource): + """Pool-backed MR that records deallocate() invocations.""" + + def __init__(self, device, options=None): + super().__init__(device, options) + self.dealloc_calls = [] + + def deallocate(self, ptr, size, *, stream): + self.dealloc_calls.append((ptr, size, stream)) + super().deallocate(ptr, size, stream=stream) + + +@pytest.mark.agent_authored(model="cursor-grok-4.6") +@pytest.mark.parametrize("pinned", [False, True], ids=["device", "pinned"]) +def test_pool_backed_mr_deallocate_called_on_close(mempool_device, pinned): + """Pool-backed mr.allocate() honors overridden deallocate() on close (#2615).""" + dev = mempool_device + stream = dev.default_stream + if pinned: + skip_if_pinned_memory_unsupported(dev) + + class RecordingMR(PinnedMemoryResource): + def __init__(self, options=None): + super().__init__(options) + self.dealloc_calls = [] + + def deallocate(self, ptr, size, *, stream): + self.dealloc_calls.append((ptr, size, stream)) + super().deallocate(ptr, size, stream=stream) + + mr = RecordingMR(PinnedMemoryResourceOptions(max_size=POOL_SIZE)) + else: + + class RecordingMR(DeviceMemoryResource): + def __init__(self, device, options=None): + super().__init__(device, options) + self.dealloc_calls = [] + + def deallocate(self, ptr, size, *, stream): + self.dealloc_calls.append((ptr, size, stream)) + super().deallocate(ptr, size, stream=stream) + + mr = RecordingMR(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE)) + + buf = mr.allocate(1024, stream=stream) + assert buf.memory_resource is mr + assert len(mr.dealloc_calls) == 0 + buf.close(stream=stream) + stream.sync() + assert len(mr.dealloc_calls) == 1 + assert mr.dealloc_calls[0][1] == 1024 + + +@pytest.mark.agent_authored(model="cursor-grok-4.6") +def test_pool_backed_mr_deallocate_called_on_gc(mempool_device): + """Pool-backed mr.allocate() honors overridden deallocate() on GC (#2615).""" + import gc + + dev = mempool_device + stream = dev.default_stream + mr = _RecordingPoolMR(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE)) + buf = mr.allocate(1024, stream=stream) + assert len(mr.dealloc_calls) == 0 + del buf + gc.collect() + stream.sync() + assert len(mr.dealloc_calls) == 1 + + +@pytest.mark.agent_authored(model="cursor-grok-4.6") +def test_pool_backed_mr_deallocate_receives_stream(mempool_device): + """Pool-backed mr.allocate() forwards close(stream) to deallocate() (#2615).""" + dev = mempool_device + stream = dev.create_stream() + received = {} + + class StreamCapturePoolMR(_RecordingPoolMR): + def deallocate(self, ptr, size, *, stream): + received["stream"] = stream + super().deallocate(ptr, size, stream=stream) + + mr = StreamCapturePoolMR(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE)) + buf = mr.allocate(1024, stream=stream) + buf.close(stream=stream) + stream.sync() + assert received["stream"].handle == stream.handle + + @pytest.mark.parametrize("buffer_type", [Buffer, ManagedBuffer]) def test_from_handle_mr_records_default_stream(buffer_type): """When a Buffer/ManagedBuffer is minted via :meth:`from_handle` with ``mr`` @@ -2242,25 +2330,36 @@ def test_dmr_handle_and_ownership(mempool_device): @pytest.mark.agent_authored(model="claude-opus-4.8") def test_dmr_deallocate_frees_pool_pointer(mempool_device): - """Closing a Buffer.from_handle(..., mr=mr) view frees the pointer via the Python - _MemPool.deallocate path; the pool's in-use bytes drop back.""" + """Closing a buffer from mr.allocate() frees via deallocate() and returns pool bytes.""" + dev = mempool_device + stream = dev.default_stream + mr = DeviceMemoryResource(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE)) + size = 256 + buf = mr.allocate(size, stream=stream) + stream.sync() + used_after_alloc = mr.attributes.used_mem_current + assert used_after_alloc >= size + buf.close(stream=stream) + stream.sync() + assert int(buf.handle) == 0 + assert mr.attributes.used_mem_current < used_after_alloc + + +@pytest.mark.agent_authored(model="cursor-grok-4.6") +def test_dmr_from_handle_deallocate_frees_pool_pointer(mempool_device): + """Buffer.from_handle(..., mr=mr) also routes teardown through deallocate().""" dev = mempool_device stream = dev.default_stream mr = DeviceMemoryResource(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE)) size = 256 - # Raw pool allocation owned by nobody else, so exactly one owner frees it (no - # double free); a Buffer.from_handle view then routes teardown through the - # Python deallocate path that mr.allocate()'s C++-direct free would skip. ptr = handle_return(driver.cuMemAllocFromPoolAsync(size, mr.handle, stream.handle)) stream.sync() used_after_alloc = mr.attributes.used_mem_current assert used_after_alloc >= size buf = Buffer.from_handle(int(ptr), size, mr=mr) - buf.close(stream) + buf.close(stream=stream) stream.sync() assert int(buf.handle) == 0 - # In-use bytes fell back, so the pointer was actually returned (buf.handle == 0 - # alone wouldn't prove it: the deleter callback swallows a failed free). assert mr.attributes.used_mem_current < used_after_alloc