Simplify green context view handles

Author: leofang

cuda_core/cuda/core/_context.pxd

@@ -24,6 +24,4 @@ cdef class Context:
     @staticmethod
     cdef Context _from_green_ctx(type cls, GreenCtxHandle h_green_ctx, int device_id)
 
-    cdef int _ensure_context_handle(self) except -1
-
     cpdef close(self)

cuda_core/cuda/core/_context.pyx

@@ -10,7 +10,6 @@ from cuda.core._resource_handles cimport (
     GreenCtxHandle,
     as_cu,
     create_context_handle_from_green_ctx,
-    ensure_context_handle,
     get_context_green_ctx,
     get_last_error,
     as_intptr,
@@ -48,29 +47,18 @@ cdef class Context:
         cdef Context ctx = cls.__new__(cls)
         ctx._h_green_ctx = h_green_ctx
         ctx._h_context = create_context_handle_from_green_ctx(h_green_ctx)
+        if not ctx._h_context:
+            HANDLE_RETURN(get_last_error())
+            raise RuntimeError("Failed to create CUDA context view from green context")
         ctx._device_id = device_id
         ctx._is_green = True
         return ctx
 
-    cdef int _ensure_context_handle(self) except -1:
-        cdef cydriver.CUcontext raw_ctx
-        if not self._h_context:
-            return 0
-        if as_cu(self._h_context) != NULL:
-            return 0
-        with nogil:
-            raw_ctx = ensure_context_handle(self._h_context)
-        if raw_ctx == NULL:
-            HANDLE_RETURN(get_last_error())
-            raise RuntimeError("Failed to materialize CUDA context from green context")
-        return 0
-
     @property
     def handle(self):
         """Return the underlying CUcontext handle."""
         if not self._h_context:
             return None
-        self._ensure_context_handle()
         if as_cu(self._h_context) == NULL:
             return None
         return as_py(self._h_context)
@@ -102,16 +90,12 @@ cdef class Context:
         if not isinstance(other, Context):
             return NotImplemented
         cdef Context _other = <Context>other
-        self._ensure_context_handle()
-        _other._ensure_context_handle()
         return as_intptr(self._h_context) == as_intptr(_other._h_context)
 
     def __hash__(self) -> int:
-        self._ensure_context_handle()
         return hash(as_intptr(self._h_context))
 
     def __repr__(self) -> str:
-        self._ensure_context_handle()
         return f"<Context handle={as_intptr(self._h_context):#x} device={self._device_id}>"
 
 

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -11,7 +11,6 @@
 #include <mutex>
 #include <stdexcept>
 #include <unordered_map>
-#include <utility>
 #include <vector>
 
 #ifndef _WIN32
@@ -227,12 +226,8 @@ void clear_last_error() noexcept {
 
 namespace {
 struct ContextBox {
-    mutable CUcontext resource;
+    CUcontext resource;
     GreenCtxHandle h_green_ctx;
-    mutable std::mutex mutex;
-
-    explicit ContextBox(CUcontext resource, GreenCtxHandle h_green_ctx = {})
-        : resource(resource), h_green_ctx(std::move(h_green_ctx)) {}
 };
 
 struct GreenCtxBox {
@@ -262,7 +257,7 @@ ContextHandle create_context_handle_ref(CUcontext ctx) {
         return h;
     }
     auto box = std::shared_ptr<const ContextBox>(
-        new ContextBox(ctx),
+        new ContextBox{ctx, {}},
         [](const ContextBox* b) {
             context_registry.unregister_handle(b->resource);
             delete b;
@@ -273,57 +268,31 @@ ContextHandle create_context_handle_ref(CUcontext ctx) {
     return h;
 }
 
-static const GreenCtxBox* get_box(const GreenCtxHandle& h) noexcept {
-    const CUgreenCtx* p = h.get();
-    return reinterpret_cast<const GreenCtxBox*>(
-        reinterpret_cast<const char*>(p) - offsetof(GreenCtxBox, resource)
-    );
-}
-
 ContextHandle create_context_handle_from_green_ctx(const GreenCtxHandle& h_green_ctx) {
     if (!h_green_ctx) {
         return {};
     }
-    auto box = std::shared_ptr<const ContextBox>(
-        new ContextBox(nullptr, h_green_ctx),
-        [](const ContextBox* b) {
-            if (b->resource) {
-                context_registry.unregister_handle(b->resource);
-            }
-            delete b;
-        }
-    );
-    return ContextHandle(box, &box->resource);
-}
-
-CUcontext ensure_context_handle(const ContextHandle& h) noexcept {
-    if (!h) {
-        err = CUDA_ERROR_INVALID_CONTEXT;
-        return nullptr;
-    }
-
-    const ContextBox* box = get_box(h);
-    std::lock_guard<std::mutex> lock(box->mutex);
-    if (box->resource) {
-        return box->resource;
-    }
-    if (!box->h_green_ctx) {
-        err = CUDA_ERROR_INVALID_CONTEXT;
-        return nullptr;
-    }
     if (!p_cuCtxFromGreenCtx) {
         err = CUDA_ERROR_NOT_SUPPORTED;
-        return nullptr;
+        return {};
     }
 
     GILReleaseGuard gil;
     CUcontext ctx = nullptr;
-    if (CUDA_SUCCESS != (err = p_cuCtxFromGreenCtx(&ctx, as_cu(box->h_green_ctx)))) {
-        return nullptr;
+    if (CUDA_SUCCESS != (err = p_cuCtxFromGreenCtx(&ctx, as_cu(h_green_ctx)))) {
+        return {};
     }
-    box->resource = ctx;
+
+    auto box = std::shared_ptr<const ContextBox>(
+        new ContextBox{ctx, h_green_ctx},
+        [](const ContextBox* b) {
+            context_registry.unregister_handle(b->resource);
+            delete b;
+        }
+    );
+    ContextHandle h(box, &box->resource);
     context_registry.register_handle(ctx, h);
-    return ctx;
+    return h;
 }
 
 GreenCtxHandle get_context_green_ctx(const ContextHandle& h) noexcept {
@@ -399,7 +368,7 @@ ContextHandle get_primary_context(int device_id) {
     }
 
     auto box = std::shared_ptr<const ContextBox>(
-        new ContextBox(ctx),
+        new ContextBox{ctx, {}},
         [device_id](const ContextBox* b) {
             context_registry.unregister_handle(b->resource);
             GILReleaseGuard gil;

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -170,15 +170,11 @@ using FileDescriptorHandle = std::shared_ptr<const int>;
 // Function to create a non-owning context handle (references existing context).
 ContextHandle create_context_handle_ref(CUcontext ctx);
 
-// Create a context handle whose CUcontext view is lazily materialized from
-// the provided green context. The returned ContextHandle keeps the green
-// context alive.
+// Create a context handle for the CUcontext view of the provided green context.
+// The returned ContextHandle keeps the green context alive, but the CUcontext
+// view is non-owning and is not destroyed independently.
 ContextHandle create_context_handle_from_green_ctx(const GreenCtxHandle& h_green_ctx);
 
-// Ensure a ContextHandle has a materialized CUcontext value. For green-context
-// views this calls cuCtxFromGreenCtx once and caches the non-owning CUcontext.
-CUcontext ensure_context_handle(const ContextHandle& h) noexcept;
-
 // Return the green context dependency associated with a ContextHandle, if any.
 GreenCtxHandle get_context_green_ctx(const ContextHandle& h) noexcept;
 

cuda_core/cuda/core/_device.pyx

@@ -1262,7 +1262,6 @@ class Device:
             if self._has_inited and self._context is not None:
                 prev_owned = self._context
             # prev_ctx is the previous context
-            ctx._ensure_context_handle()
             curr_ctx = as_cu(ctx._h_context)
             prev_ctx = NULL
             with nogil:
@@ -1271,8 +1270,6 @@ class Device:
             self._has_inited = True
             self._context = ctx  # Store owning context reference
             if prev_ctx != NULL:
-                if prev_owned is not None:
-                    prev_owned._ensure_context_handle()
                 if prev_owned is not None and as_cu(prev_owned._h_context) == prev_ctx:
                     return prev_owned
                 return Context._from_handle(Context, create_context_handle_ref(prev_ctx), self._device_id)
@@ -1418,7 +1415,6 @@ class Device:
         """
         self._check_context_initialized()
         cdef Context ctx = self._context
-        ctx._ensure_context_handle()
         return cyEvent._init(cyEvent, self._device_id, ctx._h_context, options, True)
 
     def allocate(self, size, stream: Stream | GraphBuilder | None = None) -> Buffer:

cuda_core/cuda/core/_resource_handles.pxd

@@ -116,7 +116,6 @@ cdef void clear_last_error() noexcept nogil
 # Context handles
 cdef ContextHandle create_context_handle_ref(cydriver.CUcontext ctx) except+ nogil
 cdef ContextHandle create_context_handle_from_green_ctx(const GreenCtxHandle& h_green_ctx) except+ nogil
-cdef cydriver.CUcontext ensure_context_handle(const ContextHandle& h) noexcept nogil
 cdef GreenCtxHandle get_context_green_ctx(const ContextHandle& h) noexcept nogil
 cdef GreenCtxHandle create_green_ctx_handle(
     cydriver.CUdevResourceDesc desc, cydriver.CUdevice dev, unsigned int flags) except+ nogil

cuda_core/cuda/core/_resource_handles.pyx

@@ -59,8 +59,6 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core":
         cydriver.CUcontext ctx) except+ nogil
     ContextHandle create_context_handle_from_green_ctx "cuda_core::create_context_handle_from_green_ctx" (
         const GreenCtxHandle& h_green_ctx) except+ nogil
-    cydriver.CUcontext ensure_context_handle "cuda_core::ensure_context_handle" (
-        const ContextHandle& h) noexcept nogil
     GreenCtxHandle get_context_green_ctx "cuda_core::get_context_green_ctx" (
         const ContextHandle& h) noexcept nogil
     GreenCtxHandle create_green_ctx_handle "cuda_core::create_green_ctx_handle" (

cuda_core/cuda/core/_stream.pyx

@@ -30,9 +30,7 @@ from cuda.core._resource_handles cimport (
     create_event_handle_noctx,
     create_stream_handle,
     create_stream_handle_with_owner,
-    ensure_context_handle,
     get_current_context,
-    get_last_error,
     get_legacy_stream,
     get_per_thread_stream,
     get_stream_context,
@@ -411,10 +409,6 @@ cdef inline int Stream_ensure_ctx(Stream self) except?-1 nogil:
     if not self._h_context:
         self._h_context = get_stream_context(self._h_stream)
     if self._h_context:
-        if as_cu(self._h_context) == NULL:
-            ctx = ensure_context_handle(self._h_context)
-            if ctx == NULL:
-                HANDLE_RETURN(get_last_error())
         return 0
     HANDLE_RETURN(cydriver.cuStreamGetCtx(as_cu(self._h_stream), &ctx))
     if ctx != NULL: