Skip to content

cuda.core: make Device methods use their bound context #2750

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Andy-Jost wants to merge 1 commit into
base: main
Choose a base branch
from
Open

cuda.core: make Device methods use their bound context #2750

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
591 changes: 397 additions & 194 deletions cuda_core/cuda/core/_cpp/resource_handles.cpp
View file
Open in desktop

Large diffs are not rendered by default.

48 changes: 38 additions & 10 deletions cuda_core/cuda/core/_cpp/resource_handles.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -246,6 +251,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
// ============================================================================
Expand Down Expand Up @@ -371,10 +387,11 @@ 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 and releases it with cuMemFree.
// Returns CUDA_ERROR_INVALID_CONTEXT for an empty handle.
CUresult deviceptr_alloc_raw(CUdeviceptr* ptr, size_t size,
const ContextHandle& h_context) noexcept;

// Allocate pinned host memory via cuMemAllocHost.
// When the last reference is released, cuMemFreeHost is called.
Expand Down Expand Up @@ -739,7 +756,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.
Expand All @@ -749,6 +766,9 @@ 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;

// 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.
Expand All @@ -758,27 +778,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);

// ============================================================================
Expand Down
44 changes: 20 additions & 24 deletions cuda_core/cuda/core/_device.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand All @@ -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
--------
Expand Down Expand Up @@ -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:

Expand All @@ -655,7 +660,7 @@ class Device:

Note
----
Device must be initialized.
Device must be initialized. New streams are created on this device.

Parameters
----------
Expand All @@ -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
----
Expand Down Expand Up @@ -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
-------
Expand All @@ -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
----
Expand All @@ -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
----
Expand All @@ -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
----
Expand All @@ -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
----
Expand Down
Loading
Loading