From 821e5c0f10b9e30d6d6cf232804357967d8ca01c Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Wed, 29 Jul 2026 13:27:20 -0700 Subject: [PATCH 1/3] fix(cuda.core): retain contexts for buffer cleanup Ensure deferred and MR-backed buffer destruction activates the allocation or deallocation-stream context, preventing invalid-context leaks after callers clear their current context. --- cuda_core/cuda/core/_cpp/resource_handles.cpp | 172 ++++++++++++++++-- cuda_core/cuda/core/_cpp/resource_handles.hpp | 8 +- cuda_core/cuda/core/_memory/_memory_pool.pyx | 5 +- cuda_core/cuda/core/_resource_handles.pyx | 12 +- .../graph/test_graph_definition_lifetime.py | 44 ++++- cuda_core/tests/test_memory.py | 46 +++++ 6 files changed, 267 insertions(+), 20 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/resource_handles.cpp b/cuda_core/cuda/core/_cpp/resource_handles.cpp index b3d2e3fe373..90df808704f 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.cpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.cpp @@ -34,6 +34,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; @@ -43,6 +44,7 @@ decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate = nullptr; decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority = nullptr; decltype(&cuStreamDestroy) p_cuStreamDestroy = nullptr; +decltype(&cuStreamGetCtx) p_cuStreamGetCtx = nullptr; decltype(&cuEventCreate) p_cuEventCreate = nullptr; decltype(&cuEventDestroy) p_cuEventDestroy = nullptr; @@ -64,6 +66,7 @@ decltype(&cuMemAllocHost) p_cuMemAllocHost = nullptr; decltype(&cuMemFreeAsync) p_cuMemFreeAsync = nullptr; decltype(&cuMemFree) p_cuMemFree = nullptr; decltype(&cuMemFreeHost) p_cuMemFreeHost = nullptr; +decltype(&cuPointerGetAttribute) p_cuPointerGetAttribute = nullptr; decltype(&cuMemPoolImportPointer) p_cuMemPoolImportPointer = nullptr; @@ -188,6 +191,41 @@ class GILAcquireGuard { bool acquired_; }; +// Temporarily make a retained context current, restoring the caller's exact +// prior context (including no current context) on scope exit. +class ScopedCurrentContext { +public: + explicit ScopedCurrentContext(ContextHandle h_context) noexcept + : h_context_(std::move(h_context)) { + CUcontext target = as_cu(h_context_); + if (!target || !p_cuCtxGetCurrent || !p_cuCtxSetCurrent) { + return; + } + + GILReleaseGuard gil; + if (p_cuCtxGetCurrent(&previous_) != CUDA_SUCCESS || + previous_ == target) { + return; + } + changed_ = p_cuCtxSetCurrent(target) == CUDA_SUCCESS; + } + + ~ScopedCurrentContext() { + if (changed_) { + GILReleaseGuard gil; + p_cuCtxSetCurrent(previous_); + } + } + + ScopedCurrentContext(const ScopedCurrentContext&) = delete; + ScopedCurrentContext& operator=(const ScopedCurrentContext&) = delete; + +private: + ContextHandle h_context_; + CUcontext previous_ = nullptr; + bool changed_ = false; +}; + } // namespace // ============================================================================ @@ -906,6 +944,10 @@ struct DevicePtrBox { // through a const DevicePtrHandle. The stream can be changed after // allocation (e.g., to synchronize deallocation with a different stream). mutable StreamHandle h_stream; + // Retain the allocation's context independently of the deallocation + // stream, whose context may change with Buffer.close(stream=...). + ContextHandle h_allocation_context; + mutable ContextHandle h_stream_context; }; } // namespace @@ -925,8 +967,89 @@ StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept { return get_box(h)->h_stream; } +namespace { + +ContextHandle context_handle_ref_best_effort(CUcontext ctx) noexcept { + if (!ctx) { + return {}; + } + try { + return cuda_core::create_context_handle_ref(ctx); + } catch (...) { + return {}; + } +} + +ContextHandle current_context_best_effort() noexcept { + if (!p_cuCtxGetCurrent) { + return {}; + } + CUcontext ctx = nullptr; + { + GILReleaseGuard gil; + if (p_cuCtxGetCurrent(&ctx) != CUDA_SUCCESS) { + return {}; + } + } + return context_handle_ref_best_effort(ctx); +} + +ContextHandle stream_context_best_effort( + const StreamHandle& h_stream) noexcept { + if (!h_stream) { + return {}; + } + if (ContextHandle h_context = get_stream_context(h_stream)) { + return h_context; + } + if (!p_cuStreamGetCtx) { + return {}; + } + + CUcontext ctx = nullptr; + { + GILReleaseGuard gil; + if (p_cuStreamGetCtx(as_cu(h_stream), &ctx) != CUDA_SUCCESS) { + return {}; + } + } + return context_handle_ref_best_effort(ctx); +} + +ContextHandle allocation_context_best_effort(CUdeviceptr ptr) noexcept { + if (ptr && p_cuPointerGetAttribute) { + CUcontext ctx = nullptr; + CUresult status; + { + GILReleaseGuard gil; + status = p_cuPointerGetAttribute( + &ctx, CU_POINTER_ATTRIBUTE_CONTEXT, ptr); + } + if (status == CUDA_SUCCESS && ctx) { + return context_handle_ref_best_effort(ctx); + } + } + // Stream-ordered allocations are not context-associated. Retain the + // wrapping thread's current context as the operational fallback for a + // later default-stream deallocation. + return current_context_best_effort(); +} + +ContextHandle operation_context(const DevicePtrBox* box) noexcept { + return box->h_stream_context + ? box->h_stream_context + : box->h_allocation_context; +} + +} // namespace + void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept { - get_box(h)->h_stream = h_stream; + DevicePtrBox* box = get_box(h); + box->h_stream = h_stream; + box->h_stream_context = stream_context_best_effort(h_stream); + if (!box->h_stream_context) { + box->h_stream_context = current_context_best_effort(); + } } DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h_pool, const StreamHandle& h_stream) { @@ -936,9 +1059,12 @@ DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h return {}; } + ContextHandle h_context = current_context_best_effort(); + ContextHandle h_stream_context = stream_context_best_effort(h_stream); auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, h_stream, h_context, h_stream_context}, [h_pool](DevicePtrBox* b) { + ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); delete b; @@ -954,9 +1080,12 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) return {}; } + ContextHandle h_context = current_context_best_effort(); + ContextHandle h_stream_context = stream_context_best_effort(h_stream); auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, h_stream, h_context, h_stream_context}, [](DevicePtrBox* b) { + ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); delete b; @@ -972,9 +1101,11 @@ DevicePtrHandle deviceptr_alloc(size_t size) { return {}; } + ContextHandle h_context = current_context_best_effort(); auto box = std::shared_ptr( - new DevicePtrBox{ptr, StreamHandle{}}, + new DevicePtrBox{ptr, {}, h_context, {}}, [](DevicePtrBox* b) { + ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFree(b->resource); delete b; @@ -990,9 +1121,12 @@ DevicePtrHandle deviceptr_alloc_host(size_t size) { return {}; } + ContextHandle h_context = current_context_best_effort(); auto box = std::shared_ptr( - new DevicePtrBox{reinterpret_cast(ptr), StreamHandle{}}, + new DevicePtrBox{ + reinterpret_cast(ptr), {}, h_context, {}}, [](DevicePtrBox* b) { + ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFreeHost(reinterpret_cast(b->resource)); delete b; @@ -1002,7 +1136,8 @@ 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, {}, {}, {}}); return DevicePtrHandle(box, &box->resource); } @@ -1018,7 +1153,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, {}, {}, {}}, [owner](DevicePtrBox* b) { GILAcquireGuard gil; if (gil.acquired()) { @@ -1035,9 +1170,12 @@ DevicePtrHandle deviceptr_create_mapped_graphics( const GraphicsResourceHandle& h_resource, const StreamHandle& h_stream ) { + ContextHandle h_context = current_context_best_effort(); + ContextHandle h_stream_context = stream_context_best_effort(h_stream); auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ptr, h_stream, h_context, h_stream_context}, [h_resource](DevicePtrBox* b) { + ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; CUgraphicsResource resource = as_cu(h_resource); p_cuGraphicsUnmapResources(1, &resource, as_cu(b->h_stream)); @@ -1067,9 +1205,11 @@ DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* return deviceptr_create_ref(ptr); } Py_INCREF(mr); + ContextHandle h_context = allocation_context_best_effort(ptr); auto box = std::shared_ptr( - new DevicePtrBox{ptr, StreamHandle{}}, + new DevicePtrBox{ptr, {}, h_context, {}}, [mr, size](DevicePtrBox* b) { + ScopedCurrentContext context(operation_context(b)); GILAcquireGuard gil; if (gil.acquired()) { if (mr_dealloc_cb) { @@ -1161,10 +1301,15 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* return {}; } + ContextHandle h_context = current_context_best_effort(); + ContextHandle h_stream_context = + stream_context_best_effort(h_stream); auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ + ptr, h_stream, h_context, h_stream_context}, [h_pool, key](DevicePtrBox* b) { ipc_ptr_cache.unregister_handle(key); + ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); delete b; @@ -1181,9 +1326,14 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* return {}; } + ContextHandle h_context = current_context_best_effort(); + ContextHandle h_stream_context = + stream_context_best_effort(h_stream); auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream}, + new DevicePtrBox{ + ptr, h_stream, h_context, h_stream_context}, [h_pool](DevicePtrBox* b) { + ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFreeAsync(b->resource, as_cu(b->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 3a9d2d75cff..679e59bb20a 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; @@ -76,6 +77,7 @@ extern decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate; extern decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority; extern decltype(&cuStreamDestroy) p_cuStreamDestroy; +extern decltype(&cuStreamGetCtx) p_cuStreamGetCtx; extern decltype(&cuEventCreate) p_cuEventCreate; extern decltype(&cuEventDestroy) p_cuEventDestroy; @@ -97,6 +99,7 @@ extern decltype(&cuMemAllocHost) p_cuMemAllocHost; extern decltype(&cuMemFreeAsync) p_cuMemFreeAsync; extern decltype(&cuMemFree) p_cuMemFree; extern decltype(&cuMemFreeHost) p_cuMemFreeHost; +extern decltype(&cuPointerGetAttribute) p_cuPointerGetAttribute; extern decltype(&cuMemPoolImportPointer) p_cuMemPoolImportPointer; @@ -402,8 +405,9 @@ using MRDeallocCallback = void (*)(PyObject* mr, CUdeviceptr ptr, 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. +// via the registered callback. The pointer's allocation context is retained +// when discoverable and made current for the 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); diff --git a/cuda_core/cuda/core/_memory/_memory_pool.pyx b/cuda_core/cuda/core/_memory/_memory_pool.pyx index fbb320e02ff..acebb626bce 100644 --- a/cuda_core/cuda/core/_memory/_memory_pool.pyx +++ b/cuda_core/cuda/core/_memory/_memory_pool.pyx @@ -347,11 +347,8 @@ cdef inline void _MP_deallocate( ) noexcept nogil: 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/cuda/core/_resource_handles.pyx b/cuda_core/cuda/core/_resource_handles.pyx index f6fb6ac4e20..6174225b32c 100644 --- a/cuda_core/cuda/core/_resource_handles.pyx +++ b/cuda_core/cuda/core/_resource_handles.pyx @@ -275,6 +275,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)" @@ -284,6 +285,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": # Stream void* p_cuStreamCreateWithPriority "reinterpret_cast(cuda_core::p_cuStreamCreateWithPriority)" void* p_cuStreamDestroy "reinterpret_cast(cuda_core::p_cuStreamDestroy)" + void* p_cuStreamGetCtx "reinterpret_cast(cuda_core::p_cuStreamGetCtx)" # Event void* p_cuEventCreate "reinterpret_cast(cuda_core::p_cuEventCreate)" @@ -310,6 +312,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": void* p_cuMemFreeAsync "reinterpret_cast(cuda_core::p_cuMemFreeAsync)" void* p_cuMemFree "reinterpret_cast(cuda_core::p_cuMemFree)" void* p_cuMemFreeHost "reinterpret_cast(cuda_core::p_cuMemFreeHost)" + void* p_cuPointerGetAttribute "reinterpret_cast(cuda_core::p_cuPointerGetAttribute)" # IPC void* p_cuMemPoolImportPointer "reinterpret_cast(cuda_core::p_cuMemPoolImportPointer)" @@ -376,16 +379,18 @@ cdef void* _get_optional_driver_fn(str name): cdef void _init_driver_fn_pointers() noexcept: - global p_cuDevicePrimaryCtxRetain, p_cuDevicePrimaryCtxRelease, p_cuCtxGetCurrent + global p_cuDevicePrimaryCtxRetain, p_cuDevicePrimaryCtxRelease + global p_cuCtxGetCurrent, p_cuCtxSetCurrent global p_cuGreenCtxCreate, p_cuGreenCtxDestroy, p_cuCtxFromGreenCtx global p_cuDevResourceGenerateDesc, p_cuGreenCtxStreamCreate - global p_cuStreamCreateWithPriority, p_cuStreamDestroy + global p_cuStreamCreateWithPriority, p_cuStreamDestroy, p_cuStreamGetCtx global p_cuEventCreate, p_cuEventDestroy, p_cuIpcOpenEventHandle global p_cuDeviceGetCount global p_cuMemPoolSetAccess, p_cuMemPoolDestroy, p_cuMemPoolCreate global p_cuDeviceGetMemPool, p_cuMemPoolImportFromShareableHandle global p_cuMemAllocFromPoolAsync, p_cuMemAllocAsync, p_cuMemAlloc, p_cuMemAllocHost global p_cuMemFreeAsync, p_cuMemFree, p_cuMemFreeHost + global p_cuPointerGetAttribute global p_cuMemPoolImportPointer global p_cuLibraryLoadFromFile, p_cuLibraryLoadData, p_cuLibraryUnload, p_cuLibraryGetKernel global p_cuGraphDestroy, p_cuGraphExecDestroy @@ -404,6 +409,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") @@ -413,6 +419,7 @@ cdef void _init_driver_fn_pointers() noexcept: # Stream p_cuStreamCreateWithPriority = _get_driver_fn("cuStreamCreateWithPriority") p_cuStreamDestroy = _get_driver_fn("cuStreamDestroy") + p_cuStreamGetCtx = _get_driver_fn("cuStreamGetCtx") # Event p_cuEventCreate = _get_driver_fn("cuEventCreate") @@ -439,6 +446,7 @@ cdef void _init_driver_fn_pointers() noexcept: p_cuMemFreeAsync = _get_driver_fn("cuMemFreeAsync") p_cuMemFree = _get_driver_fn("cuMemFree") p_cuMemFreeHost = _get_driver_fn("cuMemFreeHost") + p_cuPointerGetAttribute = _get_driver_fn("cuPointerGetAttribute") # IPC p_cuMemPoolImportPointer = _get_driver_fn("cuMemPoolImportPointer") diff --git a/cuda_core/tests/graph/test_graph_definition_lifetime.py b/cuda_core/tests/graph/test_graph_definition_lifetime.py index 804cccc1923..5d35ccd0767 100644 --- a/cuda_core/tests/graph/test_graph_definition_lifetime.py +++ b/cuda_core/tests/graph/test_graph_definition_lifetime.py @@ -76,7 +76,14 @@ def _wait_until(predicate, timeout=None, interval=0.02): raise AssertionError(f"condition not satisfied within {timeout}s") -from cuda.core import Device, DeviceMemoryResource, EventOptions, Kernel, LaunchConfig +from cuda.core import ( + Device, + DeviceMemoryResource, + EventOptions, + Kernel, + LaunchConfig, + LegacyPinnedMemoryResource, +) from cuda.core.graph import ( ChildGraphNode, ConditionalNode, @@ -610,6 +617,41 @@ def test_user_object_cleanup_is_coalesced_on_python_thread(init_cuda): assert set(finalized_threads) == {main_thread} +@pytest.mark.agent_authored(model="gpt-5.6") +def test_deferred_buffer_cleanup_restores_no_current_context(init_cuda, capsys): + """Graph-deferred MR cleanup activates its retained allocation context.""" + from cuda.core._utils.cuda_utils import driver, handle_return + + class TrackingPinnedMR(LegacyPinnedMemoryResource): + def __init__(self): + self.deallocations = [] + + def deallocate(self, ptr, size, *, stream=None): + super().deallocate(ptr, size, stream=stream) + self.deallocations.append((int(ptr), size)) + + mr = TrackingPinnedMR() + src = mr.allocate(4) + dst = mr.allocate(4) + graph_def = GraphDefinition() + graph_def.memcpy(dst, src, 4) + + src.close() + dst.close() + assert mr.deallocations == [] + capsys.readouterr() + + previous = handle_return(driver.cuCtxPopCurrent()) + assert int(previous) != 0 + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + + del graph_def + _wait_until(lambda: len(mr.deallocations) == 2) + + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + assert "mr.deallocate() failed" not in capsys.readouterr().err + + @pytest.mark.agent_authored(model="gpt-5.6") def test_pending_call_queue_saturation_preserves_cleanup(tmp_path): """A full CPython queue neither strands nor mis-threads cleanup.""" diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 4427b899765..188e4446012 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -495,6 +495,27 @@ def test_mr_deallocate_called_on_gc(): assert len(mr.active) == 0 +@pytest.mark.agent_authored(model="gpt-5.6") +@pytest.mark.parametrize("replace_stream", [False, True]) +def test_from_handle_mr_deallocation_restores_no_current_context(init_cuda, capsys, replace_stream): + """MR teardown uses retained pointer or replacement-stream context.""" + mr = TrackingMR() + buf = mr.allocate(1024) + stream = init_cuda.create_stream() if replace_stream else None + assert len(mr.active) == 1 + capsys.readouterr() + + previous = handle_return(driver.cuCtxPopCurrent()) + assert int(previous) != 0 + 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 + + def test_mr_deallocate_receives_stream(): """Buffer.close(stream) forwards the stream to mr.deallocate() (issue #1619).""" device = Device() @@ -563,6 +584,31 @@ def deallocate(self, ptr, size, *, stream): assert captured["stream"].handle == default_stream().handle +@pytest.mark.agent_authored(model="gpt-5.6") +def test_from_handle_pool_pointer_deallocates_without_current_context(mempool_device, capsys): + """A context-independent pool pointer uses its deallocation stream context.""" + dev = mempool_device + stream = dev.create_stream() + mr = DeviceMemoryResource(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE)) + size = 256 + ptr = handle_return(driver.cuMemAllocFromPoolAsync(size, mr.handle, stream.handle)) + stream.sync() + used_after_alloc = mr.attributes.used_mem_current + buf = Buffer.from_handle(int(ptr), size, mr=mr) + capsys.readouterr() + + previous = handle_return(driver.cuCtxPopCurrent()) + assert int(previous) != 0 + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + + buf.close(stream) + stream.sync() + + assert mr.attributes.used_mem_current < used_after_alloc + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + assert "mr.deallocate() failed" not in capsys.readouterr().err + + def test_memory_resource_and_owner_disallowed(): with pytest.raises(ValueError, match="cannot be both specified together"): a = (ctypes.c_byte * 20)() From 5bf078fdfffe046c524c36a2d9d680b41ca05483 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Wed, 29 Jul 2026 16:16:15 -0700 Subject: [PATCH 2/3] test(cuda.core): isolate context cleanup state Restore fixture-managed contexts in direct cleanup tests and run deferred pending-call cleanup in a subprocess so parallel workers cannot leak thread-local state. --- .../graph/test_graph_definition_lifetime.py | 78 ++++++++++++------- cuda_core/tests/test_memory.py | 30 +++---- 2 files changed, 69 insertions(+), 39 deletions(-) diff --git a/cuda_core/tests/graph/test_graph_definition_lifetime.py b/cuda_core/tests/graph/test_graph_definition_lifetime.py index 5d35ccd0767..8ad5e23f76c 100644 --- a/cuda_core/tests/graph/test_graph_definition_lifetime.py +++ b/cuda_core/tests/graph/test_graph_definition_lifetime.py @@ -82,7 +82,6 @@ def _wait_until(predicate, timeout=None, interval=0.02): EventOptions, Kernel, LaunchConfig, - LegacyPinnedMemoryResource, ) from cuda.core.graph import ( ChildGraphNode, @@ -618,38 +617,65 @@ def test_user_object_cleanup_is_coalesced_on_python_thread(init_cuda): @pytest.mark.agent_authored(model="gpt-5.6") -def test_deferred_buffer_cleanup_restores_no_current_context(init_cuda, capsys): +def test_deferred_buffer_cleanup_restores_no_current_context(tmp_path): """Graph-deferred MR cleanup activates its retained allocation context.""" - from cuda.core._utils.cuda_utils import driver, handle_return - - class TrackingPinnedMR(LegacyPinnedMemoryResource): - def __init__(self): - self.deallocations = [] + code = f"timeout = {_FINALIZE_TIMEOUT!r}\n" + textwrap.dedent( + """ + import gc + import time - def deallocate(self, ptr, size, *, stream=None): - super().deallocate(ptr, size, stream=stream) - self.deallocations.append((int(ptr), size)) + from cuda.core import Device, LegacyPinnedMemoryResource + from cuda.core.graph import GraphDefinition + from cuda.core._utils.cuda_utils import driver, handle_return - mr = TrackingPinnedMR() - src = mr.allocate(4) - dst = mr.allocate(4) - graph_def = GraphDefinition() - graph_def.memcpy(dst, src, 4) + class TrackingPinnedMR(LegacyPinnedMemoryResource): + def __init__(self): + self.deallocations = [] - src.close() - dst.close() - assert mr.deallocations == [] - capsys.readouterr() + def deallocate(self, ptr, size, *, stream=None): + super().deallocate(ptr, size, stream=stream) + self.deallocations.append((int(ptr), size)) - previous = handle_return(driver.cuCtxPopCurrent()) - assert int(previous) != 0 - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + Device(0).set_current() + mr = TrackingPinnedMR() + src = mr.allocate(4) + dst = mr.allocate(4) + graph_def = GraphDefinition() + graph_def.memcpy(dst, src, 4) + + src.close() + dst.close() + assert mr.deallocations == [] + + previous = handle_return(driver.cuCtxPopCurrent()) + assert int(previous) != 0 + try: + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + del graph_def - del graph_def - _wait_until(lambda: len(mr.deallocations) == 2) + deadline = time.monotonic() + timeout + while len(mr.deallocations) != 2: + gc.collect() + if time.monotonic() >= deadline: + raise AssertionError("deferred buffer cleanup timed out") + time.sleep(0) + time.sleep(0.02) - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 - assert "mr.deallocate() failed" not in capsys.readouterr().err + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + finally: + handle_return(driver.cuCtxSetCurrent(previous)) + """ + ) + result = subprocess.run( # noqa: S603 - controlled interpreter probe + [sys.executable, "-c", code], + capture_output=True, + text=True, + timeout=60, + # Avoid shadowing the installed package with cuda_core/cuda/core. + cwd=tmp_path, + ) + assert result.returncode == 0, result.stderr + assert "mr.deallocate() failed" not in result.stderr @pytest.mark.agent_authored(model="gpt-5.6") diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 188e4446012..1e7bc825f06 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -503,17 +503,19 @@ def test_from_handle_mr_deallocation_restores_no_current_context(init_cuda, caps buf = mr.allocate(1024) stream = init_cuda.create_stream() if replace_stream else None assert len(mr.active) == 1 - capsys.readouterr() previous = handle_return(driver.cuCtxPopCurrent()) assert int(previous) != 0 - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + try: + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 - buf.close(stream) + buf.close(stream) - assert len(mr.active) == 0 - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 - assert "mr.deallocate() failed" not in capsys.readouterr().err + 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)) def test_mr_deallocate_receives_stream(): @@ -595,18 +597,20 @@ def test_from_handle_pool_pointer_deallocates_without_current_context(mempool_de stream.sync() used_after_alloc = mr.attributes.used_mem_current buf = Buffer.from_handle(int(ptr), size, mr=mr) - capsys.readouterr() previous = handle_return(driver.cuCtxPopCurrent()) assert int(previous) != 0 - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + try: + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 - buf.close(stream) - stream.sync() + buf.close(stream) + stream.sync() - assert mr.attributes.used_mem_current < used_after_alloc - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 - assert "mr.deallocate() failed" not in capsys.readouterr().err + assert mr.attributes.used_mem_current < used_after_alloc + assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + assert "mr.deallocate() failed" not in capsys.readouterr().err + finally: + handle_return(driver.cuCtxSetCurrent(previous)) def test_memory_resource_and_owner_disallowed(): From 11e19aefec207be383d98ca790416a47fca944aa Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Wed, 29 Jul 2026 16:56:53 -0700 Subject: [PATCH 3/3] refactor(cuda.core): resolve buffer contexts in Cython Resolve allocation and stream contexts at the Buffer boundary so the C++ resource layer only retains and restores explicit context handles. --- cuda_core/cuda/core/_cpp/resource_handles.cpp | 149 +++--------------- cuda_core/cuda/core/_cpp/resource_handles.hpp | 20 ++- cuda_core/cuda/core/_memory/_buffer.pyx | 71 ++++++++- cuda_core/cuda/core/_resource_handles.pxd | 7 +- cuda_core/cuda/core/_resource_handles.pyx | 13 +- 5 files changed, 117 insertions(+), 143 deletions(-) diff --git a/cuda_core/cuda/core/_cpp/resource_handles.cpp b/cuda_core/cuda/core/_cpp/resource_handles.cpp index 90df808704f..34763135892 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.cpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.cpp @@ -44,7 +44,6 @@ decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate = nullptr; decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority = nullptr; decltype(&cuStreamDestroy) p_cuStreamDestroy = nullptr; -decltype(&cuStreamGetCtx) p_cuStreamGetCtx = nullptr; decltype(&cuEventCreate) p_cuEventCreate = nullptr; decltype(&cuEventDestroy) p_cuEventDestroy = nullptr; @@ -66,7 +65,6 @@ decltype(&cuMemAllocHost) p_cuMemAllocHost = nullptr; decltype(&cuMemFreeAsync) p_cuMemFreeAsync = nullptr; decltype(&cuMemFree) p_cuMemFree = nullptr; decltype(&cuMemFreeHost) p_cuMemFreeHost = nullptr; -decltype(&cuPointerGetAttribute) p_cuPointerGetAttribute = nullptr; decltype(&cuMemPoolImportPointer) p_cuMemPoolImportPointer = nullptr; @@ -198,7 +196,7 @@ class ScopedCurrentContext { explicit ScopedCurrentContext(ContextHandle h_context) noexcept : h_context_(std::move(h_context)) { CUcontext target = as_cu(h_context_); - if (!target || !p_cuCtxGetCurrent || !p_cuCtxSetCurrent) { + if (!target) { return; } @@ -944,10 +942,8 @@ struct DevicePtrBox { // through a const DevicePtrHandle. The stream can be changed after // allocation (e.g., to synchronize deallocation with a different stream). mutable StreamHandle h_stream; - // Retain the allocation's context independently of the deallocation - // stream, whose context may change with Buffer.close(stream=...). - ContextHandle h_allocation_context; - mutable ContextHandle h_stream_context; + // Used by MR-backed pointer deleters; native deleters ignore it. + mutable ContextHandle h_context; }; } // namespace @@ -967,88 +963,14 @@ StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept { return get_box(h)->h_stream; } -namespace { - -ContextHandle context_handle_ref_best_effort(CUcontext ctx) noexcept { - if (!ctx) { - return {}; - } - try { - return cuda_core::create_context_handle_ref(ctx); - } catch (...) { - return {}; - } -} - -ContextHandle current_context_best_effort() noexcept { - if (!p_cuCtxGetCurrent) { - return {}; - } - CUcontext ctx = nullptr; - { - GILReleaseGuard gil; - if (p_cuCtxGetCurrent(&ctx) != CUDA_SUCCESS) { - return {}; - } - } - return context_handle_ref_best_effort(ctx); -} - -ContextHandle stream_context_best_effort( - const StreamHandle& h_stream) noexcept { - if (!h_stream) { - return {}; - } - if (ContextHandle h_context = get_stream_context(h_stream)) { - return h_context; - } - if (!p_cuStreamGetCtx) { - return {}; - } - - CUcontext ctx = nullptr; - { - GILReleaseGuard gil; - if (p_cuStreamGetCtx(as_cu(h_stream), &ctx) != CUDA_SUCCESS) { - return {}; - } - } - return context_handle_ref_best_effort(ctx); -} - -ContextHandle allocation_context_best_effort(CUdeviceptr ptr) noexcept { - if (ptr && p_cuPointerGetAttribute) { - CUcontext ctx = nullptr; - CUresult status; - { - GILReleaseGuard gil; - status = p_cuPointerGetAttribute( - &ctx, CU_POINTER_ATTRIBUTE_CONTEXT, ptr); - } - if (status == CUDA_SUCCESS && ctx) { - return context_handle_ref_best_effort(ctx); - } - } - // Stream-ordered allocations are not context-associated. Retain the - // wrapping thread's current context as the operational fallback for a - // later default-stream deallocation. - return current_context_best_effort(); -} - -ContextHandle operation_context(const DevicePtrBox* box) noexcept { - return box->h_stream_context - ? box->h_stream_context - : box->h_allocation_context; -} - -} // namespace - -void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept { +void set_deallocation_stream( + const DevicePtrHandle& h, + const StreamHandle& h_stream, + const ContextHandle& h_context) noexcept { DevicePtrBox* box = get_box(h); box->h_stream = h_stream; - box->h_stream_context = stream_context_best_effort(h_stream); - if (!box->h_stream_context) { - box->h_stream_context = current_context_best_effort(); + if (!box->h_context) { + box->h_context = h_context; } } @@ -1059,12 +981,9 @@ DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h return {}; } - ContextHandle h_context = current_context_best_effort(); - ContextHandle h_stream_context = stream_context_best_effort(h_stream); auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream, h_context, h_stream_context}, + new DevicePtrBox{ptr, h_stream, {}}, [h_pool](DevicePtrBox* b) { - ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); delete b; @@ -1080,12 +999,9 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) return {}; } - ContextHandle h_context = current_context_best_effort(); - ContextHandle h_stream_context = stream_context_best_effort(h_stream); auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream, h_context, h_stream_context}, + new DevicePtrBox{ptr, h_stream, {}}, [](DevicePtrBox* b) { - ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); delete b; @@ -1101,11 +1017,9 @@ DevicePtrHandle deviceptr_alloc(size_t size) { return {}; } - ContextHandle h_context = current_context_best_effort(); auto box = std::shared_ptr( - new DevicePtrBox{ptr, {}, h_context, {}}, + new DevicePtrBox{ptr, {}, {}}, [](DevicePtrBox* b) { - ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFree(b->resource); delete b; @@ -1121,12 +1035,9 @@ DevicePtrHandle deviceptr_alloc_host(size_t size) { return {}; } - ContextHandle h_context = current_context_best_effort(); auto box = std::shared_ptr( - new DevicePtrBox{ - reinterpret_cast(ptr), {}, h_context, {}}, + new DevicePtrBox{reinterpret_cast(ptr), {}, {}}, [](DevicePtrBox* b) { - ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFreeHost(reinterpret_cast(b->resource)); delete b; @@ -1137,7 +1048,7 @@ DevicePtrHandle deviceptr_alloc_host(size_t size) { DevicePtrHandle deviceptr_create_ref(CUdeviceptr ptr) { auto box = std::make_shared( - DevicePtrBox{ptr, {}, {}, {}}); + DevicePtrBox{ptr, {}, {}}); return DevicePtrHandle(box, &box->resource); } @@ -1153,7 +1064,7 @@ DevicePtrHandle deviceptr_create_with_owner(CUdeviceptr ptr, PyObject* owner) { } Py_INCREF(owner); auto box = std::shared_ptr( - new DevicePtrBox{ptr, {}, {}, {}}, + new DevicePtrBox{ptr, {}, {}}, [owner](DevicePtrBox* b) { GILAcquireGuard gil; if (gil.acquired()) { @@ -1170,12 +1081,9 @@ DevicePtrHandle deviceptr_create_mapped_graphics( const GraphicsResourceHandle& h_resource, const StreamHandle& h_stream ) { - ContextHandle h_context = current_context_best_effort(); - ContextHandle h_stream_context = stream_context_best_effort(h_stream); auto box = std::shared_ptr( - new DevicePtrBox{ptr, h_stream, h_context, h_stream_context}, + new DevicePtrBox{ptr, h_stream, {}}, [h_resource](DevicePtrBox* b) { - ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; CUgraphicsResource resource = as_cu(h_resource); p_cuGraphicsUnmapResources(1, &resource, as_cu(b->h_stream)); @@ -1195,7 +1103,11 @@ void register_mr_dealloc_callback(MRDeallocCallback cb) { mr_dealloc_cb = cb; } -DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* mr) { +DevicePtrHandle deviceptr_create_with_mr( + CUdeviceptr ptr, + size_t size, + PyObject* mr, + const ContextHandle& h_context) { if (!mr) { return deviceptr_create_ref(ptr); } @@ -1205,11 +1117,10 @@ DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* return deviceptr_create_ref(ptr); } Py_INCREF(mr); - ContextHandle h_context = allocation_context_best_effort(ptr); auto box = std::shared_ptr( - new DevicePtrBox{ptr, {}, h_context, {}}, + new DevicePtrBox{ptr, {}, h_context}, [mr, size](DevicePtrBox* b) { - ScopedCurrentContext context(operation_context(b)); + ScopedCurrentContext context(b->h_context); GILAcquireGuard gil; if (gil.acquired()) { if (mr_dealloc_cb) { @@ -1301,15 +1212,10 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* return {}; } - ContextHandle h_context = current_context_best_effort(); - ContextHandle h_stream_context = - stream_context_best_effort(h_stream); auto box = std::shared_ptr( - new DevicePtrBox{ - ptr, h_stream, h_context, h_stream_context}, + new DevicePtrBox{ptr, h_stream, {}}, [h_pool, key](DevicePtrBox* b) { ipc_ptr_cache.unregister_handle(key); - ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFreeAsync(b->resource, as_cu(b->h_stream)); delete b; @@ -1326,14 +1232,9 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* return {}; } - ContextHandle h_context = current_context_best_effort(); - ContextHandle h_stream_context = - stream_context_best_effort(h_stream); auto box = std::shared_ptr( - new DevicePtrBox{ - ptr, h_stream, h_context, h_stream_context}, + new DevicePtrBox{ptr, h_stream, {}}, [h_pool](DevicePtrBox* b) { - ScopedCurrentContext context(operation_context(b)); GILReleaseGuard gil; p_cuMemFreeAsync(b->resource, as_cu(b->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 679e59bb20a..6bdbdc0b316 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.hpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.hpp @@ -77,7 +77,6 @@ extern decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate; extern decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority; extern decltype(&cuStreamDestroy) p_cuStreamDestroy; -extern decltype(&cuStreamGetCtx) p_cuStreamGetCtx; extern decltype(&cuEventCreate) p_cuEventCreate; extern decltype(&cuEventDestroy) p_cuEventDestroy; @@ -99,7 +98,6 @@ extern decltype(&cuMemAllocHost) p_cuMemAllocHost; extern decltype(&cuMemFreeAsync) p_cuMemFreeAsync; extern decltype(&cuMemFree) p_cuMemFree; extern decltype(&cuMemFreeHost) p_cuMemFreeHost; -extern decltype(&cuPointerGetAttribute) p_cuPointerGetAttribute; extern decltype(&cuMemPoolImportPointer) p_cuMemPoolImportPointer; @@ -405,11 +403,15 @@ using MRDeallocCallback = void (*)(PyObject* mr, CUdeviceptr ptr, void register_mr_dealloc_callback(MRDeallocCallback cb); // Create a device pointer handle whose destructor calls mr.deallocate() -// via the registered callback. The pointer's allocation context is retained -// when discoverable and made current for the callback. The mr's refcount is -// incremented and decremented when the handle is released. +// via the registered callback. The supplied context is retained and made +// current for the 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); +DevicePtrHandle deviceptr_create_with_mr( + CUdeviceptr ptr, + size_t size, + PyObject* mr, + const ContextHandle& h_context); // Import a device pointer from IPC via cuMemPoolImportPointer. // When the last reference is released, cuMemFreeAsync is called on the stored stream. @@ -425,7 +427,11 @@ DevicePtrHandle deviceptr_import_ipc( StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept; // Set the deallocation stream for a device pointer handle. -void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept; +// If the handle has no cleanup context, use the supplied stream context. +void set_deallocation_stream( + const DevicePtrHandle& h, + const StreamHandle& h_stream, + const ContextHandle& h_context) noexcept; // ============================================================================ // Library handle functions diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index d5ecac6cc09..b5f12c2cf8e 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -14,10 +14,13 @@ from cuda.core._memory._pinned_memory_resource import PinnedMemoryResource from cuda.core._memory._ipc cimport IPCBufferDescriptor, IPCDataForBuffer from cuda.core._memory cimport _ipc from cuda.core._resource_handles cimport ( + ContextHandle, DevicePtrHandle, StreamHandle, + create_context_handle_ref, deviceptr_create_with_owner, deviceptr_create_with_mr, + get_stream_context, register_mr_dealloc_callback, as_intptr, as_cu, @@ -77,6 +80,65 @@ register_mr_dealloc_callback(_mr_dealloc_callback) __all__ = ['Buffer', 'MemoryResource'] +# Context Resolution Helpers +# -------------------------- +cdef ContextHandle _current_context_handle(): + cdef ContextHandle h_context + cdef cydriver.CUcontext context = NULL + with nogil: + HANDLE_RETURN(cydriver.cuCtxGetCurrent(&context)) + if context != NULL: + h_context = create_context_handle_ref(context) + return h_context + + +cdef ContextHandle _pointer_context_handle(cydriver.CUdeviceptr ptr): + """Return the pointer's context, or the current context for an async allocation.""" + cdef ContextHandle h_context + cdef cydriver.CUcontext context = NULL + cdef cydriver.CUresult result + if ptr == 0: + return h_context + + with nogil: + result = cydriver.cuPointerGetAttribute( + &context, + cydriver.CUpointer_attribute.CU_POINTER_ATTRIBUTE_CONTEXT, + ptr, + ) + if result == cydriver.CUDA_SUCCESS: + if context != NULL: + return create_context_handle_ref(context) + # Stream-ordered allocations are context-independent. Retain the + # wrapping thread's current context for default-stream deallocation. + return _current_context_handle() + if result in ( + cydriver.CUresult.CUDA_ERROR_INVALID_VALUE, + cydriver.CUresult.CUDA_ERROR_NOT_INITIALIZED, + ): + # A foreign host pointer has no CUDA context to retain. + return h_context + with nogil: + HANDLE_RETURN(result) + return h_context + + +cdef ContextHandle _stream_context_handle(Stream stream): + cdef ContextHandle h_context = stream._h_context + cdef cydriver.CUcontext context = NULL + if h_context: + return h_context + h_context = get_stream_context(stream._h_stream) + if h_context: + return h_context + + with nogil: + HANDLE_RETURN(cydriver.cuStreamGetCtx(as_cu(stream._h_stream), &context)) + if context != NULL: + h_context = create_context_handle_ref(context) + return h_context + + # Memory Attribute Query Helpers # ------------------------------ cdef inline int _query_memory_attrs( @@ -188,8 +250,11 @@ cdef class Buffer: raise ValueError("owner and memory resource cannot be both specified together") cdef Buffer self = Buffer.__new__(cls) cdef uintptr_t c_ptr = (int(ptr)) + cdef ContextHandle h_context if mr is not None: - self._h_ptr = deviceptr_create_with_mr(c_ptr, size, mr) + h_context = _pointer_context_handle(c_ptr) + self._h_ptr = deviceptr_create_with_mr( + c_ptr, size, mr, h_context) else: self._h_ptr = deviceptr_create_with_owner(c_ptr, owner) self._size = size @@ -619,12 +684,14 @@ cdef Buffer Buffer_from_deviceptr_handle( cdef inline void Buffer_close(Buffer self, object stream): """Close a buffer, freeing its memory.""" cdef Stream s + cdef ContextHandle h_context if not self._h_ptr: return # Update deallocation stream if provided if stream is not None: s = Stream_accept(stream) - set_deallocation_stream(self._h_ptr, s._h_stream) + h_context = _stream_context_handle(s) + set_deallocation_stream(self._h_ptr, s._h_stream, h_context) # Reset handle - RAII deleter will free the memory (and release owner ref in C++) self._h_ptr.reset() self._size = 0 diff --git a/cuda_core/cuda/core/_resource_handles.pxd b/cuda_core/cuda/core/_resource_handles.pxd index 2f481fee4f8..f925c5be66b 100644 --- a/cuda_core/cuda/core/_resource_handles.pxd +++ b/cuda_core/cuda/core/_resource_handles.pxd @@ -211,7 +211,8 @@ cdef DevicePtrHandle deviceptr_create_mapped_graphics( const GraphicsResourceHandle& h_resource, const StreamHandle& h_stream) except+ nogil cdef DevicePtrHandle deviceptr_create_with_mr( - cydriver.CUdeviceptr ptr, size_t size, object mr) except+ nogil + cydriver.CUdeviceptr ptr, size_t size, object mr, + const ContextHandle& h_context) except+ nogil # MR deallocation callback type and registration ctypedef void (*MRDeallocCallback)( @@ -222,7 +223,9 @@ 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 -cdef void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil +cdef void set_deallocation_stream( + const DevicePtrHandle& h, const StreamHandle& h_stream, + const ContextHandle& h_context) noexcept nogil # Library handles cdef LibraryHandle create_library_handle_from_file(const char* path) except+ nogil diff --git a/cuda_core/cuda/core/_resource_handles.pyx b/cuda_core/cuda/core/_resource_handles.pyx index 6174225b32c..e3dd7ec68d7 100644 --- a/cuda_core/cuda/core/_resource_handles.pyx +++ b/cuda_core/cuda/core/_resource_handles.pyx @@ -122,14 +122,16 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": 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 + cydriver.CUdeviceptr ptr, size_t size, object mr, + const ContextHandle& h_context) 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" ( const DevicePtrHandle& h) noexcept nogil void set_deallocation_stream "cuda_core::set_deallocation_stream" ( - const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept nogil + const DevicePtrHandle& h, const StreamHandle& h_stream, + const ContextHandle& h_context) noexcept nogil # Library handles LibraryHandle create_library_handle_from_file "cuda_core::create_library_handle_from_file" ( @@ -285,7 +287,6 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": # Stream void* p_cuStreamCreateWithPriority "reinterpret_cast(cuda_core::p_cuStreamCreateWithPriority)" void* p_cuStreamDestroy "reinterpret_cast(cuda_core::p_cuStreamDestroy)" - void* p_cuStreamGetCtx "reinterpret_cast(cuda_core::p_cuStreamGetCtx)" # Event void* p_cuEventCreate "reinterpret_cast(cuda_core::p_cuEventCreate)" @@ -312,7 +313,6 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": void* p_cuMemFreeAsync "reinterpret_cast(cuda_core::p_cuMemFreeAsync)" void* p_cuMemFree "reinterpret_cast(cuda_core::p_cuMemFree)" void* p_cuMemFreeHost "reinterpret_cast(cuda_core::p_cuMemFreeHost)" - void* p_cuPointerGetAttribute "reinterpret_cast(cuda_core::p_cuPointerGetAttribute)" # IPC void* p_cuMemPoolImportPointer "reinterpret_cast(cuda_core::p_cuMemPoolImportPointer)" @@ -383,14 +383,13 @@ cdef void _init_driver_fn_pointers() noexcept: global p_cuCtxGetCurrent, p_cuCtxSetCurrent global p_cuGreenCtxCreate, p_cuGreenCtxDestroy, p_cuCtxFromGreenCtx global p_cuDevResourceGenerateDesc, p_cuGreenCtxStreamCreate - global p_cuStreamCreateWithPriority, p_cuStreamDestroy, p_cuStreamGetCtx + global p_cuStreamCreateWithPriority, p_cuStreamDestroy global p_cuEventCreate, p_cuEventDestroy, p_cuIpcOpenEventHandle global p_cuDeviceGetCount global p_cuMemPoolSetAccess, p_cuMemPoolDestroy, p_cuMemPoolCreate global p_cuDeviceGetMemPool, p_cuMemPoolImportFromShareableHandle global p_cuMemAllocFromPoolAsync, p_cuMemAllocAsync, p_cuMemAlloc, p_cuMemAllocHost global p_cuMemFreeAsync, p_cuMemFree, p_cuMemFreeHost - global p_cuPointerGetAttribute global p_cuMemPoolImportPointer global p_cuLibraryLoadFromFile, p_cuLibraryLoadData, p_cuLibraryUnload, p_cuLibraryGetKernel global p_cuGraphDestroy, p_cuGraphExecDestroy @@ -419,7 +418,6 @@ cdef void _init_driver_fn_pointers() noexcept: # Stream p_cuStreamCreateWithPriority = _get_driver_fn("cuStreamCreateWithPriority") p_cuStreamDestroy = _get_driver_fn("cuStreamDestroy") - p_cuStreamGetCtx = _get_driver_fn("cuStreamGetCtx") # Event p_cuEventCreate = _get_driver_fn("cuEventCreate") @@ -446,7 +444,6 @@ cdef void _init_driver_fn_pointers() noexcept: p_cuMemFreeAsync = _get_driver_fn("cuMemFreeAsync") p_cuMemFree = _get_driver_fn("cuMemFree") p_cuMemFreeHost = _get_driver_fn("cuMemFreeHost") - p_cuPointerGetAttribute = _get_driver_fn("cuPointerGetAttribute") # IPC p_cuMemPoolImportPointer = _get_driver_fn("cuMemPoolImportPointer")