cuda.core: make Device methods use their bound context

Run context-sensitive Device operations against the Device's bound context while preserving caller state. Centralize context-aware cleanup and synchronous allocation handling so resource lifetimes remain correct.

Author: Andy-Jost

cuda_core/cuda/core/_cpp/resource_handles.cpp

@@ -64,10 +64,15 @@ void clear_last_error() noexcept;
 // function pointers extracted from cuda.bindings.cydriver.__pyx_capi__.
 // ============================================================================
 
+extern decltype(&cuGetErrorName) p_cuGetErrorName;
+extern decltype(&cuGetErrorString) p_cuGetErrorString;
+
 extern decltype(&cuDevicePrimaryCtxRetain) p_cuDevicePrimaryCtxRetain;
 extern decltype(&cuDevicePrimaryCtxRelease) p_cuDevicePrimaryCtxRelease;
 extern decltype(&cuCtxGetCurrent) p_cuCtxGetCurrent;
 extern decltype(&cuCtxSetCurrent) p_cuCtxSetCurrent;
+extern decltype(&cuCtxSynchronize) p_cuCtxSynchronize;
+extern decltype(&cuCtxGetStreamPriorityRange) p_cuCtxGetStreamPriorityRange;
 extern decltype(&cuGreenCtxCreate) p_cuGreenCtxCreate;
 extern decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy;
 extern decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx;
@@ -77,6 +82,7 @@ extern decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate;
 
 extern decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority;
 extern decltype(&cuStreamDestroy) p_cuStreamDestroy;
+extern decltype(&cuStreamSynchronize) p_cuStreamSynchronize;
 
 extern decltype(&cuEventCreate) p_cuEventCreate;
 extern decltype(&cuEventDestroy) p_cuEventDestroy;
@@ -128,6 +134,7 @@ extern decltype(&cuGraphicsUnregisterResource) p_cuGraphicsUnregisterResource;
 
 // Texture / surface / array (PR #467)
 extern decltype(&cuArray3DCreate) p_cuArray3DCreate;
+extern decltype(&cuArray3DGetDescriptor) p_cuArray3DGetDescriptor;
 extern decltype(&cuArrayDestroy) p_cuArrayDestroy;
 extern decltype(&cuMipmappedArrayCreate) p_cuMipmappedArrayCreate;
 extern decltype(&cuMipmappedArrayDestroy) p_cuMipmappedArrayDestroy;
@@ -246,6 +253,17 @@ ContextHandle get_primary_context(int device_id);
 // Returns empty handle if no context is current (caller must check)
 ContextHandle get_current_context();
 
+// Synchronize the provided context.
+// Returns CUDA_ERROR_INVALID_CONTEXT for an empty handle.
+CUresult context_synchronize(const ContextHandle& h_context) noexcept;
+
+// Query the stream priority range for the provided context.
+// Returns CUDA_ERROR_INVALID_CONTEXT for an empty handle.
+CUresult context_get_stream_priority_range(
+    const ContextHandle& h_context,
+    int* least_priority,
+    int* greatest_priority) noexcept;
+
 // ============================================================================
 // Stream handle functions
 // ============================================================================
@@ -371,10 +389,15 @@ DevicePtrHandle deviceptr_alloc_from_pool(
 // Returns empty handle on error (caller must check).
 DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream);
 
-// Allocate device memory synchronously via cuMemAlloc.
-// When the last reference is released, cuMemFree is called.
-// Returns empty handle on error (caller must check).
-DevicePtrHandle deviceptr_alloc(size_t size);
+// Allocate device memory synchronously via cuMemAlloc with the provided
+// context current. The caller owns the pointer; release it with deviceptr_free.
+// Returns CUDA_ERROR_INVALID_CONTEXT for an empty handle.
+CUresult deviceptr_alloc_raw(CUdeviceptr* ptr, size_t size,
+                             const ContextHandle& h_context) noexcept;
+
+// Synchronize the provided stream, if any, then release a synchronous
+// allocation via cuMemFree.
+CUresult deviceptr_free(CUdeviceptr ptr, const StreamHandle& h_stream) noexcept;
 
 // Allocate pinned host memory via cuMemAllocHost.
 // When the last reference is released, cuMemFreeHost is called.
@@ -739,7 +762,7 @@ FileDescriptorHandle create_fd_handle_ref(int fd);
 // Create an owning CUDA array via cuArray3DCreate.
 // When the last reference is released, cuArrayDestroy is called automatically.
 // Returns empty handle on error (caller must check).
-OpaqueArrayHandle create_array_handle(const CUDA_ARRAY3D_DESCRIPTOR& desc);
+OpaqueArrayHandle create_array_handle(const ContextHandle& h_context, const CUDA_ARRAY3D_DESCRIPTOR& desc);
 
 // Create a non-owning array handle (references an existing CUarray).
 // Use for arrays owned elsewhere (e.g. graphics interop). Never destroyed here.
@@ -749,6 +772,12 @@ OpaqueArrayHandle create_array_handle_ref(CUarray arr);
 // When the last reference is released, cuArrayDestroy is called automatically.
 OpaqueArrayHandle create_array_handle_owning(CUarray arr);
 
+// Return the context dependency associated with an array, if known.
+ContextHandle get_array_context(const OpaqueArrayHandle& h) noexcept;
+
+// Query an array descriptor.
+CUresult get_array_descriptor(const OpaqueArrayHandle& h, CUDA_ARRAY3D_DESCRIPTOR* desc) noexcept;
+
 // Create a non-owning handle to a mipmap level via cuMipmappedArrayGetLevel.
 // The level CUarray is owned by the mipmap; the parent MipmappedArrayHandle is
 // embedded in the box so it outlives the level view. No destroy in the deleter.
@@ -758,27 +787,35 @@ OpaqueArrayHandle create_array_level_handle(const MipmappedArrayHandle& h_mip, u
 // Create an owning mipmapped array via cuMipmappedArrayCreate.
 // When the last reference is released, cuMipmappedArrayDestroy is called.
 // Returns empty handle on error (caller must check).
-MipmappedArrayHandle create_mipmapped_array_handle(const CUDA_ARRAY3D_DESCRIPTOR& desc,
+MipmappedArrayHandle create_mipmapped_array_handle(const ContextHandle& h_context,
+                                                   const CUDA_ARRAY3D_DESCRIPTOR& desc,
                                                    unsigned int num_levels);
 
+// Return the context dependency associated with a mipmapped array, if known.
+ContextHandle get_mipmapped_array_context(const MipmappedArrayHandle& h) noexcept;
+
 // Create an owning texture object via cuTexObjectCreate, embedding the backing
 // resource handle (array / mipmapped array / linear-or-pitch2d device pointer)
 // so the backing always outlives the texture. cuTexObjectDestroy runs in the
 // deleter. Returns empty handle on error (caller must check).
-TexObjectHandle create_tex_object_handle_array(const CUDA_RESOURCE_DESC& res,
+TexObjectHandle create_tex_object_handle_array(const ContextHandle& h_context,
+                                               const CUDA_RESOURCE_DESC& res,
                                                const CUDA_TEXTURE_DESC& tex,
                                                const OpaqueArrayHandle& h_backing);
-TexObjectHandle create_tex_object_handle_mipmap(const CUDA_RESOURCE_DESC& res,
+TexObjectHandle create_tex_object_handle_mipmap(const ContextHandle& h_context,
+                                                const CUDA_RESOURCE_DESC& res,
                                                 const CUDA_TEXTURE_DESC& tex,
                                                 const MipmappedArrayHandle& h_backing);
-TexObjectHandle create_tex_object_handle_linear(const CUDA_RESOURCE_DESC& res,
+TexObjectHandle create_tex_object_handle_linear(const ContextHandle& h_context,
+                                                const CUDA_RESOURCE_DESC& res,
                                                 const CUDA_TEXTURE_DESC& tex,
                                                 const DevicePtrHandle& h_backing);
 
 // Create an owning surface object via cuSurfObjectCreate, embedding the backing
 // array handle so it outlives the surface. cuSurfObjectDestroy runs in the
 // deleter. Returns empty handle on error (caller must check).
-SurfObjectHandle create_surf_object_handle(const CUDA_RESOURCE_DESC& res,
+SurfObjectHandle create_surf_object_handle(const ContextHandle& h_context,
+                                           const CUDA_RESOURCE_DESC& res,
                                            const OpaqueArrayHandle& h_backing);
 
 // ============================================================================

cuda_core/cuda/core/_cpp/resource_handles.hpp

@@ -579,14 +579,17 @@ class Device:
     def memory_resource(self, mr: MemoryResource) -> None: ...
     @property
     def default_stream(self) -> Stream:
-        """Return default CUDA :obj:`~_stream.Stream` associated with this device.
+        """Return a default CUDA :obj:`~_stream.Stream` token.
 
         The type of default stream returned depends on if the environment
         variable CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM is set.
 
         If set, returns a per-thread default stream. Otherwise returns
         the legacy stream.
 
+        A default-stream token uses the device that is current when the token
+        is used.
+
         """
     def __int__(self) -> int:
         """Return device_id."""
@@ -611,7 +614,9 @@ class Device:
         Returns
         -------
         :obj:`~_context.Context`, optional
-            Popped context.
+            The previous context, or ``None`` if no context was current. When
+            returned, its ``device_id`` identifies the device that was
+            previously current.
 
         Examples
         --------
@@ -643,7 +648,7 @@ class Device:
 
         """
     def create_stream(self, obj: IsStreamType | None=None, options: StreamOptions | None=None) -> Stream:
-        """Create a :obj:`~_stream.Stream` object.
+        """Create or wrap a :obj:`~_stream.Stream` object.
 
         New stream objects can be created in two different ways:
 
@@ -655,7 +660,7 @@ class Device:
 
         Note
         ----
-        Device must be initialized.
+        Device must be initialized. New streams are created on this device.
 
         Parameters
         ----------
@@ -671,7 +676,7 @@ class Device:
 
         """
     def create_event(self, options: EventOptions | None=None) -> Event:
-        """Create an :obj:`~_event.Event` object without recording it to a :obj:`~_stream.Stream`.
+        """Create an :obj:`~_event.Event` on this device without recording it to a :obj:`~_stream.Stream`.
 
         Note
         ----
@@ -714,15 +719,15 @@ class Device:
 
         """
     def sync(self) -> None:
-        """Synchronize the device.
+        """Synchronize this device.
 
         Note
         ----
         Device must be initialized.
 
         """
     def create_graph_builder(self) -> GraphBuilder:
-        """Create a new :obj:`~graph.GraphBuilder` object.
+        """Create a new :obj:`~graph.GraphBuilder` on this device.
 
         Returns
         -------
@@ -731,12 +736,10 @@ class Device:
 
         """
     def create_opaque_array(self, options: OpaqueArrayOptions) -> OpaqueArray:
-        """Create an :obj:`~cuda.core.texture.OpaqueArray` on the current device.
+        """Create an :obj:`~cuda.core.texture.OpaqueArray` on this device.
 
         Allocates an opaque, hardware-laid-out CUDA array for texture/surface
-        access. The array is created in the current CUDA context, so make this
-        device current with :meth:`set_current` before calling (mirroring
-        :meth:`create_stream` / :meth:`create_event`).
+        access.
 
         Note
         ----
@@ -755,12 +758,10 @@ class Device:
         .. versionadded:: 1.1.0
         """
     def create_mipmapped_array(self, options: MipmappedArrayOptions) -> MipmappedArray:
-        """Create a :obj:`~cuda.core.texture.MipmappedArray` on the current device.
+        """Create a :obj:`~cuda.core.texture.MipmappedArray` on this device.
 
         Allocates a mipmapped CUDA array for texture/surface access across
-        levels. The array is created in the current CUDA context, so make this
-        device current with :meth:`set_current` before calling (mirroring
-        :meth:`create_stream` / :meth:`create_event`).
+        levels.
 
         Note
         ----
@@ -779,15 +780,13 @@ class Device:
         .. versionadded:: 1.1.0
         """
     def create_texture_object(self, *, resource: ResourceDescriptor, options: TextureObjectOptions | None=None) -> TextureObject:
-        """Create a :obj:`~cuda.core.texture.TextureObject` on the current device.
+        """Create a :obj:`~cuda.core.texture.TextureObject` on this device.
 
         Binds a resource (an :obj:`~cuda.core.texture.OpaqueArray` /
         :obj:`~cuda.core.texture.MipmappedArray` / linear or pitch2d
         :obj:`~cuda.core.Buffer`, wrapped in a
         :obj:`~cuda.core.texture.ResourceDescriptor`) as a bindless texture for
-        kernel-side sampled reads. The object is created in the current CUDA
-        context, so make this device current with :meth:`set_current` before
-        calling (mirroring :meth:`create_stream` / :meth:`create_event`).
+        kernel-side sampled reads. The resource must belong to this device.
 
         Note
         ----
@@ -808,15 +807,12 @@ class Device:
         .. versionadded:: 1.1.0
         """
     def create_surface_object(self, *, resource: ResourceDescriptor) -> SurfaceObject:
-        """Create a :obj:`~cuda.core.texture.SurfaceObject` on the current device.
+        """Create a :obj:`~cuda.core.texture.SurfaceObject` on this device.
 
         Binds an :obj:`~cuda.core.texture.OpaqueArray` (via a
         :obj:`~cuda.core.texture.ResourceDescriptor`) as a bindless surface for
         kernel-side typed load/store. The backing array must have been created
-        with ``is_surface_load_store=True``. The object is created in the
-        current CUDA context, so make this device current with
-        :meth:`set_current` before calling (mirroring :meth:`create_stream` /
-        :meth:`create_event`).
+        with ``is_surface_load_store=True`` and must belong to this device.
 
         Note
         ----