Skip to content

Commit f84a7cf

Browse files
Andy-Jostclaude
Andy-Jost
and
claude
authored
cuda.core: make Device methods use their bound context (#2750)
* 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. * cuda.core: address review of #2750 - Fix _SynchronousMemoryResource to record a deallocation token bound to its own context, so Buffer teardown enters the right context regardless of what is current, and works with no context current. Move the class to its own module and resolve the primary context lazily. - LegacyPinnedMemoryResource.device_id returns -1 as documented; texture creation over a pinned buffer works again. - Device.set_current delegates a foreign-device context to its owning device instead of raising, so the save/restore idiom round-trips across devices. - Guard empty context handles once in invoke_in_context(_or_undo); warn when an undo is skipped because the target context is gone. - context_synchronize and context_get_stream_priority_range release the GIL like the other helpers. Rename array/mipmap box accessors to get_box. - Tests: query the driver (cuStreamGetCtx, cross-context cuEventRecord) instead of comparing cached metadata; add sync(), set_current round-trip, pinned texture, and synchronous-resource teardown tests; register device_x2 with the parallel-test plugin. - Move the release note to 1.3.0 and describe sync() as acting on the bound context. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * cuda.core tests: fix driver-call plumbing in the bound-context helper handle_return() takes the whole result tuple, and cuCtxGetDevice() returns a CUdevice that never compares equal to an int; both made the cross-context event-record check raise instead of run (or skip) as intended. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * cuda.core: regenerate stub with stubgen-pyx 0.2.22 Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * cuda.core: create stream-ordering events in the recorded stream's context The uniform empty-context guard added for review item 6 broke create_event_handle_noctx, which relied on an empty handle meaning "current context". That helper was itself a #2311-class bug: Stream.wait(stream) and the foreign-array/tensor import paths created their temporary ordering event in whatever context was current, and cuEventRecord rejects an event from a different context than the stream it is recorded on, so cross-device waits failed unless the right device happened to be current. Replace it with create_event_handle_for_stream, which resolves the stream's owning context via cuStreamGetCtx and creates the event there. Callers now check the returned handle and surface the real creation error. With that, every creation helper requires a context and no exception remains. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
1 parent f48022a commit f84a7cf

34 files changed

Lines changed: 1511 additions & 480 deletions

‎cuda_core/cuda/core/_cpp/resource_handles.cpp‎

Lines changed: 429 additions & 196 deletions
Large diffs are not rendered by default.

‎cuda_core/cuda/core/_cpp/resource_handles.hpp‎

Lines changed: 53 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -64,10 +64,15 @@ void clear_last_error() noexcept;
6464
// function pointers extracted from cuda.bindings.cydriver.__pyx_capi__.
6565
// ============================================================================
6666

67+
extern decltype(&cuGetErrorName) p_cuGetErrorName;
68+
extern decltype(&cuGetErrorString) p_cuGetErrorString;
69+
6770
extern decltype(&cuDevicePrimaryCtxRetain) p_cuDevicePrimaryCtxRetain;
6871
extern decltype(&cuDevicePrimaryCtxRelease) p_cuDevicePrimaryCtxRelease;
6972
extern decltype(&cuCtxGetCurrent) p_cuCtxGetCurrent;
7073
extern decltype(&cuCtxSetCurrent) p_cuCtxSetCurrent;
74+
extern decltype(&cuCtxSynchronize) p_cuCtxSynchronize;
75+
extern decltype(&cuCtxGetStreamPriorityRange) p_cuCtxGetStreamPriorityRange;
7176
extern decltype(&cuGreenCtxCreate) p_cuGreenCtxCreate;
7277
extern decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy;
7378
extern decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx;
@@ -77,6 +82,7 @@ extern decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate;
7782

7883
extern decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority;
7984
extern decltype(&cuStreamDestroy) p_cuStreamDestroy;
85+
extern decltype(&cuStreamGetCtx) p_cuStreamGetCtx;
8086

8187
extern decltype(&cuEventCreate) p_cuEventCreate;
8288
extern decltype(&cuEventDestroy) p_cuEventDestroy;
@@ -246,6 +252,17 @@ ContextHandle get_primary_context(int device_id);
246252
// Returns empty handle if no context is current (caller must check)
247253
ContextHandle get_current_context();
248254

255+
// Synchronize the provided context. Releases the GIL around the driver call.
256+
// Returns CUDA_ERROR_INVALID_CONTEXT for an empty handle.
257+
CUresult context_synchronize(const ContextHandle& h_context) noexcept;
258+
259+
// Query the stream priority range for the provided context.
260+
// Returns CUDA_ERROR_INVALID_CONTEXT for an empty handle.
261+
CUresult context_get_stream_priority_range(
262+
const ContextHandle& h_context,
263+
int* least_priority,
264+
int* greatest_priority) noexcept;
265+
249266
// ============================================================================
250267
// Stream handle functions
251268
// ============================================================================
@@ -287,6 +304,14 @@ StreamHandle get_legacy_stream();
287304
// Note: Per-thread stream has no specific context dependency.
288305
StreamHandle get_per_thread_stream();
289306

307+
// Wrap CU_STREAM_LEGACY with an explicit context, bypassing the "bind to
308+
// whatever is current" resolution that a bare default-stream token uses (see
309+
// make_deallocation_stream). Lets a resource that always operates in one
310+
// known context (e.g. a synchronous, non-pooled allocator) record a correct
311+
// deallocation context without requiring that context to be current when the
312+
// token is created. Returns an empty handle for an empty h_context.
313+
StreamHandle create_context_bound_legacy_stream(const ContextHandle& h_context);
314+
290315
// ============================================================================
291316
// Event handle functions
292317
// ============================================================================
@@ -300,11 +325,14 @@ EventHandle create_event_handle(const ContextHandle& h_ctx, unsigned int flags,
300325
bool timing_enabled, bool is_blocking_sync,
301326
bool ipc_enabled, int device_id);
302327

303-
// Create an owning event handle without context dependency.
304-
// Use for temporary events that are created and destroyed in the same scope.
328+
// Create an owning event in the context that owns `stream`, so it can be
329+
// recorded on that stream regardless of which context is current. Default-
330+
// stream tokens resolve to the current context (cuStreamGetCtx semantics).
331+
// Use for temporary ordering events that are created and destroyed in the
332+
// same scope; the handle carries no device id.
305333
// When the last reference is released, cuEventDestroy is called automatically.
306334
// Returns empty handle on error (caller must check).
307-
EventHandle create_event_handle_noctx(unsigned int flags);
335+
EventHandle create_event_handle_for_stream(CUstream stream, unsigned int flags);
308336

309337
// Create an owning event handle from an IPC handle.
310338
// The originating process owns the event and its context.
@@ -371,10 +399,11 @@ DevicePtrHandle deviceptr_alloc_from_pool(
371399
// Returns empty handle on error (caller must check).
372400
DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream);
373401

374-
// Allocate device memory synchronously via cuMemAlloc.
375-
// When the last reference is released, cuMemFree is called.
376-
// Returns empty handle on error (caller must check).
377-
DevicePtrHandle deviceptr_alloc(size_t size);
402+
// Allocate device memory synchronously via cuMemAlloc with the provided
403+
// context current. The caller owns the pointer and releases it with cuMemFree.
404+
// Returns CUDA_ERROR_INVALID_CONTEXT for an empty handle.
405+
CUresult deviceptr_alloc_raw(CUdeviceptr* ptr, size_t size,
406+
const ContextHandle& h_context) noexcept;
378407

379408
// Allocate pinned host memory via cuMemAllocHost.
380409
// When the last reference is released, cuMemFreeHost is called.
@@ -739,7 +768,7 @@ FileDescriptorHandle create_fd_handle_ref(int fd);
739768
// Create an owning CUDA array via cuArray3DCreate.
740769
// When the last reference is released, cuArrayDestroy is called automatically.
741770
// Returns empty handle on error (caller must check).
742-
OpaqueArrayHandle create_array_handle(const CUDA_ARRAY3D_DESCRIPTOR& desc);
771+
OpaqueArrayHandle create_array_handle(const ContextHandle& h_context, const CUDA_ARRAY3D_DESCRIPTOR& desc);
743772

744773
// Create a non-owning array handle (references an existing CUarray).
745774
// Use for arrays owned elsewhere (e.g. graphics interop). Never destroyed here.
@@ -749,6 +778,9 @@ OpaqueArrayHandle create_array_handle_ref(CUarray arr);
749778
// When the last reference is released, cuArrayDestroy is called automatically.
750779
OpaqueArrayHandle create_array_handle_owning(CUarray arr);
751780

781+
// Return the context dependency associated with an array, if known.
782+
ContextHandle get_array_context(const OpaqueArrayHandle& h) noexcept;
783+
752784
// Create a non-owning handle to a mipmap level via cuMipmappedArrayGetLevel.
753785
// The level CUarray is owned by the mipmap; the parent MipmappedArrayHandle is
754786
// embedded in the box so it outlives the level view. No destroy in the deleter.
@@ -758,27 +790,35 @@ OpaqueArrayHandle create_array_level_handle(const MipmappedArrayHandle& h_mip, u
758790
// Create an owning mipmapped array via cuMipmappedArrayCreate.
759791
// When the last reference is released, cuMipmappedArrayDestroy is called.
760792
// Returns empty handle on error (caller must check).
761-
MipmappedArrayHandle create_mipmapped_array_handle(const CUDA_ARRAY3D_DESCRIPTOR& desc,
793+
MipmappedArrayHandle create_mipmapped_array_handle(const ContextHandle& h_context,
794+
const CUDA_ARRAY3D_DESCRIPTOR& desc,
762795
unsigned int num_levels);
763796

797+
// Return the context dependency associated with a mipmapped array, if known.
798+
ContextHandle get_mipmapped_array_context(const MipmappedArrayHandle& h) noexcept;
799+
764800
// Create an owning texture object via cuTexObjectCreate, embedding the backing
765801
// resource handle (array / mipmapped array / linear-or-pitch2d device pointer)
766802
// so the backing always outlives the texture. cuTexObjectDestroy runs in the
767803
// deleter. Returns empty handle on error (caller must check).
768-
TexObjectHandle create_tex_object_handle_array(const CUDA_RESOURCE_DESC& res,
804+
TexObjectHandle create_tex_object_handle_array(const ContextHandle& h_context,
805+
const CUDA_RESOURCE_DESC& res,
769806
const CUDA_TEXTURE_DESC& tex,
770807
const OpaqueArrayHandle& h_backing);
771-
TexObjectHandle create_tex_object_handle_mipmap(const CUDA_RESOURCE_DESC& res,
808+
TexObjectHandle create_tex_object_handle_mipmap(const ContextHandle& h_context,
809+
const CUDA_RESOURCE_DESC& res,
772810
const CUDA_TEXTURE_DESC& tex,
773811
const MipmappedArrayHandle& h_backing);
774-
TexObjectHandle create_tex_object_handle_linear(const CUDA_RESOURCE_DESC& res,
812+
TexObjectHandle create_tex_object_handle_linear(const ContextHandle& h_context,
813+
const CUDA_RESOURCE_DESC& res,
775814
const CUDA_TEXTURE_DESC& tex,
776815
const DevicePtrHandle& h_backing);
777816

778817
// Create an owning surface object via cuSurfObjectCreate, embedding the backing
779818
// array handle so it outlives the surface. cuSurfObjectDestroy runs in the
780819
// deleter. Returns empty handle on error (caller must check).
781-
SurfObjectHandle create_surf_object_handle(const CUDA_RESOURCE_DESC& res,
820+
SurfObjectHandle create_surf_object_handle(const ContextHandle& h_context,
821+
const CUDA_RESOURCE_DESC& res,
782822
const OpaqueArrayHandle& h_backing);
783823

784824
// ============================================================================

‎cuda_core/cuda/core/_device.pyi‎

Lines changed: 31 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -583,14 +583,17 @@ class Device:
583583
def memory_resource(self, mr: MemoryResource) -> None: ...
584584
@property
585585
def default_stream(self) -> Stream:
586-
"""Return default CUDA :obj:`~_stream.Stream` associated with this device.
586+
"""Return a default CUDA :obj:`~_stream.Stream` token.
587587
588588
The type of default stream returned depends on if the environment
589589
variable CUDA_PYTHON_CUDA_PER_THREAD_DEFAULT_STREAM is set.
590590
591591
If set, returns a per-thread default stream. Otherwise returns
592592
the legacy stream.
593593
594+
A default-stream token uses the device that is current when the token
595+
is used.
596+
594597
"""
595598
def __int__(self) -> int:
596599
"""Return device_id."""
@@ -607,6 +610,12 @@ class Device:
607610
608611
Providing a `ctx` causes the previous set context to be popped and returned.
609612
613+
If `ctx` was created on a different device than this receiver, the call
614+
is delegated to that device's own :meth:`set_current`. This keeps the
615+
owning device's bookkeeping consistent and lets a context this method
616+
handed out for a foreign device be pushed back through any ``Device``
617+
object, matching the CUDA context stack's own thread-wide semantics.
618+
610619
Parameters
611620
----------
612621
ctx : :obj:`~_context.Context`, optional
@@ -615,7 +624,9 @@ class Device:
615624
Returns
616625
-------
617626
:obj:`~_context.Context`, optional
618-
Popped context.
627+
The previous context, or ``None`` if no context was current. When
628+
returned, its ``device_id`` identifies the device that was
629+
previously current.
619630
620631
Examples
621632
--------
@@ -647,7 +658,7 @@ class Device:
647658
648659
"""
649660
def create_stream(self, obj: IsStreamType | None=None, options: StreamOptions | None=None) -> Stream:
650-
"""Create a :obj:`~_stream.Stream` object.
661+
"""Create or wrap a :obj:`~_stream.Stream` object.
651662
652663
New stream objects can be created in two different ways:
653664
@@ -659,7 +670,7 @@ class Device:
659670
660671
Note
661672
----
662-
Device must be initialized.
673+
Device must be initialized. New streams are created on this device.
663674
664675
Parameters
665676
----------
@@ -675,7 +686,7 @@ class Device:
675686
676687
"""
677688
def create_event(self, options: EventOptions | None=None) -> Event:
678-
"""Create an :obj:`~_event.Event` object without recording it to a :obj:`~_stream.Stream`.
689+
"""Create an :obj:`~_event.Event` on this device without recording it to a :obj:`~_stream.Stream`.
679690
680691
Note
681692
----
@@ -718,15 +729,20 @@ class Device:
718729
719730
"""
720731
def sync(self) -> None:
721-
"""Synchronize the device.
732+
"""Synchronize this device's bound context.
733+
734+
Waits for all preceding work in this device's bound :obj:`~_context.Context`
735+
to complete. Only that context is synchronized, not the device as a
736+
whole; work queued in a different context on the same device (e.g. a
737+
green context) is unaffected.
722738
723739
Note
724740
----
725741
Device must be initialized.
726742
727743
"""
728744
def create_graph_builder(self) -> GraphBuilder:
729-
"""Create a new :obj:`~graph.GraphBuilder` object.
745+
"""Create a new :obj:`~graph.GraphBuilder` on this device.
730746
731747
Returns
732748
-------
@@ -735,12 +751,10 @@ class Device:
735751
736752
"""
737753
def create_opaque_array(self, options: OpaqueArrayOptions) -> OpaqueArray:
738-
"""Create an :obj:`~cuda.core.texture.OpaqueArray` on the current device.
754+
"""Create an :obj:`~cuda.core.texture.OpaqueArray` on this device.
739755
740756
Allocates an opaque, hardware-laid-out CUDA array for texture/surface
741-
access. The array is created in the current CUDA context, so make this
742-
device current with :meth:`set_current` before calling (mirroring
743-
:meth:`create_stream` / :meth:`create_event`).
757+
access.
744758
745759
Note
746760
----
@@ -759,12 +773,10 @@ class Device:
759773
.. versionadded:: 1.1.0
760774
"""
761775
def create_mipmapped_array(self, options: MipmappedArrayOptions) -> MipmappedArray:
762-
"""Create a :obj:`~cuda.core.texture.MipmappedArray` on the current device.
776+
"""Create a :obj:`~cuda.core.texture.MipmappedArray` on this device.
763777
764778
Allocates a mipmapped CUDA array for texture/surface access across
765-
levels. The array is created in the current CUDA context, so make this
766-
device current with :meth:`set_current` before calling (mirroring
767-
:meth:`create_stream` / :meth:`create_event`).
779+
levels.
768780
769781
Note
770782
----
@@ -783,15 +795,13 @@ class Device:
783795
.. versionadded:: 1.1.0
784796
"""
785797
def create_texture_object(self, *, resource: ResourceDescriptor, options: TextureObjectOptions | None=None) -> TextureObject:
786-
"""Create a :obj:`~cuda.core.texture.TextureObject` on the current device.
798+
"""Create a :obj:`~cuda.core.texture.TextureObject` on this device.
787799
788800
Binds a resource (an :obj:`~cuda.core.texture.OpaqueArray` /
789801
:obj:`~cuda.core.texture.MipmappedArray` / linear or pitch2d
790802
:obj:`~cuda.core.Buffer`, wrapped in a
791803
:obj:`~cuda.core.texture.ResourceDescriptor`) as a bindless texture for
792-
kernel-side sampled reads. The object is created in the current CUDA
793-
context, so make this device current with :meth:`set_current` before
794-
calling (mirroring :meth:`create_stream` / :meth:`create_event`).
804+
kernel-side sampled reads. The resource must belong to this device.
795805
796806
Note
797807
----
@@ -812,15 +822,12 @@ class Device:
812822
.. versionadded:: 1.1.0
813823
"""
814824
def create_surface_object(self, *, resource: ResourceDescriptor) -> SurfaceObject:
815-
"""Create a :obj:`~cuda.core.texture.SurfaceObject` on the current device.
825+
"""Create a :obj:`~cuda.core.texture.SurfaceObject` on this device.
816826
817827
Binds an :obj:`~cuda.core.texture.OpaqueArray` (via a
818828
:obj:`~cuda.core.texture.ResourceDescriptor`) as a bindless surface for
819829
kernel-side typed load/store. The backing array must have been created
820-
with ``is_surface_load_store=True``. The object is created in the
821-
current CUDA context, so make this device current with
822-
:meth:`set_current` before calling (mirroring :meth:`create_stream` /
823-
:meth:`create_event`).
830+
with ``is_surface_load_store=True`` and must belong to this device.
824831
825832
Note
826833
----

0 commit comments

Comments
 (0)