From 7363024648d22f0632ad2feba188a9f80a8b9ab2 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 6 Aug 2026 12:14:29 -0700 Subject: [PATCH 1/4] cuda.core: capture bound contexts for buffer deallocation streams Record a DeallocationStream at device-pointer creation so default-stream tokens pin the allocation context (and PTDS the allocating thread) instead of relying on ambient state at free time. --- cuda_core/cuda/core/_cpp/resource_handles.cpp | 133 +++++++++++++++--- 1 file changed, 110 insertions(+), 23 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/resource_handles.cpp b/cuda_core/cuda/core/_cpp/resource_handles.cpp index ef1b8d0f2f8..6e694b6c584 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.cpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.cpp @@ -9,12 +9,14 @@ #include #include #include +#include #include #include #include #include #include #include +#include #include #include @@ -726,6 +728,84 @@ StreamHandle get_per_thread_stream() { return handle; } +// ============================================================================ +// Deallocation streams +// +// A DeallocationStream is a StreamHandle used for ordering frees. It differs +// from an ordinary StreamHandle only for default-stream tokens, for which it +// stores the (de)allocation context. Ordinarily, the LEGACY and PER_THREAD +// default streams resolve to whichever context is active at the time they are +// used, but for storing deallocation recipes we need to pin the context. With +// the PER_THREAD token, it is not possible to restore the original stream when +// deallocation runs on a different thread. Therefore, in that case the +// allocating host thread id is also stored so that cross-thread frees can be +// detected and warnings can be issued. +// ============================================================================ + +namespace { + +bool is_default_stream_token(CUstream stream) noexcept { + return stream == nullptr + || stream == CU_STREAM_LEGACY + || stream == CU_STREAM_PER_THREAD; +} + +} // namespace + +// ptds_tid is std::thread::id{} except for CU_STREAM_PER_THREAD. +struct DeallocationStream { + StreamHandle h_stream; + std::thread::id ptds_tid{}; +}; + +// Real streams are returned unchanged. Default-stream tokens without an +// embedded context are bound to the current context when one is current; +// otherwise the ambient token is left as-is. +static DeallocationStream make_deallocation_stream(const StreamHandle& h) { + if (!h) { + return {}; + } + + const CUstream stream = as_cu(h); + if (!is_default_stream_token(stream)) { + return DeallocationStream{h, {}}; + } + + StreamHandle h_bound = h; + if (!get_stream_context(h)) { + ContextHandle h_ctx = get_current_context(); + if (h_ctx) { + // Do not register in stream_registry: the token value alone is not + // a unique stream identity (context is part of the meaning). + auto box = std::shared_ptr( + new StreamBox{stream, h_ctx}); + h_bound = StreamHandle(box, &box->resource); + } + // else: leave the ambient token; later work can fail loudly + } + + std::thread::id ptds_tid{}; + if (stream == CU_STREAM_PER_THREAD) { + ptds_tid = std::this_thread::get_id(); + } + return DeallocationStream{std::move(h_bound), ptds_tid}; +} + +static void warn_if_ptds_cross_thread(const DeallocationStream& stream) noexcept { + if (stream.ptds_tid == std::thread::id{}) { + return; + } + if (stream.ptds_tid == std::this_thread::get_id()) { + return; + } + std::fprintf( + stderr, + "Warning: Buffer deallocation for a per-thread default stream " + "is running on a different host thread than the one that recorded " + "the deallocation stream; ordering relative to the allocating " + "thread's PTDS is not preserved\n"); +} + // ============================================================================ // Event Handles // ============================================================================ @@ -913,10 +993,10 @@ MemoryPoolHandle create_mempool_handle_ipc(int fd, CUmemAllocationHandleType han namespace { struct DevicePtrBox { CUdeviceptr resource; - // Mutable to allow set_deallocation_stream() to update the stream - // through a const DevicePtrHandle. The stream can be changed after - // allocation (e.g., to synchronize deallocation with a different stream). - mutable StreamHandle h_stream; + // Mutable so set_deallocation_stream() can update free ordering through a + // const DevicePtrHandle. Built with make_deallocation_stream so default- + // stream tokens carry a bound context. + mutable DeallocationStream deallocation; }; } // namespace @@ -924,7 +1004,7 @@ struct DevicePtrBox { // This works because DevicePtrHandle is a shared_ptr alias pointing to // &box->resource, so we can compute the containing struct using offsetof. // The const_cast is safe because we only use this to access the mutable -// h_stream member or in the deleter (where the box is being destroyed). +// deallocation member or in the deleter (where the box is being destroyed). static DevicePtrBox* get_box(const DevicePtrHandle& h) { const CUdeviceptr* p = h.get(); return reinterpret_cast( @@ -933,11 +1013,11 @@ static DevicePtrBox* get_box(const DevicePtrHandle& h) { } StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept { - return get_box(h)->h_stream; + return get_box(h)->deallocation.h_stream; } void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept { - get_box(h)->h_stream = h_stream; + get_box(h)->deallocation = make_deallocation_stream(h_stream); } DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h_pool, const StreamHandle& h_stream) { @@ -948,10 +1028,11 @@ DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h } auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, make_deallocation_stream(h_stream)}, [h_pool](DevicePtrBox* b) { GILReleaseGuard gil; - p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); + warn_if_ptds_cross_thread(b->deallocation); + p_cuMemFreeAsync(b->resource, as_cu(b->deallocation.h_stream)); delete b; } ); @@ -966,10 +1047,11 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) } auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, make_deallocation_stream(h_stream)}, [](DevicePtrBox* b) { GILReleaseGuard gil; - p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); + warn_if_ptds_cross_thread(b->deallocation); + p_cuMemFreeAsync(b->resource, as_cu(b->deallocation.h_stream)); delete b; } ); @@ -984,7 +1066,7 @@ DevicePtrHandle deviceptr_alloc(size_t size) { } auto box = std::shared_ptr( - new DevicePtrBox{ptr, StreamHandle{}}, + new DevicePtrBox{ptr, DeallocationStream{}}, [](DevicePtrBox* b) { GILReleaseGuard gil; p_cuMemFree(b->resource); @@ -1002,7 +1084,7 @@ DevicePtrHandle deviceptr_alloc_host(size_t size) { } auto box = std::shared_ptr( - new DevicePtrBox{reinterpret_cast(ptr), StreamHandle{}}, + new DevicePtrBox{reinterpret_cast(ptr), DeallocationStream{}}, [](DevicePtrBox* b) { GILReleaseGuard gil; p_cuMemFreeHost(reinterpret_cast(b->resource)); @@ -1013,7 +1095,7 @@ DevicePtrHandle deviceptr_alloc_host(size_t size) { } DevicePtrHandle deviceptr_create_ref(CUdeviceptr ptr) { - auto box = std::make_shared(DevicePtrBox{ptr, StreamHandle{}}); + auto box = std::make_shared(DevicePtrBox{ptr, DeallocationStream{}}); return DevicePtrHandle(box, &box->resource); } @@ -1029,7 +1111,7 @@ DevicePtrHandle deviceptr_create_with_owner(CUdeviceptr ptr, PyObject* owner) { } Py_INCREF(owner); auto box = std::shared_ptr( - new DevicePtrBox{ptr, StreamHandle{}}, + new DevicePtrBox{ptr, DeallocationStream{}}, [owner](DevicePtrBox* b) { GILAcquireGuard gil; if (gil.acquired()) { @@ -1047,11 +1129,13 @@ DevicePtrHandle deviceptr_create_mapped_graphics( const StreamHandle& h_stream ) { auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, make_deallocation_stream(h_stream)}, [h_resource](DevicePtrBox* b) { GILReleaseGuard gil; CUgraphicsResource resource = as_cu(h_resource); - p_cuGraphicsUnmapResources(1, &resource, as_cu(b->h_stream)); + warn_if_ptds_cross_thread(b->deallocation); + p_cuGraphicsUnmapResources( + 1, &resource, as_cu(b->deallocation.h_stream)); delete b; } ); @@ -1079,12 +1163,13 @@ DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* } Py_INCREF(mr); auto box = std::shared_ptr( - new DevicePtrBox{ptr, StreamHandle{}}, + new DevicePtrBox{ptr, DeallocationStream{}}, [mr, size](DevicePtrBox* b) { GILAcquireGuard gil; if (gil.acquired()) { if (mr_dealloc_cb) { - mr_dealloc_cb(mr, b->resource, size, b->h_stream); + warn_if_ptds_cross_thread(b->deallocation); + mr_dealloc_cb(mr, b->resource, size, b->deallocation.h_stream); } Py_DECREF(mr); } @@ -1173,11 +1258,12 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* } auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, make_deallocation_stream(h_stream)}, [h_pool, key](DevicePtrBox* b) { ipc_ptr_cache.unregister_handle(key); GILReleaseGuard gil; - p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); + warn_if_ptds_cross_thread(b->deallocation); + p_cuMemFreeAsync(b->resource, as_cu(b->deallocation.h_stream)); delete b; } ); @@ -1193,10 +1279,11 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* } auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, make_deallocation_stream(h_stream)}, [h_pool](DevicePtrBox* b) { GILReleaseGuard gil; - p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); + warn_if_ptds_cross_thread(b->deallocation); + p_cuMemFreeAsync(b->resource, as_cu(b->deallocation.h_stream)); delete b; } ); From aa9ed7ebd3c1880be15a75f9e39884d6ebdf95ca Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 6 Aug 2026 12:47:50 -0700 Subject: [PATCH 2/4] cuda.core: activate bound context during device-pointer teardown Make the deallocation stream's context current around free/unmap/MR cleanup so destruction no longer depends on ambient CUDA context, and wire cuCtxSetCurrent into the resource-handles driver table. --- cuda_core/cuda/core/_cpp/resource_handles.cpp | 154 ++++++++++++++---- cuda_core/cuda/core/_cpp/resource_handles.hpp | 1 + cuda_core/cuda/core/_resource_handles.pyx | 3 + 3 files changed, 123 insertions(+), 35 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/resource_handles.cpp b/cuda_core/cuda/core/_cpp/resource_handles.cpp index 6e694b6c584..d25160b914d 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.cpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.cpp @@ -36,6 +36,7 @@ namespace cuda_core { decltype(&cuDevicePrimaryCtxRetain) p_cuDevicePrimaryCtxRetain = nullptr; decltype(&cuDevicePrimaryCtxRelease) p_cuDevicePrimaryCtxRelease = nullptr; decltype(&cuCtxGetCurrent) p_cuCtxGetCurrent = nullptr; +decltype(&cuCtxSetCurrent) p_cuCtxSetCurrent = nullptr; decltype(&cuGreenCtxCreate) p_cuGreenCtxCreate = nullptr; decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy = nullptr; decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx = nullptr; @@ -192,6 +193,53 @@ class GILAcquireGuard { bool acquired_; }; +// Temporarily make a context current, restoring the caller's prior binding +// (including having no context current) on scope exit. The handle is held for +// the duration so the context cannot be destroyed mid-scope. +class ScopedCurrentContext { +public: + explicit ScopedCurrentContext(ContextHandle h_context) noexcept + : h_context_(std::move(h_context)) { + CUcontext target = as_cu(h_context_); + if (!target) { + return; + } + + GILReleaseGuard gil; + status_ = p_cuCtxGetCurrent(&previous_); + if (status_ != CUDA_SUCCESS || previous_ == target) { + return; + } + status_ = p_cuCtxSetCurrent(target); + changed_ = status_ == CUDA_SUCCESS; + } + + ~ScopedCurrentContext() { + if (changed_) { + GILReleaseGuard gil; + CUresult status = p_cuCtxSetCurrent(previous_); + if (status != CUDA_SUCCESS) { + std::fprintf( + stderr, + "Warning: cuCtxSetCurrent (restoring the caller's context) " + "failed (CUDA error %d)\n", + static_cast(status)); + } + } + } + + CUresult status() const noexcept { return status_; } + + ScopedCurrentContext(const ScopedCurrentContext&) = delete; + ScopedCurrentContext& operator=(const ScopedCurrentContext&) = delete; + +private: + ContextHandle h_context_; + CUcontext previous_ = nullptr; + bool changed_ = false; + CUresult status_ = CUDA_SUCCESS; +}; + } // namespace // ============================================================================ @@ -742,16 +790,6 @@ StreamHandle get_per_thread_stream() { // detected and warnings can be issued. // ============================================================================ -namespace { - -bool is_default_stream_token(CUstream stream) noexcept { - return stream == nullptr - || stream == CU_STREAM_LEGACY - || stream == CU_STREAM_PER_THREAD; -} - -} // namespace - // ptds_tid is std::thread::id{} except for CU_STREAM_PER_THREAD. struct DeallocationStream { StreamHandle h_stream; @@ -767,7 +805,9 @@ static DeallocationStream make_deallocation_stream(const StreamHandle& h) { } const CUstream stream = as_cu(h); - if (!is_default_stream_token(stream)) { + if (stream != nullptr + && stream != CU_STREAM_LEGACY + && stream != CU_STREAM_PER_THREAD) { return DeallocationStream{h, {}}; } @@ -791,19 +831,33 @@ static DeallocationStream make_deallocation_stream(const StreamHandle& h) { return DeallocationStream{std::move(h_bound), ptds_tid}; } -static void warn_if_ptds_cross_thread(const DeallocationStream& stream) noexcept { - if (stream.ptds_tid == std::thread::id{}) { - return; +template +CUresult with_deallocation_context( + const DeallocationStream& stream, + const char* operation, + Fn&& fn) noexcept { + if (stream.ptds_tid != std::thread::id{} + && stream.ptds_tid != std::this_thread::get_id()) { + std::fprintf( + stderr, + "Warning: Buffer deallocation for a per-thread default stream " + "is running on a different host thread than the one that recorded " + "the deallocation stream; ordering relative to the allocating " + "thread's PTDS is not preserved\n"); + } + ScopedCurrentContext context(get_stream_context(stream.h_stream)); + CUresult status = context.status(); + if (status == CUDA_SUCCESS) { + status = fn(stream); } - if (stream.ptds_tid == std::this_thread::get_id()) { - return; + if (status != CUDA_SUCCESS) { + std::fprintf( + stderr, + "Warning: %s failed during resource destruction (CUDA error %d)\n", + operation, + static_cast(status)); } - std::fprintf( - stderr, - "Warning: Buffer deallocation for a per-thread default stream " - "is running on a different host thread than the one that recorded " - "the deallocation stream; ordering relative to the allocating " - "thread's PTDS is not preserved\n"); + return status; } // ============================================================================ @@ -1031,8 +1085,13 @@ DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h new DevicePtrBox{ptr, make_deallocation_stream(h_stream)}, [h_pool](DevicePtrBox* b) { GILReleaseGuard gil; - warn_if_ptds_cross_thread(b->deallocation); - p_cuMemFreeAsync(b->resource, as_cu(b->deallocation.h_stream)); + with_deallocation_context( + b->deallocation, + "cuMemFreeAsync", + [b](const DeallocationStream& stream) { + return p_cuMemFreeAsync( + b->resource, as_cu(stream.h_stream)); + }); delete b; } ); @@ -1050,8 +1109,13 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) new DevicePtrBox{ptr, make_deallocation_stream(h_stream)}, [](DevicePtrBox* b) { GILReleaseGuard gil; - warn_if_ptds_cross_thread(b->deallocation); - p_cuMemFreeAsync(b->resource, as_cu(b->deallocation.h_stream)); + with_deallocation_context( + b->deallocation, + "cuMemFreeAsync", + [b](const DeallocationStream& stream) { + return p_cuMemFreeAsync( + b->resource, as_cu(stream.h_stream)); + }); delete b; } ); @@ -1133,9 +1197,13 @@ DevicePtrHandle deviceptr_create_mapped_graphics( [h_resource](DevicePtrBox* b) { GILReleaseGuard gil; CUgraphicsResource resource = as_cu(h_resource); - warn_if_ptds_cross_thread(b->deallocation); - p_cuGraphicsUnmapResources( - 1, &resource, as_cu(b->deallocation.h_stream)); + with_deallocation_context( + b->deallocation, + "cuGraphicsUnmapResources", + [b, &resource](const DeallocationStream& stream) { + return p_cuGraphicsUnmapResources( + 1, &resource, as_cu(stream.h_stream)); + }); delete b; } ); @@ -1168,8 +1236,14 @@ DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* GILAcquireGuard gil; if (gil.acquired()) { if (mr_dealloc_cb) { - warn_if_ptds_cross_thread(b->deallocation); - mr_dealloc_cb(mr, b->resource, size, b->deallocation.h_stream); + with_deallocation_context( + b->deallocation, + "MemoryResource deallocate", + [mr, size, b](const DeallocationStream& stream) { + mr_dealloc_cb( + mr, b->resource, size, stream.h_stream); + return CUDA_SUCCESS; + }); } Py_DECREF(mr); } @@ -1262,8 +1336,13 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* [h_pool, key](DevicePtrBox* b) { ipc_ptr_cache.unregister_handle(key); GILReleaseGuard gil; - warn_if_ptds_cross_thread(b->deallocation); - p_cuMemFreeAsync(b->resource, as_cu(b->deallocation.h_stream)); + with_deallocation_context( + b->deallocation, + "cuMemFreeAsync", + [b](const DeallocationStream& stream) { + return p_cuMemFreeAsync( + b->resource, as_cu(stream.h_stream)); + }); delete b; } ); @@ -1282,8 +1361,13 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* new DevicePtrBox{ptr, make_deallocation_stream(h_stream)}, [h_pool](DevicePtrBox* b) { GILReleaseGuard gil; - warn_if_ptds_cross_thread(b->deallocation); - p_cuMemFreeAsync(b->resource, as_cu(b->deallocation.h_stream)); + with_deallocation_context( + b->deallocation, + "cuMemFreeAsync", + [b](const DeallocationStream& stream) { + return p_cuMemFreeAsync( + b->resource, as_cu(stream.h_stream)); + }); delete b; } ); diff --git a/cuda_core/cuda/core/_cpp/resource_handles.hpp b/cuda_core/cuda/core/_cpp/resource_handles.hpp index 6a1a0edd6c7..e1b6b589b29 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.hpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.hpp @@ -67,6 +67,7 @@ void clear_last_error() noexcept; extern decltype(&cuDevicePrimaryCtxRetain) p_cuDevicePrimaryCtxRetain; extern decltype(&cuDevicePrimaryCtxRelease) p_cuDevicePrimaryCtxRelease; extern decltype(&cuCtxGetCurrent) p_cuCtxGetCurrent; +extern decltype(&cuCtxSetCurrent) p_cuCtxSetCurrent; extern decltype(&cuGreenCtxCreate) p_cuGreenCtxCreate; extern decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy; extern decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx; diff --git a/cuda_core/cuda/core/_resource_handles.pyx b/cuda_core/cuda/core/_resource_handles.pyx index 464fad6c1bf..52fd702c888 100644 --- a/cuda_core/cuda/core/_resource_handles.pyx +++ b/cuda_core/cuda/core/_resource_handles.pyx @@ -293,6 +293,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": void* p_cuDevicePrimaryCtxRetain "reinterpret_cast(cuda_core::p_cuDevicePrimaryCtxRetain)" void* p_cuDevicePrimaryCtxRelease "reinterpret_cast(cuda_core::p_cuDevicePrimaryCtxRelease)" void* p_cuCtxGetCurrent "reinterpret_cast(cuda_core::p_cuCtxGetCurrent)" + void* p_cuCtxSetCurrent "reinterpret_cast(cuda_core::p_cuCtxSetCurrent)" void* p_cuGreenCtxCreate "reinterpret_cast(cuda_core::p_cuGreenCtxCreate)" void* p_cuGreenCtxDestroy "reinterpret_cast(cuda_core::p_cuGreenCtxDestroy)" void* p_cuCtxFromGreenCtx "reinterpret_cast(cuda_core::p_cuCtxFromGreenCtx)" @@ -397,6 +398,7 @@ cdef void* _get_optional_driver_fn(str name): cdef void _init_driver_fn_pointers() noexcept: global p_cuDevicePrimaryCtxRetain, p_cuDevicePrimaryCtxRelease, p_cuCtxGetCurrent + global p_cuCtxSetCurrent global p_cuGreenCtxCreate, p_cuGreenCtxDestroy, p_cuCtxFromGreenCtx global p_cuDevResourceGenerateDesc, p_cuGreenCtxStreamCreate global p_cuStreamCreateWithPriority, p_cuStreamDestroy @@ -425,6 +427,7 @@ cdef void _init_driver_fn_pointers() noexcept: p_cuDevicePrimaryCtxRetain = _get_driver_fn("cuDevicePrimaryCtxRetain") p_cuDevicePrimaryCtxRelease = _get_driver_fn("cuDevicePrimaryCtxRelease") p_cuCtxGetCurrent = _get_driver_fn("cuCtxGetCurrent") + p_cuCtxSetCurrent = _get_driver_fn("cuCtxSetCurrent") p_cuGreenCtxCreate = _get_optional_driver_fn("cuGreenCtxCreate") p_cuGreenCtxDestroy = _get_optional_driver_fn("cuGreenCtxDestroy") p_cuCtxFromGreenCtx = _get_optional_driver_fn("cuCtxFromGreenCtx") From 348d120bdbfbeebf4a4c6520fd68b344db1a0f04 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 6 Aug 2026 13:42:16 -0700 Subject: [PATCH 3/4] cuda.core: record from_handle deallocation streams at creation Add keyword-only stream= on Buffer/ManagedBuffer.from_handle when mr owns the pointer, bind it at construction, and cover teardown with no or foreign current context. --- cuda_core/cuda/core/_memory/_buffer.pyi | 9 +- cuda_core/cuda/core/_memory/_buffer.pyx | 46 +++-- .../cuda/core/_memory/_managed_buffer.py | 7 +- cuda_core/tests/test_memory.py | 179 +++++++++++++++++- 4 files changed, 211 insertions(+), 30 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_buffer.pyi b/cuda_core/cuda/core/_memory/_buffer.pyi index 1d824cf6fc0..9ad745f04b5 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyi +++ b/cuda_core/cuda/core/_memory/_buffer.pyi @@ -39,12 +39,14 @@ class Buffer: ... @classmethod - def _init(cls, ptr: DevicePointerType, size: int, mr: MemoryResource | None=None, ipc_descriptor: IPCBufferDescriptor | None=None, owner: object | None=None) -> Buffer: + def _init(cls, ptr: DevicePointerType, size: int, mr: MemoryResource | None=None, ipc_descriptor: IPCBufferDescriptor | None=None, owner: object | None=None, *, stream: Stream | GraphBuilder | None=None) -> Buffer: """Create a Buffer from a raw pointer. 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. + When ``mr`` is provided, a deallocation stream is recorded at creation + (``stream`` if given, otherwise ``default_stream()``). """ @staticmethod @@ -55,7 +57,7 @@ class Buffer: ... @staticmethod - def from_handle(ptr: DevicePointerType, size: int, mr: MemoryResource | None=None, owner: object | None=None) -> Buffer: + def from_handle(ptr: DevicePointerType, size: int, mr: MemoryResource | None=None, owner: object | None=None, *, stream: Stream | GraphBuilder | None=None) -> Buffer: """Create a new :class:`Buffer` object from a pointer. Parameters @@ -72,6 +74,9 @@ class Buffer: 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. + stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional + Keyword-only. Deallocation stream to record when ``mr`` owns the + pointer. Defaults to the calling thread's default stream. Note ---- diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index 2506331d0fd..6e2ef2f6e1d 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -49,23 +49,21 @@ cdef void _mr_dealloc_callback( size_t size, const StreamHandle& h_stream, ) noexcept: - """Called by the C++ deleter to deallocate via MemoryResource.deallocate. - - This is the C++ teardown path: there is no Python caller frame from - which to obtain a stream. If the device-pointer handle was created - without ``set_deallocation_stream`` being called (e.g. buffers minted - via ``Buffer.from_handle(ptr, size, mr=mr)`` from DLPack import, - third-party adapters, or other foreign sources), ``h_stream`` is - empty here. Stream-ordered MR ``deallocate`` overrides reject - ``stream=None`` (issue #2001), so without a fallback the destructor - would print a warning and leak the allocation. Fall back to the - legacy/per-thread default stream so the free still happens; this is - the unique exception to the "no implicit default-stream fallback" - policy because the teardown has no other source of truth. - """ + """Called by the C++ deleter to deallocate via MemoryResource.deallocate.""" cdef Stream stream try: - stream = Stream._from_handle(Stream, h_stream) if h_stream else default_stream() + if not h_stream: + print( + "Warning: no deallocation stream was recorded; falling back to " + "the default stream for mr.deallocate() during Buffer " + "destruction. This is an internal cuda-core error; please " + "report it with your CUDA driver, CUDA Toolkit, and " + "cuda-python versions.", + file=sys.stderr, + ) + stream = default_stream() + else: + stream = Stream._from_handle(Stream, h_stream) mr.deallocate(int(ptr), size, stream=stream) except Exception as exc: print(f"Warning: mr.deallocate() failed during Buffer destruction: {exc}", @@ -176,20 +174,29 @@ cdef class Buffer: def _init( cls, ptr: DevicePointerType, size_t size, mr: MemoryResource | None = None, ipc_descriptor: IPCBufferDescriptor | None = None, - owner : object | None = None + owner : object | None = None, + *, + stream: Stream | GraphBuilder | None = None, ) -> Buffer: """Create a Buffer from a raw pointer. 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. + When ``mr`` is provided, a deallocation stream is recorded at creation + (``stream`` if given, otherwise ``default_stream()``). """ if mr is not None and owner is not None: raise ValueError("owner and memory resource cannot be both specified together") + if stream is not None and mr is None: + raise ValueError("stream requires a memory resource (mr)") cdef Buffer self = Buffer.__new__(cls) cdef uintptr_t c_ptr = (int(ptr)) + cdef Stream s if mr is not None: self._h_ptr = deviceptr_create_with_mr(c_ptr, size, mr) + s = Stream_accept(default_stream() if stream is None else stream) + set_deallocation_stream(self._h_ptr, s._h_stream) else: self._h_ptr = deviceptr_create_with_owner(c_ptr, owner) self._size = size @@ -217,6 +224,8 @@ cdef class Buffer: def from_handle( ptr: DevicePointerType, size_t size, mr: MemoryResource | None = None, owner: object | None = None, + *, + stream: Stream | GraphBuilder | None = None, ) -> Buffer: """Create a new :class:`Buffer` object from a pointer. @@ -234,6 +243,9 @@ cdef class Buffer: 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. + stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional + Keyword-only. Deallocation stream to record when ``mr`` owns the + pointer. Defaults to the calling thread's default stream. Note ---- @@ -241,7 +253,7 @@ cdef class Buffer: 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) + return Buffer._init(ptr, size, mr=mr, owner=owner, stream=stream) @classmethod def from_ipc_descriptor( diff --git a/cuda_core/cuda/core/_memory/_managed_buffer.py b/cuda_core/cuda/core/_memory/_managed_buffer.py index 83a6c618864..dbd676340ec 100644 --- a/cuda_core/cuda/core/_memory/_managed_buffer.py +++ b/cuda_core/cuda/core/_memory/_managed_buffer.py @@ -154,6 +154,8 @@ def from_handle( size: int, mr: MemoryResource | None = None, owner: object | None = None, + *, + stream: Stream | GraphBuilder | None = None, ) -> Buffer: """Wrap an existing managed-memory pointer in a :class:`ManagedBuffer`. @@ -173,8 +175,11 @@ def from_handle( owner : object, optional An object that keeps the underlying allocation alive. ``owner`` and ``mr`` cannot both be specified. + stream : Stream | GraphBuilder, optional + Keyword-only. Deallocation stream to record when ``mr`` owns the + pointer. Defaults to the calling thread's default stream. """ - return cls._init(ptr, size, mr=mr, owner=owner) + return cls._init(ptr, size, mr=mr, owner=owner, stream=stream) @property def read_mostly(self) -> bool: diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 6af4d025b03..adaa00559f7 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -513,15 +513,11 @@ def deallocate(self, ptr, size, *, stream=None): assert received["stream"].handle == stream.handle -def test_mr_dealloc_callback_falls_back_to_default_stream(): - """When a Buffer's device-pointer handle has no attached deallocation - stream (e.g. buffers minted via :meth:`Buffer.from_handle` from DLPack - import, IPC import, or third-party adapters), the C++ deleter callback - must fall back to the default stream rather than passing ``stream=None`` - to ``mr.deallocate``. Stream-ordered MRs validate the stream and would - otherwise raise ``TypeError`` from inside the ``noexcept`` callback, - which only logs a warning and silently leaks the allocation. See - `#2001 `__. +def test_from_handle_mr_records_default_stream(): + """When a Buffer is minted via :meth:`Buffer.from_handle` with ``mr`` but + without an explicit ``stream=``, the deallocation stream is recorded at + creation as the calling thread's default stream (not chosen later in the + destructor). See `#2497`. """ import gc @@ -553,7 +549,6 @@ def deallocate(self, ptr, size, *, stream): captured["stream"] = Stream_accept(stream) mr = StrictCapturingMR() - # Buffer.from_handle binds mr but does not attach a deallocation stream. # ptr=1 is fine because StrictCapturingMR.deallocate does not free. buf = Buffer.from_handle(1, 1024, mr=mr) del buf @@ -563,6 +558,170 @@ def deallocate(self, ptr, size, *, stream): assert captured["stream"].handle == default_stream().handle +@pytest.mark.agent_authored(model="cursor-grok-4.5") +def test_from_handle_mr_records_explicit_stream(): + """Buffer.from_handle(..., mr=mr, stream=s) stores s for teardown.""" + import gc + + from cuda.core._stream import Stream_accept + + device = Device() + device.set_current() + stream = device.create_stream() + captured = {} + + class StrictCapturingMR(MemoryResource): + @property + def is_device_accessible(self): + return True + + @property + def is_host_accessible(self): + return False + + @property + def device_id(self): + return device.device_id + + def allocate(self, size, *, stream): + raise NotImplementedError + + def deallocate(self, ptr, size, *, stream): + captured["stream"] = Stream_accept(stream) + + mr = StrictCapturingMR() + buf = Buffer.from_handle(1, 1024, mr=mr, stream=stream) + del buf + gc.collect() + + assert captured["stream"].handle == stream.handle + + +@pytest.mark.agent_authored(model="cursor-grok-4.5") +def test_from_handle_stream_requires_mr(): + device = Device() + device.set_current() + stream = device.create_stream() + with pytest.raises(ValueError, match="stream requires a memory resource"): + Buffer.from_handle(1, 1024, stream=stream) + + +@pytest.mark.agent_authored(model="cursor-grok-4.5") +@pytest.mark.parametrize("replace_stream", [False, True]) +def test_mr_deallocation_without_current_context(init_cuda, capsys, replace_stream): + """MR-backed Buffer teardown activates the recorded context when none is current.""" + mr = TrackingMR() + buf = mr.allocate(1024) + stream = init_cuda.create_stream() if replace_stream else None + assert len(mr.active) == 1 + + previous = handle_return(driver.cuCtxPopCurrent()) + assert int(previous) != 0 + try: + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + + buf.close(stream) + + assert len(mr.active) == 0 + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + assert "mr.deallocate() failed" not in capsys.readouterr().err + finally: + handle_return(driver.cuCtxSetCurrent(previous)) + + +@pytest.mark.agent_authored(model="cursor-grok-4.5") +@pytest.mark.parametrize("replace_stream", [False, True]) +def test_mr_deallocation_with_foreign_context(capsys, replace_stream): + """MR-backed Buffer teardown switches away from an unrelated current context.""" + if ccx_system.get_num_devices() < 2: + pytest.skip("Test requires at least 2 GPUs") + + alloc_dev = Device(0) + alloc_dev.set_current() + mr = TrackingMR() + buf = mr.allocate(1024) + stream = alloc_dev.create_stream() if replace_stream else None + assert len(mr.active) == 1 + alloc_ctx = int(handle_return(driver.cuCtxGetCurrent())) + + foreign_dev = Device(1) + foreign_dev.set_current() + foreign_ctx = int(handle_return(driver.cuCtxGetCurrent())) + assert foreign_ctx != 0 + assert foreign_ctx != alloc_ctx + + try: + buf.close(stream) + + assert len(mr.active) == 0 + assert int(handle_return(driver.cuCtxGetCurrent())) == foreign_ctx + assert "mr.deallocate() failed" not in capsys.readouterr().err + finally: + alloc_dev.set_current() + + +@pytest.mark.agent_authored(model="cursor-grok-4.5") +def test_pool_buffer_deallocates_without_current_context(mempool_device, capsys): + """Pool Buffer.close frees on the recorded stream with no current context.""" + dev = mempool_device + stream = dev.create_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 + + previous = handle_return(driver.cuCtxPopCurrent()) + assert int(previous) != 0 + try: + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + + buf.close() + stream.sync() + + assert mr.attributes.used_mem_current < used_after_alloc + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + err = capsys.readouterr().err + assert "failed during resource destruction" not in err + assert "mr.deallocate() failed" not in err + finally: + handle_return(driver.cuCtxSetCurrent(previous)) + + +@pytest.mark.agent_authored(model="cursor-grok-4.5") +def test_pool_buffer_deallocates_with_foreign_context(mempool_device_x2, capsys): + """Pool Buffer.close frees under the recorded context while another is current.""" + alloc_dev, foreign_dev = mempool_device_x2 + alloc_dev.set_current() + stream = alloc_dev.create_stream() + mr = DeviceMemoryResource(alloc_dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE)) + size = 256 + buf = mr.allocate(size, stream=stream) + stream.sync() + used_after_alloc = mr.attributes.used_mem_current + alloc_ctx = int(handle_return(driver.cuCtxGetCurrent())) + + foreign_dev.set_current() + foreign_ctx = int(handle_return(driver.cuCtxGetCurrent())) + assert foreign_ctx != 0 + assert foreign_ctx != alloc_ctx + + try: + buf.close() + assert int(handle_return(driver.cuCtxGetCurrent())) == foreign_ctx + + # Observe the free on the allocation device, then restore the foreign context. + alloc_dev.set_current() + stream.sync() + assert mr.attributes.used_mem_current < used_after_alloc + foreign_dev.set_current() + + err = capsys.readouterr().err + assert "failed during resource destruction" not in err + finally: + alloc_dev.set_current() + + def test_memory_resource_and_owner_disallowed(): with pytest.raises(ValueError, match="cannot be both specified together"): a = (ctypes.c_byte * 20)() From 9566f5badf48aa9547241ebe9f39cedcdc1eee14 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 6 Aug 2026 15:08:38 -0700 Subject: [PATCH 4/4] cuda.core: fail loudly on MemoryResource free errors Stop treating CUDA_ERROR_INVALID_CONTEXT as a successful pool free, and let explicit mr.deallocate() raise; destruction still contains errors in the callback. Document PTDS deallocation ordering on the stream parameters and note the context-safe Buffer teardown fix in the 1.2.0 release notes. --- cuda_core/cuda/core/_memory/_buffer.pyi | 13 ++++++++++--- cuda_core/cuda/core/_memory/_buffer.pyx | 13 ++++++++++--- .../core/_memory/_graph_memory_resource.pyx | 2 +- .../cuda/core/_memory/_managed_buffer.py | 7 +++++-- cuda_core/cuda/core/_memory/_memory_pool.pyx | 7 ++----- cuda_core/docs/source/release/1.2.0-notes.rst | 13 +++++++++++++ cuda_core/tests/test_memory.py | 19 +++++++++++++++++-- 7 files changed, 58 insertions(+), 16 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_buffer.pyi b/cuda_core/cuda/core/_memory/_buffer.pyi index 9ad745f04b5..c136c6bec5f 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyi +++ b/cuda_core/cuda/core/_memory/_buffer.pyi @@ -75,8 +75,11 @@ class Buffer: The reference is kept as long as the buffer is alive. The ``owner`` and ``mr`` cannot be specified together. stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional - Keyword-only. Deallocation stream to record when ``mr`` owns the - pointer. Defaults to the calling thread's default stream. + Keyword-only. The stream used to order the buffer's deallocation + when ``mr`` owns the pointer. Defaults to ``default_stream()``. If + the buffer may be freed from a different host thread, pass a stream + other than the per-thread default stream, which refers to a + different stream on each thread. Note ---- @@ -269,7 +272,11 @@ class MemoryResource: stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder` Keyword-only. The stream on which to perform the allocation asynchronously. Must be passed explicitly; pass - ``device.default_stream`` to use the default stream. + ``device.default_stream`` to use the default stream. This stream + also orders the buffer's eventual deallocation, so if the buffer may + be freed from a different host thread, prefer a stream other than + the per-thread default stream, which refers to a different stream on + each thread. Returns ------- diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index 6e2ef2f6e1d..2d26ac61d3e 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -244,8 +244,11 @@ cdef class Buffer: The reference is kept as long as the buffer is alive. The ``owner`` and ``mr`` cannot be specified together. stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder`, optional - Keyword-only. Deallocation stream to record when ``mr`` owns the - pointer. Defaults to the calling thread's default stream. + Keyword-only. The stream used to order the buffer's deallocation + when ``mr`` owns the pointer. Defaults to ``default_stream()``. If + the buffer may be freed from a different host thread, pass a stream + other than the per-thread default stream, which refers to a + different stream on each thread. Note ---- @@ -558,7 +561,11 @@ cdef class MemoryResource: stream : :obj:`~_stream.Stream` | :obj:`~graph.GraphBuilder` Keyword-only. The stream on which to perform the allocation asynchronously. Must be passed explicitly; pass - ``device.default_stream`` to use the default stream. + ``device.default_stream`` to use the default stream. This stream + also orders the buffer's eventual deallocation, so if the buffer may + be freed from a different host thread, prefer a stream other than + the per-thread default stream, which refers to a different stream on + each thread. Returns ------- diff --git a/cuda_core/cuda/core/_memory/_graph_memory_resource.pyx b/cuda_core/cuda/core/_memory/_graph_memory_resource.pyx index e845a47b080..67ecf97f58c 100644 --- a/cuda_core/cuda/core/_memory/_graph_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_graph_memory_resource.pyx @@ -225,7 +225,7 @@ cdef inline Buffer GMR_allocate(cyGraphMemoryResource self, size_t size, Stream return Buffer_from_deviceptr_handle(h_ptr, size, self, None) -cdef inline void GMR_deallocate(intptr_t ptr, size_t size, Stream stream) noexcept: +cdef inline void GMR_deallocate(intptr_t ptr, size_t size, Stream stream) except *: cdef cydriver.CUstream s = as_cu(stream._h_stream) cdef cydriver.CUdeviceptr devptr = ptr with nogil: diff --git a/cuda_core/cuda/core/_memory/_managed_buffer.py b/cuda_core/cuda/core/_memory/_managed_buffer.py index dbd676340ec..448ed4a206c 100644 --- a/cuda_core/cuda/core/_memory/_managed_buffer.py +++ b/cuda_core/cuda/core/_memory/_managed_buffer.py @@ -176,8 +176,11 @@ def from_handle( An object that keeps the underlying allocation alive. ``owner`` and ``mr`` cannot both be specified. stream : Stream | GraphBuilder, optional - Keyword-only. Deallocation stream to record when ``mr`` owns the - pointer. Defaults to the calling thread's default stream. + Keyword-only. The stream used to order the buffer's deallocation + when ``mr`` owns the pointer. Defaults to ``default_stream()``. If + the buffer may be freed from a different host thread, pass a stream + other than the per-thread default stream, which refers to a + different stream on each thread. """ return cls._init(ptr, size, mr=mr, owner=owner, stream=stream) diff --git a/cuda_core/cuda/core/_memory/_memory_pool.pyx b/cuda_core/cuda/core/_memory/_memory_pool.pyx index 8f9a4354b84..cccc95a01a2 100644 --- a/cuda_core/cuda/core/_memory/_memory_pool.pyx +++ b/cuda_core/cuda/core/_memory/_memory_pool.pyx @@ -347,14 +347,11 @@ cdef Buffer _MP_allocate(_MemPool self, size_t size, Stream stream, type cls = B cdef inline void _MP_deallocate( _MemPool self, uintptr_t ptr, size_t size, Stream stream -) noexcept nogil: +) except *: cdef cydriver.CUstream s = as_cu(stream._h_stream) cdef cydriver.CUdeviceptr devptr = ptr - cdef cydriver.CUresult r with nogil: - r = cydriver.cuMemFreeAsync(devptr, s) - if r != cydriver.CUDA_ERROR_INVALID_CONTEXT: - HANDLE_RETURN(r) + HANDLE_RETURN(cydriver.cuMemFreeAsync(devptr, s)) cdef inline _MP_close(_MemPool self): 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 4bb81cec759..7482ae9eefd 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -9,6 +9,19 @@ Fixes and enhancements ---------------------- +- A :class:`Buffer` is now freed correctly even when the CUDA context current + at teardown is not the one it was allocated in, or when no context is current + at all. This happens routinely when a buffer is released by the garbage + collector on another thread or by deferred CUDA graph cleanup; previously the + free could fail or be skipped, leaking the allocation. + (`#2497 `__) + +- :meth:`Buffer.from_handle` and :meth:`ManagedBuffer.from_handle` accept a + keyword-only ``stream`` that records the stream used to order the buffer's + deallocation when the memory resource owns the pointer. It defaults to + ``default_stream()``. + (`#2497 `__) + - 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 adaa00559f7..04f7b6ec785 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -516,8 +516,8 @@ def deallocate(self, ptr, size, *, stream=None): def test_from_handle_mr_records_default_stream(): """When a Buffer is minted via :meth:`Buffer.from_handle` with ``mr`` but without an explicit ``stream=``, the deallocation stream is recorded at - creation as the calling thread's default stream (not chosen later in the - destructor). See `#2497`. + creation as ``default_stream()`` (not chosen later in the destructor). + See `#2497`. """ import gc @@ -660,6 +660,21 @@ def test_mr_deallocation_with_foreign_context(capsys, replace_stream): alloc_dev.set_current() +@pytest.mark.agent_authored(model="claude-opus-5") +def test_mr_deallocate_raises_on_driver_error(mempool_device): + """An explicit mr.deallocate() call propagates driver errors to the caller. + + Buffer teardown must not raise, so the containment lives in the destruction + callback rather than in deallocate() itself. See `#2497`. + """ + dev = mempool_device + stream = dev.create_stream() + mr = DeviceMemoryResource(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE)) + + with pytest.raises(CUDAError): + mr.deallocate(0xDEADBEEF, 256, stream=stream) + + @pytest.mark.agent_authored(model="cursor-grok-4.5") def test_pool_buffer_deallocates_without_current_context(mempool_device, capsys): """Pool Buffer.close frees on the recorded stream with no current context."""