cuda.core: make buffer cleanup context-safe

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;
@@ -188,6 +189,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) {
+            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 +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;
+    // Used by MR-backed pointer deleters; native deleters ignore it.
+    mutable ContextHandle h_context;
 };
 }  // namespace
 
@@ -925,8 +963,15 @@ StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept {
     return get_box(h)->h_stream;
 }
 
-void set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept {
-    get_box(h)->h_stream = h_stream;
+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;
+    if (!box->h_context) {
+        box->h_context = h_context;
+    }
 }
 
 DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h_pool, const StreamHandle& h_stream) {
@@ -937,7 +982,7 @@ DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h
     }
 
     auto box = std::shared_ptr<DevicePtrBox>(
-        new DevicePtrBox{ptr, h_stream},
+        new DevicePtrBox{ptr, h_stream, {}},
         [h_pool](DevicePtrBox* b) {
             GILReleaseGuard gil;
             p_cuMemFreeAsync(b->resource, as_cu(b->h_stream));
@@ -955,7 +1000,7 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream)
     }
 
     auto box = std::shared_ptr<DevicePtrBox>(
-        new DevicePtrBox{ptr, h_stream},
+        new DevicePtrBox{ptr, h_stream, {}},
         [](DevicePtrBox* b) {
             GILReleaseGuard gil;
             p_cuMemFreeAsync(b->resource, as_cu(b->h_stream));
@@ -973,7 +1018,7 @@ DevicePtrHandle deviceptr_alloc(size_t size) {
     }
 
     auto box = std::shared_ptr<DevicePtrBox>(
-        new DevicePtrBox{ptr, StreamHandle{}},
+        new DevicePtrBox{ptr, {}, {}},
         [](DevicePtrBox* b) {
             GILReleaseGuard gil;
             p_cuMemFree(b->resource);
@@ -991,7 +1036,7 @@ DevicePtrHandle deviceptr_alloc_host(size_t size) {
     }
 
     auto box = std::shared_ptr<DevicePtrBox>(
-        new DevicePtrBox{reinterpret_cast<CUdeviceptr>(ptr), StreamHandle{}},
+        new DevicePtrBox{reinterpret_cast<CUdeviceptr>(ptr), {}, {}},
         [](DevicePtrBox* b) {
             GILReleaseGuard gil;
             p_cuMemFreeHost(reinterpret_cast<void*>(b->resource));
@@ -1002,7 +1047,8 @@ DevicePtrHandle deviceptr_alloc_host(size_t size) {
 }
 
 DevicePtrHandle deviceptr_create_ref(CUdeviceptr ptr) {
-    auto box = std::make_shared<DevicePtrBox>(DevicePtrBox{ptr, StreamHandle{}});
+    auto box = std::make_shared<DevicePtrBox>(
+        DevicePtrBox{ptr, {}, {}});
     return DevicePtrHandle(box, &box->resource);
 }
 
@@ -1018,7 +1064,7 @@ DevicePtrHandle deviceptr_create_with_owner(CUdeviceptr ptr, PyObject* owner) {
     }
     Py_INCREF(owner);
     auto box = std::shared_ptr<DevicePtrBox>(
-        new DevicePtrBox{ptr, StreamHandle{}},
+        new DevicePtrBox{ptr, {}, {}},
         [owner](DevicePtrBox* b) {
             GILAcquireGuard gil;
             if (gil.acquired()) {
@@ -1036,7 +1082,7 @@ DevicePtrHandle deviceptr_create_mapped_graphics(
     const StreamHandle& h_stream
 ) {
     auto box = std::shared_ptr<DevicePtrBox>(
-        new DevicePtrBox{ptr, h_stream},
+        new DevicePtrBox{ptr, h_stream, {}},
         [h_resource](DevicePtrBox* b) {
             GILReleaseGuard gil;
             CUgraphicsResource resource = as_cu(h_resource);
@@ -1057,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);
     }
@@ -1068,8 +1118,9 @@ DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject*
     }
     Py_INCREF(mr);
     auto box = std::shared_ptr<DevicePtrBox>(
-        new DevicePtrBox{ptr, StreamHandle{}},
+        new DevicePtrBox{ptr, {}, h_context},
         [mr, size](DevicePtrBox* b) {
+            ScopedCurrentContext context(b->h_context);
             GILAcquireGuard gil;
             if (gil.acquired()) {
                 if (mr_dealloc_cb) {
@@ -1162,7 +1213,7 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
         }
 
         auto box = std::shared_ptr<DevicePtrBox>(
-            new DevicePtrBox{ptr, h_stream},
+            new DevicePtrBox{ptr, h_stream, {}},
             [h_pool, key](DevicePtrBox* b) {
                 ipc_ptr_cache.unregister_handle(key);
                 GILReleaseGuard gil;
@@ -1182,7 +1233,7 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void*
         }
 
         auto box = std::shared_ptr<DevicePtrBox>(
-            new DevicePtrBox{ptr, h_stream},
+            new DevicePtrBox{ptr, h_stream, {}},
             [h_pool](DevicePtrBox* b) {
                 GILReleaseGuard gil;
                 p_cuMemFreeAsync(b->resource, as_cu(b->h_stream));

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;
@@ -402,10 +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 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.
@@ -421,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

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 = <uintptr_t>(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

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 = <cydriver.CUdeviceptr>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):

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