diff --git a/cuda_core/cuda/core/_cpp/resource_handles.cpp b/cuda_core/cuda/core/_cpp/resource_handles.cpp index ee116a9f353..ee4aef8e2e9 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.cpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.cpp @@ -12,12 +12,15 @@ #include #include #include +#include #include #include #include #include #include +#include #include +#include #include #ifndef _WIN32 @@ -33,10 +36,15 @@ namespace cuda_core { // function pointers extracted from cuda.bindings.cydriver.__pyx_capi__. // ============================================================================ +decltype(&cuGetErrorName) p_cuGetErrorName = nullptr; +decltype(&cuGetErrorString) p_cuGetErrorString = nullptr; + decltype(&cuDevicePrimaryCtxRetain) p_cuDevicePrimaryCtxRetain = nullptr; decltype(&cuDevicePrimaryCtxRelease) p_cuDevicePrimaryCtxRelease = nullptr; decltype(&cuCtxGetCurrent) p_cuCtxGetCurrent = nullptr; decltype(&cuCtxSetCurrent) p_cuCtxSetCurrent = nullptr; +decltype(&cuCtxSynchronize) p_cuCtxSynchronize = nullptr; +decltype(&cuCtxGetStreamPriorityRange) p_cuCtxGetStreamPriorityRange = nullptr; decltype(&cuGreenCtxCreate) p_cuGreenCtxCreate = nullptr; decltype(&cuGreenCtxDestroy) p_cuGreenCtxDestroy = nullptr; decltype(&cuCtxFromGreenCtx) p_cuCtxFromGreenCtx = nullptr; @@ -128,46 +136,34 @@ NvvmDestroyProgramFn p_nvvmDestroyProgram = nullptr; NvJitLinkDestroyFn p_nvJitLinkDestroy = nullptr; // ============================================================================ -// GIL management helpers +// GIL and scoped-context management helpers // ============================================================================ namespace { -// Helper to release the GIL while calling into the CUDA driver. -// This guard is *conditional*: if the caller already dropped the GIL, -// we avoid calling PyEval_SaveThread (which requires holding the GIL). -// It also handles the case where Python is finalizing and GIL operations -// are no longer safe. +// Conditionally release the GIL while calling into the CUDA driver. class GILReleaseGuard { public: - GILReleaseGuard() : tstate_(nullptr), released_(false) { - // Don't try to manipulate GIL if Python is finalizing + GILReleaseGuard() noexcept { if (!Py_IsInitialized() || py_is_finalizing()) { return; } - // PyGILState_Check() returns 1 if the GIL is held by this thread. if (PyGILState_Check()) { tstate_ = PyEval_SaveThread(); - released_ = true; } - // Note: If the GIL is not released (finalizing, or not held): - // - Reduces parallelism (other Python threads remain blocked) - // - No deadlock risk as long as the guarded code doesn't call back into Python } ~GILReleaseGuard() { - if (released_) { + if (tstate_) { PyEval_RestoreThread(tstate_); } } - // Non-copyable, non-movable GILReleaseGuard(const GILReleaseGuard&) = delete; GILReleaseGuard& operator=(const GILReleaseGuard&) = delete; private: - PyThreadState* tstate_; - bool released_; + PyThreadState* tstate_ = nullptr; }; // Helper to acquire the GIL when we might not hold it. @@ -200,55 +196,228 @@ class GILAcquireGuard { bool acquired_; }; -// Temporarily make a context current, restoring the caller's prior binding -// (including having no context current) on scope exit. The handle is held for -// the duration so the context cannot be destroyed mid-scope. -class ScopedCurrentContext { -public: - explicit ScopedCurrentContext(ContextHandle h_context) noexcept - : h_context_(std::move(h_context)) { - CUcontext target = as_cu(h_context_); - if (!target) { - return; - } +void warn_on_cuda_error(const char* operation, CUresult status, const char* detail = nullptr) noexcept; + +// Make a context current and record the state needed to restore it. +// An empty handle is a no-op: the operation runs in the caller's current +// context, and nothing is restored on exit. +CUresult enter_context(const ContextHandle& h_context, CUcontext* previous, int* changed) noexcept { + *previous = nullptr; + *changed = 0; + CUcontext target = as_cu(h_context); + if (!target) { + return CUDA_SUCCESS; + } + + GILReleaseGuard gil; + CUresult status = p_cuCtxGetCurrent(previous); + if (status != CUDA_SUCCESS || *previous == target) { + return status; + } + status = p_cuCtxSetCurrent(target); + *changed = status == CUDA_SUCCESS; + return status; +} +// Restore the previous context and preserve an earlier operation error. +CUresult exit_context(CUcontext previous, int changed, CUresult operation_status) noexcept { + CUresult restore_status = CUDA_SUCCESS; + if (changed) { GILReleaseGuard gil; - status_ = p_cuCtxGetCurrent(&previous_); - if (status_ != CUDA_SUCCESS || previous_ == target) { - return; + restore_status = p_cuCtxSetCurrent(previous); + } + if (operation_status != CUDA_SUCCESS && restore_status != CUDA_SUCCESS) { + warn_on_cuda_error("cuCtxSetCurrent (restoring the caller's context)", restore_status); + } + return operation_status != CUDA_SUCCESS ? operation_status : restore_status; +} + +// Require a callable to be invocable without throwing. +#define ASSERT_NOTHROW_INVOCABLE(...) \ + static_assert(std::is_nothrow_invocable_v<__VA_ARGS__>, "operation must be noexcept") + +// Store a stream and any state needed to preserve deallocation ordering. +struct DeallocationStream { + StreamHandle h_stream; + std::thread::id ptds_tid{}; +}; + +// Return whether a stream handle needs a current context to resolve it. +bool is_default_stream(CUstream stream) noexcept { + return stream == nullptr || stream == CU_STREAM_LEGACY || stream == CU_STREAM_PER_THREAD; +} + +// Return the context a deallocation-stream token must run under. Real streams +// resolve their own context; default-stream tokens use the context bound at +// allocation time. Warn when PTDS deallocation crosses host threads. +ContextHandle deallocation_context(const DeallocationStream& stream) noexcept { + if (!is_default_stream(as_cu(stream.h_stream))) { + return {}; + } + if (stream.ptds_tid != std::thread::id{} + && stream.ptds_tid != std::this_thread::get_id()) { + std::fprintf( + stderr, + "Warning: Buffer deallocation for a per-thread default stream " + "is running on a different host thread than the one that recorded " + "the deallocation stream; ordering relative to the allocating " + "thread's PTDS is not preserved\n"); + } + return get_stream_context(stream.h_stream); +} + +// Run an operation with the requested context current. +template +CUresult invoke_in_context(const ContextHandle& h_context, Fn&& operation, Args&&... args) noexcept { + ASSERT_NOTHROW_INVOCABLE(Fn&&, Args&&...); + CUcontext previous = nullptr; + int changed = 0; + CUresult status = enter_context(h_context, &previous, &changed); + if (status == CUDA_SUCCESS) { + status = std::invoke(std::forward(operation), std::forward(args)...); + } + return exit_context(previous, changed, status); +} + +// Run a creation operation and undo it if context restoration fails. +// Context-independent undo always runs. Context-sensitive undo runs only +// after verifying that the target context remains current; otherwise the +// resource leaks rather than risking cleanup in the wrong context. +template +CUresult invoke_in_context_or_undo(const ContextHandle& h_context, Fn&& operation, + Undo&& undo, bool undo_requires_target_context) noexcept { + ASSERT_NOTHROW_INVOCABLE(Fn&&); + ASSERT_NOTHROW_INVOCABLE(Undo&&); + CUcontext previous = nullptr; + int changed = 0; + CUresult status = enter_context(h_context, &previous, &changed); + if (status != CUDA_SUCCESS) { + return status; + } + status = std::invoke(std::forward(operation)); + CUresult composite = exit_context(previous, changed, status); + if (status == CUDA_SUCCESS && composite != CUDA_SUCCESS) { + bool undo_ok = true; + if (undo_requires_target_context) { + CUcontext current = nullptr; + undo_ok = p_cuCtxGetCurrent(¤t) == CUDA_SUCCESS + && current == as_cu(h_context); + } + if (undo_ok) { + std::invoke(std::forward(undo)); } - status_ = p_cuCtxSetCurrent(target); - changed_ = status_ == CUDA_SUCCESS; } + return composite; +} - ~ScopedCurrentContext() { - if (changed_) { - GILReleaseGuard gil; - CUresult status = p_cuCtxSetCurrent(previous_); - if (status != CUDA_SUCCESS) { - std::fprintf( - stderr, - "Warning: cuCtxSetCurrent (restoring the caller's context) " - "failed (CUDA error %d)\n", - static_cast(status)); - } +// Write a warning that includes the CUDA error name and description. +void warn_on_cuda_error(const char* operation, CUresult status, const char* detail) noexcept { + const char* error_name = nullptr; + const char* error_description = nullptr; + CUresult name_status = p_cuGetErrorName(status, &error_name); + CUresult description_status = p_cuGetErrorString(status, &error_description); + + if (name_status == CUDA_SUCCESS && description_status == CUDA_SUCCESS) { + if (detail) { + std::fprintf(stderr, "Warning: %s %s: %s: %s\n", + operation, detail, error_name, error_description); + } else { + std::fprintf(stderr, "Warning: %s failed: %s: %s\n", + operation, error_name, error_description); + } + } else { + if (detail) { + std::fprintf(stderr, "Warning: %s %s (CUDA error %d)\n", + operation, detail, static_cast(status)); + } else { + std::fprintf(stderr, "Warning: %s failed (CUDA error %d)\n", + operation, static_cast(status)); + } + } +} + +// Run cleanup with the requested context current. Warn and skip the operation +// if activation fails, and independently warn on operation or restoration +// failure. Return the operation or activation status; restoration never +// changes the return value. +template +CUresult cleanup_in_context(const ContextHandle& h_context, const char* name, + Fn&& operation, Args&&... args) noexcept { + ASSERT_NOTHROW_INVOCABLE(Fn&&, Args&&...); + CUcontext previous = nullptr; + int changed = 0; + CUresult status = enter_context(h_context, &previous, &changed); + if (status != CUDA_SUCCESS) { + warn_on_cuda_error(name, status, + "skipped (context activation failed; resource leaked)"); + } else { + status = std::invoke(std::forward(operation), std::forward(args)...); + if (status != CUDA_SUCCESS) { + warn_on_cuda_error(name, status); } } + CUresult restore = exit_context(previous, changed, CUDA_SUCCESS); + if (restore != CUDA_SUCCESS) { + warn_on_cuda_error(name, restore, "failed while restoring the caller's context"); + } + return status; +} + +#undef ASSERT_NOTHROW_INVOCABLE - CUresult status() const noexcept { return status_; } +// Decorate a CUDA operation to warn whenever it returns an error. +template +class WarnOnFailure { +public: + explicit WarnOnFailure(const char* operation) noexcept : operation_(operation) {} - ScopedCurrentContext(const ScopedCurrentContext&) = delete; - ScopedCurrentContext& operator=(const ScopedCurrentContext&) = delete; + template + CUresult operator()(Args&&... args) const noexcept { + CUresult status = Function(std::forward(args)...); + if (status != CUDA_SUCCESS) { + warn_on_cuda_error(operation_, status); + } + return status; + } private: - ContextHandle h_context_; - CUcontext previous_ = nullptr; - bool changed_ = false; - CUresult status_ = CUDA_SUCCESS; + const char* operation_; }; +// Warning-decorated CUDA operations used by non-throwing cleanup paths. +const WarnOnFailure pw_cuStreamDestroy{"cuStreamDestroy"}; +const WarnOnFailure pw_cuEventDestroy{"cuEventDestroy"}; +const WarnOnFailure pw_cuMemFree{"cuMemFree"}; +const WarnOnFailure pw_cuMemFreeAsync{"cuMemFreeAsync"}; +const WarnOnFailure pw_cuArrayDestroy{"cuArrayDestroy"}; +const WarnOnFailure pw_cuMipmappedArrayDestroy{"cuMipmappedArrayDestroy"}; +const WarnOnFailure pw_cuTexObjectDestroy{"cuTexObjectDestroy"}; +const WarnOnFailure pw_cuSurfObjectDestroy{"cuSurfObjectDestroy"}; + } // namespace +// Synchronize the provided context. +CUresult context_synchronize(const ContextHandle& h_context) noexcept { + if (!h_context) { + return CUDA_ERROR_INVALID_CONTEXT; + } + return invoke_in_context(h_context, []() noexcept { + return p_cuCtxSynchronize(); + }); +} + +// Query the stream priority range for the provided context. +CUresult context_get_stream_priority_range(const ContextHandle& h_context, + int* least_priority, + int* greatest_priority) noexcept { + if (!h_context) { + return CUDA_ERROR_INVALID_CONTEXT; + } + return invoke_in_context(h_context, [&]() noexcept { + return p_cuCtxGetStreamPriorityRange(least_priority, greatest_priority); + }); +} + // ============================================================================ // CUDA user-object deferred cleanup // @@ -478,12 +647,14 @@ class HandleRegistry { // Thread-local status of the most recent CUDA API call in this module. static thread_local CUresult err = CUDA_SUCCESS; +// Return and clear the calling thread's most recent CUDA error. CUresult get_last_error() noexcept { CUresult e = err; err = CUDA_SUCCESS; return e; } +// Return the calling thread's most recent CUDA error without clearing it. CUresult peek_last_error() noexcept { return err; } @@ -676,22 +847,21 @@ static HandleRegistry stream_registry; StreamHandle create_stream_handle(const ContextHandle& h_ctx, unsigned int flags, int priority) { GILReleaseGuard gil; - CUstream stream; - - // Dispatch: green context uses cuGreenCtxStreamCreate, primary uses cuStreamCreateWithPriority + CUstream stream = nullptr; GreenCtxHandle h_green = get_context_green_ctx(h_ctx); if (h_green) { - if (!p_cuGreenCtxStreamCreate) { - err = CUDA_ERROR_NOT_SUPPORTED; - return {}; - } - if (CUDA_SUCCESS != (err = p_cuGreenCtxStreamCreate(&stream, as_cu(h_green), flags, priority))) { - return {}; - } + err = p_cuGreenCtxStreamCreate + ? p_cuGreenCtxStreamCreate(&stream, as_cu(h_green), flags, priority) + : CUDA_ERROR_NOT_SUPPORTED; } else { - if (CUDA_SUCCESS != (err = p_cuStreamCreateWithPriority(&stream, flags, priority))) { - return {}; - } + err = invoke_in_context_or_undo( + h_ctx, + [&]() noexcept { return p_cuStreamCreateWithPriority(&stream, flags, priority); }, + [&]() noexcept { pw_cuStreamDestroy(stream); }, + /*undo_requires_target_context=*/false); + } + if (err != CUDA_SUCCESS) { + return {}; } auto box = std::shared_ptr( @@ -699,7 +869,7 @@ StreamHandle create_stream_handle(const ContextHandle& h_ctx, unsigned int flags [](const StreamBox* b) { stream_registry.unregister_handle(b->resource); GILReleaseGuard gil; - p_cuStreamDestroy(b->resource); + pw_cuStreamDestroy(b->resource); delete b; } ); @@ -769,6 +939,7 @@ void py_object_user_object_destroy(void* py_object) noexcept { Py_DECREF(reinterpret_cast(py_object)); } +// Return the context retained by a stream handle. ContextHandle get_stream_context(const StreamHandle& h) noexcept { return h ? get_box(h)->h_context : ContextHandle{}; } @@ -797,12 +968,6 @@ StreamHandle get_per_thread_stream() { // detected and warnings can be issued. // ============================================================================ -// ptds_tid is std::thread::id{} except for CU_STREAM_PER_THREAD. -struct DeallocationStream { - StreamHandle h_stream; - std::thread::id ptds_tid{}; -}; - // Real streams are copied unchanged. Default-stream tokens without an embedded // context are bound to the current context. Returns false (and sets err) when a // default-stream token cannot be bound because no context is current. @@ -814,9 +979,7 @@ static bool make_deallocation_stream( } const CUstream stream = as_cu(h); - if (stream != nullptr - && stream != CU_STREAM_LEGACY - && stream != CU_STREAM_PER_THREAD) { + if (!is_default_stream(stream)) { out = DeallocationStream{h, {}}; return true; } @@ -845,35 +1008,6 @@ static bool make_deallocation_stream( return true; } -template -CUresult with_deallocation_context( - const DeallocationStream& stream, - const char* operation, - Fn&& fn) noexcept { - if (stream.ptds_tid != std::thread::id{} - && stream.ptds_tid != std::this_thread::get_id()) { - std::fprintf( - stderr, - "Warning: Buffer deallocation for a per-thread default stream " - "is running on a different host thread than the one that recorded " - "the deallocation stream; ordering relative to the allocating " - "thread's PTDS is not preserved\n"); - } - ScopedCurrentContext context(get_stream_context(stream.h_stream)); - CUresult status = context.status(); - if (status == CUDA_SUCCESS) { - status = fn(stream); - } - if (status != CUDA_SUCCESS) { - std::fprintf( - stderr, - "Warning: %s failed during resource destruction (CUDA error %d)\n", - operation, - static_cast(status)); - } - return status; -} - // ============================================================================ // Event Handles // ============================================================================ @@ -912,6 +1046,7 @@ int get_event_device_id(const EventHandle& h) noexcept { return h ? get_box(h)->device_id : -1; } +// Return the context retained by an event handle. ContextHandle get_event_context(const EventHandle& h) noexcept { return h ? get_box(h)->h_context : ContextHandle{}; } @@ -923,17 +1058,22 @@ EventHandle create_event_handle(const ContextHandle& h_ctx, unsigned int flags, bool timing_enabled, bool is_blocking_sync, bool ipc_enabled, int device_id) { GILReleaseGuard gil; - CUevent event; - if (CUDA_SUCCESS != (err = p_cuEventCreate(&event, flags))) { + CUevent event = nullptr; + err = invoke_in_context_or_undo( + h_ctx, + [&]() noexcept { return p_cuEventCreate(&event, flags); }, + [&]() noexcept { pw_cuEventDestroy(event); }, + /*undo_requires_target_context=*/false); + if (err != CUDA_SUCCESS) { return {}; } auto box = std::shared_ptr( new EventBox{event, timing_enabled, is_blocking_sync, ipc_enabled, device_id, h_ctx}, - [h_ctx](const EventBox* b) { + [](const EventBox* b) { event_registry.unregister_handle(b->resource); GILReleaseGuard gil; - p_cuEventDestroy(b->resource); + pw_cuEventDestroy(b->resource); delete b; } ); @@ -967,7 +1107,7 @@ EventHandle create_event_handle_ipc(const CUipcEventHandle& ipc_handle, [](const EventBox* b) { event_registry.unregister_handle(b->resource); GILReleaseGuard gil; - p_cuEventDestroy(b->resource); + pw_cuEventDestroy(b->resource); delete b; } ); @@ -1080,12 +1220,13 @@ static DevicePtrBox* get_box(const DevicePtrHandle& h) { ); } +// Return the stream that orders a device pointer's deallocation. StreamHandle deallocation_stream(const DevicePtrHandle& h) noexcept { return get_box(h)->deallocation.h_stream; } -CUresult set_deallocation_stream( - const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept { +// Replace the stream that orders a device pointer's deallocation. +CUresult set_deallocation_stream(const DevicePtrHandle& h, const StreamHandle& h_stream) noexcept { if (!h) { return CUDA_ERROR_INVALID_VALUE; } @@ -1106,7 +1247,7 @@ DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h DeallocationStream ds; if (!make_deallocation_stream(h_stream, ds)) { - p_cuMemFreeAsync(ptr, as_cu(h_stream)); + pw_cuMemFreeAsync(ptr, as_cu(h_stream)); return {}; } @@ -1114,10 +1255,10 @@ DevicePtrHandle deviceptr_alloc_from_pool(size_t size, const MemoryPoolHandle& h new DevicePtrBox{ptr, std::move(ds)}, [h_pool](DevicePtrBox* b) { GILReleaseGuard gil; - with_deallocation_context( - b->deallocation, - "cuMemFreeAsync", - [b](const DeallocationStream& stream) { + const DeallocationStream& stream = b->deallocation; + cleanup_in_context( + deallocation_context(stream), "cuMemFreeAsync", + [&]() noexcept { return p_cuMemFreeAsync( b->resource, as_cu(stream.h_stream)); }); @@ -1136,7 +1277,7 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) DeallocationStream ds; if (!make_deallocation_stream(h_stream, ds)) { - p_cuMemFreeAsync(ptr, as_cu(h_stream)); + pw_cuMemFreeAsync(ptr, as_cu(h_stream)); return {}; } @@ -1144,10 +1285,10 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) new DevicePtrBox{ptr, std::move(ds)}, [](DevicePtrBox* b) { GILReleaseGuard gil; - with_deallocation_context( - b->deallocation, - "cuMemFreeAsync", - [b](const DeallocationStream& stream) { + const DeallocationStream& stream = b->deallocation; + cleanup_in_context( + deallocation_context(stream), "cuMemFreeAsync", + [&]() noexcept { return p_cuMemFreeAsync( b->resource, as_cu(stream.h_stream)); }); @@ -1157,22 +1298,18 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) return DevicePtrHandle(box, &box->resource); } -DevicePtrHandle deviceptr_alloc(size_t size) { - GILReleaseGuard gil; - CUdeviceptr ptr; - if (CUDA_SUCCESS != (err = p_cuMemAlloc(&ptr, size))) { - return {}; +// Allocate device memory synchronously with the provided context current. +CUresult deviceptr_alloc_raw(CUdeviceptr* ptr, size_t size, + const ContextHandle& h_context) noexcept { + if (!h_context) { + return CUDA_ERROR_INVALID_CONTEXT; } - - auto box = std::shared_ptr( - new DevicePtrBox{ptr, DeallocationStream{}}, - [](DevicePtrBox* b) { - GILReleaseGuard gil; - p_cuMemFree(b->resource); - delete b; - } - ); - return DevicePtrHandle(box, &box->resource); + GILReleaseGuard gil; + return invoke_in_context_or_undo( + h_context, + [&]() noexcept { return p_cuMemAlloc(ptr, size); }, + [&]() noexcept { pw_cuMemFree(*ptr); }, + /*undo_requires_target_context=*/false); } DevicePtrHandle deviceptr_alloc_host(size_t size) { @@ -1236,10 +1373,10 @@ DevicePtrHandle deviceptr_create_mapped_graphics( [h_resource](DevicePtrBox* b) { GILReleaseGuard gil; CUgraphicsResource resource = as_cu(h_resource); - with_deallocation_context( - b->deallocation, - "cuGraphicsUnmapResources", - [b, &resource](const DeallocationStream& stream) { + const DeallocationStream& stream = b->deallocation; + cleanup_in_context( + deallocation_context(stream), "cuGraphicsUnmapResources", + [&]() noexcept { return p_cuGraphicsUnmapResources( 1, &resource, as_cu(stream.h_stream)); }); @@ -1275,12 +1412,11 @@ DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* GILAcquireGuard gil; if (gil.acquired()) { if (mr_dealloc_cb) { - with_deallocation_context( - b->deallocation, - "MemoryResource deallocate", - [mr, size, b](const DeallocationStream& stream) { - mr_dealloc_cb( - mr, b->resource, size, stream.h_stream); + const DeallocationStream& stream = b->deallocation; + cleanup_in_context( + deallocation_context(stream), "MemoryResource.deallocate", + [&]() noexcept { + mr_dealloc_cb(mr, b->resource, size, stream.h_stream); return CUDA_SUCCESS; }); } @@ -1372,7 +1508,7 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* DeallocationStream ds; if (!make_deallocation_stream(h_stream, ds)) { - p_cuMemFreeAsync(ptr, as_cu(h_stream)); + pw_cuMemFreeAsync(ptr, as_cu(h_stream)); return {}; } @@ -1381,10 +1517,10 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* [h_pool, key](DevicePtrBox* b) { ipc_ptr_cache.unregister_handle(key); GILReleaseGuard gil; - with_deallocation_context( - b->deallocation, - "cuMemFreeAsync", - [b](const DeallocationStream& stream) { + const DeallocationStream& stream = b->deallocation; + cleanup_in_context( + deallocation_context(stream), "cuMemFreeAsync", + [&]() noexcept { return p_cuMemFreeAsync( b->resource, as_cu(stream.h_stream)); }); @@ -1404,7 +1540,7 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* DeallocationStream ds; if (!make_deallocation_stream(h_stream, ds)) { - p_cuMemFreeAsync(ptr, as_cu(h_stream)); + pw_cuMemFreeAsync(ptr, as_cu(h_stream)); return {}; } @@ -1412,10 +1548,10 @@ DevicePtrHandle deviceptr_import_ipc(const MemoryPoolHandle& h_pool, const void* new DevicePtrBox{ptr, std::move(ds)}, [h_pool](DevicePtrBox* b) { GILReleaseGuard gil; - with_deallocation_context( - b->deallocation, - "cuMemFreeAsync", - [b](const DeallocationStream& stream) { + const DeallocationStream& stream = b->deallocation; + cleanup_in_context( + deallocation_context(stream), "cuMemFreeAsync", + [&]() noexcept { return p_cuMemFreeAsync( b->resource, as_cu(stream.h_stream)); }); @@ -2671,12 +2807,18 @@ struct ArrayBox { // Non-null only for a mipmap-level view: keeps the parent mipmap (the real // owner of the level's storage) alive for as long as the level is held. MipmappedArrayHandle h_parent; + ContextHandle h_context; }; struct MipmappedArrayBox { CUmipmappedArray resource; + ContextHandle h_context; }; +// Texture and surface objects are per-context pool indices. Destroying one +// with the wrong context current can silently succeed without freeing it or +// can free an unrelated object, so destruction must enter the creating +// context. Handle-based resources resolve their own context and must not. struct TexObjectBox { // Tagged so TexObjectHandle is a distinct C++ type from DevicePtrHandle / // SurfObjectHandle (all wrap `unsigned long long`). @@ -2685,31 +2827,64 @@ struct TexObjectBox { // DevicePtrHandle). The texture's resource is a union; we only need to keep // whichever backing it was built from alive, never to dereference it. std::shared_ptr h_backing; + ContextHandle h_context; }; struct SurfObjectBox { SurfObjectValue resource; OpaqueArrayHandle h_array; // surfaces are always array-backed + ContextHandle h_context; }; + +// Recover an array's owning box from its aliased resource pointer. +const ArrayBox* get_array_box(const OpaqueArrayHandle& h) noexcept { + const CUarray* p = h.get(); + return reinterpret_cast( + reinterpret_cast(p) - offsetof(ArrayBox, resource)); +} + +// Recover a mipmapped array's owning box from its aliased resource pointer. +const MipmappedArrayBox* get_mipmapped_array_box(const MipmappedArrayHandle& h) noexcept { + const CUmipmappedArray* p = h.get(); + return reinterpret_cast( + reinterpret_cast(p) + - offsetof(MipmappedArrayBox, resource)); +} + +// Wrap an array with shared owning-destruction behavior. +static OpaqueArrayHandle wrap_array_owned(CUarray arr, ContextHandle h_context) { + auto box = std::shared_ptr( + new ArrayBox{arr, {}, std::move(h_context)}, + [](const ArrayBox* b) { + GILReleaseGuard gil; + pw_cuArrayDestroy(b->resource); + delete b; + } + ); + return OpaqueArrayHandle(box, &box->resource); +} + } // namespace -OpaqueArrayHandle create_array_handle(const CUDA_ARRAY3D_DESCRIPTOR& desc) { +OpaqueArrayHandle create_array_handle(const ContextHandle& h_context, const CUDA_ARRAY3D_DESCRIPTOR& desc) { GILReleaseGuard gil; - CUarray arr; - if (CUDA_SUCCESS != (err = p_cuArray3DCreate(&arr, &desc))) { + CUarray arr = nullptr; + err = invoke_in_context_or_undo( + h_context, + [&]() noexcept { return p_cuArray3DCreate(&arr, &desc); }, + [&]() noexcept { pw_cuArrayDestroy(arr); }, + /*undo_requires_target_context=*/false); + if (err != CUDA_SUCCESS) { return {}; } - // Allocation and adoption share the same owning lifetime; the only - // difference is who calls cuArray3DCreate. Delegate so the owning box and - // its destroy-on-last-ref deleter are defined in exactly one place. - return create_array_handle_owning(arr); + return wrap_array_owned(arr, h_context); } OpaqueArrayHandle create_array_handle_ref(CUarray arr) { if (!arr) { return {}; } - auto box = std::make_shared(ArrayBox{arr, {}}); + auto box = std::make_shared(ArrayBox{arr, {}, {}}); return OpaqueArrayHandle(box, &box->resource); } @@ -2717,64 +2892,81 @@ OpaqueArrayHandle create_array_handle_owning(CUarray arr) { if (!arr) { return {}; } - auto box = std::shared_ptr( - new ArrayBox{arr, {}}, - [](const ArrayBox* b) { - GILReleaseGuard gil; - p_cuArrayDestroy(b->resource); - delete b; - } - ); - return OpaqueArrayHandle(box, &box->resource); + return wrap_array_owned(arr, {}); +} + +// Return the context retained by an array handle. +ContextHandle get_array_context(const OpaqueArrayHandle& h) noexcept { + return h ? get_array_box(h)->h_context : ContextHandle{}; } OpaqueArrayHandle create_array_level_handle(const MipmappedArrayHandle& h_mip, unsigned int level) { GILReleaseGuard gil; CUarray arr; + ContextHandle h_context = h_mip ? get_mipmapped_array_box(h_mip)->h_context : ContextHandle{}; if (CUDA_SUCCESS != (err = p_cuMipmappedArrayGetLevel(&arr, as_cu(h_mip), level))) { return {}; } // Non-owning level view: storage belongs to the mipmap. Embed the mipmap // handle so the parent outlives this level; the deleter does not destroy. auto box = std::shared_ptr( - new ArrayBox{arr, h_mip}, + new ArrayBox{arr, h_mip, h_context}, [](const ArrayBox* b) { delete b; } ); return OpaqueArrayHandle(box, &box->resource); } -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) { GILReleaseGuard gil; - CUmipmappedArray mip; - if (CUDA_SUCCESS != (err = p_cuMipmappedArrayCreate(&mip, &desc, num_levels))) { + CUmipmappedArray mip = nullptr; + err = invoke_in_context_or_undo( + h_context, + [&]() noexcept { return p_cuMipmappedArrayCreate(&mip, &desc, num_levels); }, + [&]() noexcept { pw_cuMipmappedArrayDestroy(mip); }, + /*undo_requires_target_context=*/false); + if (err != CUDA_SUCCESS) { return {}; } auto box = std::shared_ptr( - new MipmappedArrayBox{mip}, + new MipmappedArrayBox{mip, h_context}, [](const MipmappedArrayBox* b) { GILReleaseGuard gil; - p_cuMipmappedArrayDestroy(b->resource); + pw_cuMipmappedArrayDestroy(b->resource); delete b; } ); return MipmappedArrayHandle(box, &box->resource); } +// Return the context retained by a mipmapped array handle. +ContextHandle get_mipmapped_array_context(const MipmappedArrayHandle& h) noexcept { + return h ? get_mipmapped_array_box(h)->h_context : ContextHandle{}; +} + namespace { TexObjectHandle make_tex_object_handle(const CUDA_RESOURCE_DESC& res, const CUDA_TEXTURE_DESC& tex, - std::shared_ptr h_backing) { + std::shared_ptr h_backing, + const ContextHandle& h_context) { GILReleaseGuard gil; - CUtexObject obj; - if (CUDA_SUCCESS != (err = p_cuTexObjectCreate(&obj, &res, &tex, nullptr))) { + CUtexObject obj = 0; + err = invoke_in_context_or_undo( + h_context, + [&]() noexcept { return p_cuTexObjectCreate(&obj, &res, &tex, nullptr); }, + [&]() noexcept { pw_cuTexObjectDestroy(obj); }, + /*undo_requires_target_context=*/true); + if (err != CUDA_SUCCESS) { return {}; } auto box = std::shared_ptr( - new TexObjectBox{TexObjectValue{obj}, std::move(h_backing)}, + new TexObjectBox{TexObjectValue{obj}, std::move(h_backing), h_context}, [](const TexObjectBox* b) { GILReleaseGuard gil; - p_cuTexObjectDestroy(b->resource.raw); + cleanup_in_context(b->h_context, "cuTexObjectDestroy", [&]() noexcept { + return p_cuTexObjectDestroy(b->resource.raw); + }); delete b; } ); @@ -2782,36 +2974,47 @@ TexObjectHandle make_tex_object_handle(const CUDA_RESOURCE_DESC& res, } } // namespace -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) { - return make_tex_object_handle(res, tex, h_backing); + return make_tex_object_handle(res, tex, h_backing, h_context); } -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) { - return make_tex_object_handle(res, tex, h_backing); + return make_tex_object_handle(res, tex, h_backing, h_context); } -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) { - return make_tex_object_handle(res, tex, h_backing); + return make_tex_object_handle(res, tex, h_backing, h_context); } -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) { GILReleaseGuard gil; - CUsurfObject obj; - if (CUDA_SUCCESS != (err = p_cuSurfObjectCreate(&obj, &res))) { + CUsurfObject obj = 0; + err = invoke_in_context_or_undo( + h_context, + [&]() noexcept { return p_cuSurfObjectCreate(&obj, &res); }, + [&]() noexcept { pw_cuSurfObjectDestroy(obj); }, + /*undo_requires_target_context=*/true); + if (err != CUDA_SUCCESS) { return {}; } auto box = std::shared_ptr( - new SurfObjectBox{SurfObjectValue{obj}, h_backing}, + new SurfObjectBox{SurfObjectValue{obj}, h_backing, h_context}, [](const SurfObjectBox* b) { GILReleaseGuard gil; - p_cuSurfObjectDestroy(b->resource.raw); + cleanup_in_context(b->h_context, "cuSurfObjectDestroy", [&]() noexcept { + return p_cuSurfObjectDestroy(b->resource.raw); + }); delete b; } ); diff --git a/cuda_core/cuda/core/_cpp/resource_handles.hpp b/cuda_core/cuda/core/_cpp/resource_handles.hpp index ff1a12a4618..57ac9a244d6 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.hpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.hpp @@ -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; @@ -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 // ============================================================================ @@ -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. @@ -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. @@ -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. @@ -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); // ============================================================================ diff --git a/cuda_core/cuda/core/_device.pyi b/cuda_core/cuda/core/_device.pyi index 369f2b198d8..39ff7d3a26e 100644 --- a/cuda_core/cuda/core/_device.pyi +++ b/cuda_core/cuda/core/_device.pyi @@ -579,7 +579,7 @@ 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. @@ -587,6 +587,9 @@ class Device: 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,7 +719,7 @@ class Device: """ def sync(self) -> None: - """Synchronize the device. + """Synchronize this device. Note ---- @@ -722,7 +727,7 @@ class Device: """ 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 ---- diff --git a/cuda_core/cuda/core/_device.pyx b/cuda_core/cuda/core/_device.pyx index a52287a2aed..78245c6107d 100644 --- a/cuda_core/cuda/core/_device.pyx +++ b/cuda_core/cuda/core/_device.pyx @@ -23,6 +23,7 @@ from cuda.core._resource_handles cimport ( GreenCtxHandle, create_context_handle_ref, create_green_ctx_handle, + context_synchronize, get_primary_context, get_last_error, as_cu, @@ -37,7 +38,9 @@ from cuda.core._utils.cuda_utils import ( handle_return, runtime, ) -from cuda.core._stream cimport default_stream +from cuda.core._stream cimport ( + default_stream, +) from typing import TYPE_CHECKING @@ -1021,6 +1024,7 @@ class Device: raise CUDAError( f"Device {self._device_id} is not yet initialized, perhaps you forgot to call .set_current() first?" ) + Context_check_open(self._context) @classmethod @@ -1202,8 +1206,11 @@ class Device: from cuda.core._memory import DeviceMemoryResource self._memory_resource = DeviceMemoryResource(self._device_id) else: - from cuda.core._memory._legacy import _SynchronousMemoryResource - self._memory_resource = _SynchronousMemoryResource(self._device_id) + from cuda.core._memory._device_memory_resource import ( + _SynchronousMemoryResource, + ) + self._memory_resource = _SynchronousMemoryResource( + self._device_id, self._context) return self._memory_resource @@ -1215,7 +1222,7 @@ class Device: @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. @@ -1223,6 +1230,9 @@ class Device: 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. + """ return default_stream() @@ -1261,7 +1271,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 -------- @@ -1276,6 +1288,7 @@ class Device: """ cdef ContextHandle h_context cdef cydriver.CUcontext prev_ctx, curr_ctx + cdef cydriver.CUdevice prev_dev cdef Context prev_owned = None if ctx is not None: @@ -1289,10 +1302,12 @@ class Device: ) if self._has_inited and self._context is not None: prev_owned = self._context - # prev_ctx is the previous context curr_ctx = as_cu(ctx._h_context) prev_ctx = NULL with nogil: + HANDLE_RETURN(cydriver.cuCtxGetCurrent(&prev_ctx)) + if prev_ctx != NULL: + HANDLE_RETURN(cydriver.cuCtxGetDevice(&prev_dev)) HANDLE_RETURN(cydriver.cuCtxPopCurrent(&prev_ctx)) HANDLE_RETURN(cydriver.cuCtxPushCurrent(curr_ctx)) self._has_inited = True @@ -1300,7 +1315,8 @@ class Device: if prev_ctx != NULL: 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) + return Context._from_handle( + Context, create_context_handle_ref(prev_ctx), prev_dev) else: # use primary ctx h_context = get_primary_context(self._device_id) @@ -1381,7 +1397,7 @@ class Device: return Context._from_green_ctx(Context, h_green, self._device_id) 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: @@ -1393,7 +1409,7 @@ class Device: Note ---- - Device must be initialized. + Device must be initialized. New streams are created on this device. Parameters ---------- @@ -1412,7 +1428,7 @@ class Device: return Stream._init(obj=obj, options=options, device_id=self._device_id, ctx=self._context) 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 ---- @@ -1462,7 +1478,7 @@ class Device: return self.memory_resource.allocate(size, stream=stream) def sync(self) -> None: - """Synchronize the device. + """Synchronize this device. Note ---- @@ -1470,10 +1486,14 @@ class Device: """ self._check_context_initialized() - handle_return(runtime.cudaDeviceSynchronize()) + cdef Context ctx = self._context + cdef cydriver.CUresult status + with nogil: + status = context_synchronize(ctx._h_context) + HANDLE_RETURN(status) def create_graph_builder(self) -> GraphBuilder: - """Create a new :obj:`~graph.GraphBuilder` object. + """Create a new :obj:`~graph.GraphBuilder` on this device. Returns ------- @@ -1487,12 +1507,10 @@ class Device: return GraphBuilder._init(self.create_stream()) 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 ---- @@ -1513,15 +1531,13 @@ class Device: from cuda.core.texture._array import _create_opaque_array self._check_context_initialized() - return _create_opaque_array(options) + return _create_opaque_array(options, self._context, self._device_id) 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 ---- @@ -1542,20 +1558,18 @@ class Device: from cuda.core.texture._mipmapped_array import _create_mipmapped_array self._check_context_initialized() - return _create_mipmapped_array(options) + return _create_mipmapped_array(options, self._context, self._device_id) 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 ---- @@ -1578,18 +1592,16 @@ class Device: from cuda.core.texture._texture import _create_texture_object self._check_context_initialized() - return _create_texture_object(resource, options) + return _create_texture_object( + resource, options, self._context, self._device_id) 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 ---- @@ -1611,7 +1623,8 @@ class Device: from cuda.core.texture._surface import _create_surface_object self._check_context_initialized() - return _create_surface_object(resource) + return _create_surface_object( + resource, self._context, self._device_id) cdef inline int Device_ensure_cuda_initialized() except? -1: diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyi b/cuda_core/cuda/core/_memory/_device_memory_resource.pyi index 897e1d03302..b74c344d2cd 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_device_memory_resource.pyi @@ -4,9 +4,13 @@ import uuid from dataclasses import dataclass from cuda.core._device import Device +from cuda.core._memory._buffer import Buffer, MemoryResource from cuda.core._memory._ipc import IPCAllocationHandle from cuda.core._memory._memory_pool import _MemPool from cuda.core._memory._peer_access_utils import PeerAccessibleBySetProxy +from cuda.core._stream import Stream +from cuda.core.graph import GraphBuilder +from cuda.core.typing import DevicePointerType __all__ = ['DeviceMemoryResource', 'DeviceMemoryResourceOptions'] @@ -28,6 +32,19 @@ class DeviceMemoryResourceOptions: ipc_enabled: bool = False max_size: int = 0 +class _SynchronousMemoryResource(MemoryResource): + __slots__ = ('_context', '_device_id') + + def __init__(self, device_id: int, context=None) -> None: ... + def allocate(self, size: int, *, stream: Stream | GraphBuilder | None=None) -> Buffer: ... + def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | GraphBuilder | None=None) -> None: ... + @property + def is_device_accessible(self) -> bool: ... + @property + def is_host_accessible(self) -> bool: ... + @property + def device_id(self) -> int: ... + class DeviceMemoryResource(_MemPool): """ A device memory resource managing a stream-ordered memory pool. diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx index d72b0e45ebc..426ddb6d8dc 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx @@ -4,7 +4,11 @@ from __future__ import annotations +from libc.stdint cimport uintptr_t + from cuda.bindings cimport cydriver +from cuda.core._context cimport Context +from cuda.core._memory._buffer cimport Buffer, MemoryResource from cuda.core._memory._location cimport cumemlocation_from_id from cuda.core._memory._memory_pool cimport ( _MemPool, MP_check_open, MP_init_create_pool, MP_raise_release_threshold, @@ -12,10 +16,14 @@ from cuda.core._memory._memory_pool cimport ( from cuda.core._memory cimport _ipc from cuda.core._memory._ipc cimport IPCAllocationHandle from cuda.core._resource_handles cimport ( + ContextHandle, as_cu, + deviceptr_alloc_raw, get_device_mempool, get_last_error, + get_primary_context, ) +from cuda.core._stream cimport Stream, Stream_accept from cuda.core._utils.cuda_utils cimport ( check_or_create_options, HANDLE_RETURN, @@ -34,6 +42,8 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: from cuda.core._device import Device + from cuda.core.graph import GraphBuilder + from cuda.core.typing import DevicePointerType __all__ = ['DeviceMemoryResource', 'DeviceMemoryResourceOptions'] @@ -57,6 +67,68 @@ cdef class DeviceMemoryResourceOptions: max_size : int = 0 +class _SynchronousMemoryResource(MemoryResource): + __slots__ = ("_context", "_device_id") + + def __init__(self, device_id: int, context=None) -> None: + cdef ContextHandle h_context + from .._device import Device + + self._device_id = Device(device_id).device_id + if context is None: + h_context = get_primary_context(self._device_id) + if not h_context: + HANDLE_RETURN(get_last_error()) + context = Context._from_handle( + Context, h_context, self._device_id) + self._context = context + + def allocate( + self, + size_t size, + *, + stream: Stream | GraphBuilder | None = None, + ) -> Buffer: + # cuMemAlloc is synchronous; stream is accepted (and validated) + # for interface conformance but not used. + if stream is not None: + Stream_accept(stream) + + cdef Context context = self._context + cdef cydriver.CUdeviceptr ptr = 0 + if size: + with nogil: + HANDLE_RETURN(deviceptr_alloc_raw(&ptr, size, context._h_context)) + return Buffer._init(ptr, size, self) + + def deallocate( + self, + ptr: DevicePointerType, + size_t size, + *, + stream: Stream | GraphBuilder | None = None, + ) -> None: + if stream is not None: + Stream_accept(stream).sync() + cdef cydriver.CUdeviceptr devptr + if size: + devptr = int(ptr) + with nogil: + HANDLE_RETURN(cydriver.cuMemFree(devptr)) + + @property + def is_device_accessible(self) -> bool: + return True + + @property + def is_host_accessible(self) -> bool: + return False + + @property + def device_id(self) -> int: + return self._device_id + + cdef class DeviceMemoryResource(_MemPool): """ A device memory resource managing a stream-ordered memory pool. diff --git a/cuda_core/cuda/core/_memory/_legacy.py b/cuda_core/cuda/core/_memory/_legacy.py index 4acbcb54e3a..a2a8843a448 100644 --- a/cuda_core/cuda/core/_memory/_legacy.py +++ b/cuda_core/cuda/core/_memory/_legacy.py @@ -96,47 +96,3 @@ def is_host_accessible(self) -> bool: def device_id(self) -> int: """This memory resource is not bound to any GPU.""" raise RuntimeError("a pinned memory resource is not bound to any GPU") - - -class _SynchronousMemoryResource(MemoryResource): - __slots__ = ("_device_id",) - - def __init__(self, device_id: int) -> None: - from .._device import Device - - self._device_id = Device(device_id).device_id - - def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> Buffer: - # cuMemAlloc is synchronous; stream is accepted (and validated) - # for interface conformance but not used. - from cuda.core._stream import Stream_accept - - if stream is not None: - Stream_accept(stream) - if size: - err, ptr = driver.cuMemAlloc(size) - raise_if_driver_error(err) - else: - ptr = 0 - return Buffer._init(ptr, size, self) - - def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | GraphBuilder | None = None) -> None: - from cuda.core._stream import Stream_accept - - if stream is not None: - Stream_accept(stream).sync() - if size: - (err,) = driver.cuMemFree(ptr) - raise_if_driver_error(err) - - @property - def is_device_accessible(self) -> bool: - return True - - @property - def is_host_accessible(self) -> bool: - return False - - @property - def device_id(self) -> int: - return self._device_id diff --git a/cuda_core/cuda/core/_resource_handles.pxd b/cuda_core/cuda/core/_resource_handles.pxd index 568af27ac2e..acf10e0fa2c 100644 --- a/cuda_core/cuda/core/_resource_handles.pxd +++ b/cuda_core/cuda/core/_resource_handles.pxd @@ -178,6 +178,12 @@ cdef GreenCtxHandle create_green_ctx_handle( cdef GreenCtxHandle create_green_ctx_handle_ref(cydriver.CUgreenCtx ctx) except+ nogil cdef ContextHandle get_primary_context(int device_id) except+ nogil cdef ContextHandle get_current_context() except+ nogil +cdef cydriver.CUresult context_synchronize( + const ContextHandle& h_context) noexcept nogil +cdef cydriver.CUresult context_get_stream_priority_range( + const ContextHandle& h_context, + int* least_priority, + int* greatest_priority) noexcept nogil # Stream handles cdef StreamHandle create_stream_handle( @@ -219,7 +225,8 @@ cdef MemoryPoolHandle create_mempool_handle_ipc( cdef DevicePtrHandle deviceptr_alloc_from_pool( size_t size, const MemoryPoolHandle& h_pool, const StreamHandle& h_stream) except+ nogil cdef DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) except+ nogil -cdef DevicePtrHandle deviceptr_alloc(size_t size) except+ nogil +cdef cydriver.CUresult deviceptr_alloc_raw( + cydriver.CUdeviceptr* ptr, size_t size, const ContextHandle& h_context) noexcept nogil cdef DevicePtrHandle deviceptr_alloc_host(size_t size) except+ nogil cdef DevicePtrHandle deviceptr_create_ref(cydriver.CUdeviceptr ptr) except+ nogil cdef DevicePtrHandle deviceptr_create_with_owner(cydriver.CUdeviceptr ptr, object owner) except+ nogil @@ -325,23 +332,29 @@ cdef FileDescriptorHandle create_fd_handle(int fd) except+ nogil cdef FileDescriptorHandle create_fd_handle_ref(int fd) except+ nogil # Array / mipmapped-array / texture / surface handles (PR #467) -cdef OpaqueArrayHandle create_array_handle(const cydriver.CUDA_ARRAY3D_DESCRIPTOR& desc) except+ nogil +cdef OpaqueArrayHandle create_array_handle( + const ContextHandle& h_context, const cydriver.CUDA_ARRAY3D_DESCRIPTOR& desc) except+ nogil cdef OpaqueArrayHandle create_array_handle_ref(cydriver.CUarray arr) except+ nogil cdef OpaqueArrayHandle create_array_handle_owning(cydriver.CUarray arr) except+ nogil +cdef ContextHandle get_array_context(const OpaqueArrayHandle& h) noexcept nogil cdef OpaqueArrayHandle create_array_level_handle(const MipmappedArrayHandle& h_mip, unsigned int level) except+ nogil cdef MipmappedArrayHandle create_mipmapped_array_handle( - const cydriver.CUDA_ARRAY3D_DESCRIPTOR& desc, unsigned int num_levels) except+ nogil + const ContextHandle& h_context, const cydriver.CUDA_ARRAY3D_DESCRIPTOR& desc, + unsigned int num_levels) except+ nogil +cdef ContextHandle get_mipmapped_array_context( + const MipmappedArrayHandle& h) noexcept nogil cdef TexObjectHandle create_tex_object_handle_array( - const cydriver.CUDA_RESOURCE_DESC& res, const cydriver.CUDA_TEXTURE_DESC& tex, - const OpaqueArrayHandle& h_backing) except+ nogil + const ContextHandle& h_context, const cydriver.CUDA_RESOURCE_DESC& res, + const cydriver.CUDA_TEXTURE_DESC& tex, const OpaqueArrayHandle& h_backing) except+ nogil cdef TexObjectHandle create_tex_object_handle_mipmap( - const cydriver.CUDA_RESOURCE_DESC& res, const cydriver.CUDA_TEXTURE_DESC& tex, - const MipmappedArrayHandle& h_backing) except+ nogil + const ContextHandle& h_context, const cydriver.CUDA_RESOURCE_DESC& res, + const cydriver.CUDA_TEXTURE_DESC& tex, const MipmappedArrayHandle& h_backing) except+ nogil cdef TexObjectHandle create_tex_object_handle_linear( - const cydriver.CUDA_RESOURCE_DESC& res, const cydriver.CUDA_TEXTURE_DESC& tex, - const DevicePtrHandle& h_backing) except+ nogil + const ContextHandle& h_context, const cydriver.CUDA_RESOURCE_DESC& res, + const cydriver.CUDA_TEXTURE_DESC& tex, const DevicePtrHandle& h_backing) except+ nogil cdef SurfObjectHandle create_surf_object_handle( - const cydriver.CUDA_RESOURCE_DESC& res, const OpaqueArrayHandle& h_backing) except+ nogil + const ContextHandle& h_context, const cydriver.CUDA_RESOURCE_DESC& res, + const OpaqueArrayHandle& h_backing) except+ nogil # SM resource split (13.1+ — calls through function pointer, safe on older bindings) # groupParams is void* here to avoid referencing CU_DEV_SM_RESOURCE_GROUP_PARAMS diff --git a/cuda_core/cuda/core/_resource_handles.pyx b/cuda_core/cuda/core/_resource_handles.pyx index c7de24666f8..beecb4b745a 100644 --- a/cuda_core/cuda/core/_resource_handles.pyx +++ b/cuda_core/cuda/core/_resource_handles.pyx @@ -51,6 +51,12 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": ContextHandle get_primary_context "cuda_core::get_primary_context" ( int device_id) except+ nogil ContextHandle get_current_context "cuda_core::get_current_context" () except+ nogil + cydriver.CUresult context_synchronize "cuda_core::context_synchronize" ( + const ContextHandle& h_context) noexcept nogil + cydriver.CUresult context_get_stream_priority_range "cuda_core::context_get_stream_priority_range" ( + const ContextHandle& h_context, + int* least_priority, + int* greatest_priority) noexcept nogil # Stream handles StreamHandle create_stream_handle "cuda_core::create_stream_handle" ( @@ -107,7 +113,8 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": size_t size, const MemoryPoolHandle& h_pool, const StreamHandle& h_stream) except+ nogil DevicePtrHandle deviceptr_alloc_async "cuda_core::deviceptr_alloc_async" ( size_t size, const StreamHandle& h_stream) except+ nogil - DevicePtrHandle deviceptr_alloc "cuda_core::deviceptr_alloc" (size_t size) except+ nogil + cydriver.CUresult deviceptr_alloc_raw "cuda_core::deviceptr_alloc_raw" ( + cydriver.CUdeviceptr* ptr, size_t size, const ContextHandle& h_context) noexcept nogil DevicePtrHandle deviceptr_alloc_host "cuda_core::deviceptr_alloc_host" (size_t size) except+ nogil DevicePtrHandle deviceptr_create_ref "cuda_core::deviceptr_create_ref" ( cydriver.CUdeviceptr ptr) except+ nogil @@ -253,26 +260,32 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": # Array / mipmapped-array / texture / surface handles (PR #467) OpaqueArrayHandle create_array_handle "cuda_core::create_array_handle" ( - const cydriver.CUDA_ARRAY3D_DESCRIPTOR& desc) except+ nogil + const ContextHandle& h_context, const cydriver.CUDA_ARRAY3D_DESCRIPTOR& desc) except+ nogil OpaqueArrayHandle create_array_handle_ref "cuda_core::create_array_handle_ref" ( cydriver.CUarray arr) except+ nogil OpaqueArrayHandle create_array_handle_owning "cuda_core::create_array_handle_owning" ( cydriver.CUarray arr) except+ nogil + ContextHandle get_array_context "cuda_core::get_array_context" ( + const OpaqueArrayHandle& h) noexcept nogil OpaqueArrayHandle create_array_level_handle "cuda_core::create_array_level_handle" ( const MipmappedArrayHandle& h_mip, unsigned int level) except+ nogil MipmappedArrayHandle create_mipmapped_array_handle "cuda_core::create_mipmapped_array_handle" ( - const cydriver.CUDA_ARRAY3D_DESCRIPTOR& desc, unsigned int num_levels) except+ nogil + const ContextHandle& h_context, const cydriver.CUDA_ARRAY3D_DESCRIPTOR& desc, + unsigned int num_levels) except+ nogil + ContextHandle get_mipmapped_array_context "cuda_core::get_mipmapped_array_context" ( + const MipmappedArrayHandle& h) noexcept nogil TexObjectHandle create_tex_object_handle_array "cuda_core::create_tex_object_handle_array" ( - const cydriver.CUDA_RESOURCE_DESC& res, const cydriver.CUDA_TEXTURE_DESC& tex, - const OpaqueArrayHandle& h_backing) except+ nogil + const ContextHandle& h_context, const cydriver.CUDA_RESOURCE_DESC& res, + const cydriver.CUDA_TEXTURE_DESC& tex, const OpaqueArrayHandle& h_backing) except+ nogil TexObjectHandle create_tex_object_handle_mipmap "cuda_core::create_tex_object_handle_mipmap" ( - const cydriver.CUDA_RESOURCE_DESC& res, const cydriver.CUDA_TEXTURE_DESC& tex, - const MipmappedArrayHandle& h_backing) except+ nogil + const ContextHandle& h_context, const cydriver.CUDA_RESOURCE_DESC& res, + const cydriver.CUDA_TEXTURE_DESC& tex, const MipmappedArrayHandle& h_backing) except+ nogil TexObjectHandle create_tex_object_handle_linear "cuda_core::create_tex_object_handle_linear" ( - const cydriver.CUDA_RESOURCE_DESC& res, const cydriver.CUDA_TEXTURE_DESC& tex, - const DevicePtrHandle& h_backing) except+ nogil + const ContextHandle& h_context, const cydriver.CUDA_RESOURCE_DESC& res, + const cydriver.CUDA_TEXTURE_DESC& tex, const DevicePtrHandle& h_backing) except+ nogil SurfObjectHandle create_surf_object_handle "cuda_core::create_surf_object_handle" ( - const cydriver.CUDA_RESOURCE_DESC& res, const OpaqueArrayHandle& h_backing) except+ nogil + const ContextHandle& h_context, const cydriver.CUDA_RESOURCE_DESC& res, + const OpaqueArrayHandle& h_backing) except+ nogil # ============================================================================= @@ -297,11 +310,17 @@ cdef const char* _CUDA_DRIVER_API_V1_NAME = b"cuda.core._resource_handles._CUDA_ # Declare extern variables with reinterpret_cast to allow void* assignment cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": + # Error formatting + void* p_cuGetErrorName "reinterpret_cast(cuda_core::p_cuGetErrorName)" + void* p_cuGetErrorString "reinterpret_cast(cuda_core::p_cuGetErrorString)" + # Context void* p_cuDevicePrimaryCtxRetain "reinterpret_cast(cuda_core::p_cuDevicePrimaryCtxRetain)" void* p_cuDevicePrimaryCtxRelease "reinterpret_cast(cuda_core::p_cuDevicePrimaryCtxRelease)" void* p_cuCtxGetCurrent "reinterpret_cast(cuda_core::p_cuCtxGetCurrent)" void* p_cuCtxSetCurrent "reinterpret_cast(cuda_core::p_cuCtxSetCurrent)" + void* p_cuCtxSynchronize "reinterpret_cast(cuda_core::p_cuCtxSynchronize)" + void* p_cuCtxGetStreamPriorityRange "reinterpret_cast(cuda_core::p_cuCtxGetStreamPriorityRange)" void* p_cuGreenCtxCreate "reinterpret_cast(cuda_core::p_cuGreenCtxCreate)" void* p_cuGreenCtxDestroy "reinterpret_cast(cuda_core::p_cuGreenCtxDestroy)" void* p_cuCtxFromGreenCtx "reinterpret_cast(cuda_core::p_cuCtxFromGreenCtx)" @@ -408,8 +427,9 @@ cdef void* _get_optional_driver_fn(str name): cdef void _init_driver_fn_pointers() noexcept: + global p_cuGetErrorName, p_cuGetErrorString global p_cuDevicePrimaryCtxRetain, p_cuDevicePrimaryCtxRelease, p_cuCtxGetCurrent - global p_cuCtxSetCurrent + global p_cuCtxSetCurrent, p_cuCtxSynchronize, p_cuCtxGetStreamPriorityRange global p_cuGreenCtxCreate, p_cuGreenCtxDestroy, p_cuCtxFromGreenCtx global p_cuDevResourceGenerateDesc, p_cuGreenCtxStreamCreate global p_cuStreamCreateWithPriority, p_cuStreamDestroy @@ -435,11 +455,17 @@ cdef void _init_driver_fn_pointers() noexcept: global p_cuTexObjectCreate, p_cuTexObjectDestroy global p_cuSurfObjectCreate, p_cuSurfObjectDestroy + # Error formatting + p_cuGetErrorName = _get_driver_fn("cuGetErrorName") + p_cuGetErrorString = _get_driver_fn("cuGetErrorString") + # Context p_cuDevicePrimaryCtxRetain = _get_driver_fn("cuDevicePrimaryCtxRetain") p_cuDevicePrimaryCtxRelease = _get_driver_fn("cuDevicePrimaryCtxRelease") p_cuCtxGetCurrent = _get_driver_fn("cuCtxGetCurrent") p_cuCtxSetCurrent = _get_driver_fn("cuCtxSetCurrent") + p_cuCtxSynchronize = _get_driver_fn("cuCtxSynchronize") + p_cuCtxGetStreamPriorityRange = _get_driver_fn("cuCtxGetStreamPriorityRange") p_cuGreenCtxCreate = _get_optional_driver_fn("cuGreenCtxCreate") p_cuGreenCtxDestroy = _get_optional_driver_fn("cuGreenCtxDestroy") p_cuCtxFromGreenCtx = _get_optional_driver_fn("cuCtxFromGreenCtx") diff --git a/cuda_core/cuda/core/_stream.pyx b/cuda_core/cuda/core/_stream.pyx index 76aee36cdd3..9c51c488a6c 100644 --- a/cuda_core/cuda/core/_stream.pyx +++ b/cuda_core/cuda/core/_stream.pyx @@ -20,7 +20,10 @@ import warnings from dataclasses import dataclass from typing import Protocol, TYPE_CHECKING -from cuda.core._context cimport Context +from cuda.core._context cimport ( + Context, + Context_check_open, +) from cuda.core._device_resources cimport DeviceResources from cuda.core._event import Event, EventOptions @@ -32,6 +35,7 @@ from cuda.core._resource_handles cimport ( create_event_handle_noctx, create_stream_handle, create_stream_handle_with_owner, + context_get_stream_priority_range, get_current_context, get_last_error, get_legacy_stream, @@ -129,10 +133,7 @@ cdef class Stream: cdef StreamHandle h_stream cdef cydriver.CUstream borrowed cdef ContextHandle h_context - - # Extract context handle if provided - if ctx is not None: - h_context = (ctx)._h_context + cdef Context context if obj is not None and options is not None: raise ValueError("obj and options cannot be both specified") @@ -144,6 +145,12 @@ cdef class Stream: h_stream = create_stream_handle_with_owner(borrowed, obj) return Stream._from_handle(cls, h_stream) + if ctx is None: + raise RuntimeError("A CUDA context is required to create a stream") + context = ctx + Context_check_open(context) + h_context = context._h_context + cdef StreamOptions opts = check_or_create_options(StreamOptions, options, "Stream options") nonblocking = opts.nonblocking priority = opts.priority @@ -154,13 +161,9 @@ cdef class Stream: cdef int high, low cdef cydriver.CUresult res_code with nogil: - res_code = cydriver.cuCtxGetStreamPriorityRange(&high, &low) - if res_code != cydriver.CUresult.CUDA_SUCCESS: - if res_code == cydriver.CUresult.CUDA_ERROR_INVALID_CONTEXT: - raise RuntimeError( - "No current CUDA context. Call dev.set_current() before creating streams." - ) - HANDLE_RETURN(res_code) + res_code = context_get_stream_priority_range( + context._h_context, &high, &low) + HANDLE_RETURN(res_code) cdef int prio if priority is not None: prio = priority diff --git a/cuda_core/cuda/core/texture/_array.pyi b/cuda_core/cuda/core/texture/_array.pyi index 1b48b8afc82..f4258d5cbd5 100644 --- a/cuda_core/cuda/core/texture/_array.pyi +++ b/cuda_core/cuda/core/texture/_array.pyi @@ -3,6 +3,7 @@ from dataclasses import dataclass from cuda.bindings import cydriver +from cuda.core._context import Context from cuda.core.typing import ArrayFormatType _ARRAYFORMAT_TO_CU = {ArrayFormatType.UINT8: int(cydriver.CU_AD_FORMAT_UNSIGNED_INT8), ArrayFormatType.UINT16: int(cydriver.CU_AD_FORMAT_UNSIGNED_INT16), ArrayFormatType.UINT32: int(cydriver.CU_AD_FORMAT_UNSIGNED_INT32), ArrayFormatType.INT8: int(cydriver.CU_AD_FORMAT_SIGNED_INT8), ArrayFormatType.INT16: int(cydriver.CU_AD_FORMAT_SIGNED_INT16), ArrayFormatType.INT32: int(cydriver.CU_AD_FORMAT_SIGNED_INT32), ArrayFormatType.FLOAT16: int(cydriver.CU_AD_FORMAT_HALF), ArrayFormatType.FLOAT32: int(cydriver.CU_AD_FORMAT_FLOAT)} @@ -161,8 +162,8 @@ def _validate_format_channels(format, num_channels): def _validate_array_shape(shape): """Coerce ``shape`` to a tuple of ints and validate rank (1-3) and that every extent is >= 1. Returns the normalized tuple.""" -def _create_opaque_array(options): - """Allocate a new :class:`OpaqueArray` on the current device. +def _create_opaque_array(options, ctx: Context, device_id: int): + """Allocate a new :class:`OpaqueArray` on the specified device. Backs :meth:`cuda.core.Device.create_opaque_array`. ``options`` is an :class:`OpaqueArrayOptions` (or a mapping accepted by it); it is validated diff --git a/cuda_core/cuda/core/texture/_array.pyx b/cuda_core/cuda/core/texture/_array.pyx index fbfa908c526..e5fc3f6c9e2 100644 --- a/cuda_core/cuda/core/texture/_array.pyx +++ b/cuda_core/cuda/core/texture/_array.pyx @@ -9,6 +9,7 @@ from libc.stdint cimport intptr_t from libc.string cimport memset from cuda.bindings cimport cydriver +from cuda.core._context cimport Context from cuda.core._memory._buffer cimport Buffer, Buffer_check_open from cuda.core._resource_handles cimport ( OpaqueArrayHandle, @@ -515,8 +516,8 @@ cdef OpaqueArray _array_from_handle(OpaqueArrayHandle h, int device_id): return self -def _create_opaque_array(options): - """Allocate a new :class:`OpaqueArray` on the current device. +def _create_opaque_array(options, Context ctx, int device_id): + """Allocate a new :class:`OpaqueArray` on the specified device. Backs :meth:`cuda.core.Device.create_opaque_array`. ``options`` is an :class:`OpaqueArrayOptions` (or a mapping accepted by it); it is validated @@ -545,7 +546,7 @@ def _create_opaque_array(options): Flags=flags, ) - cdef OpaqueArrayHandle h = create_array_handle(desc3d) + cdef OpaqueArrayHandle h = create_array_handle(ctx._h_context, desc3d) if not h: HANDLE_RETURN(get_last_error()) @@ -555,5 +556,5 @@ def _create_opaque_array(options): self._format = c_format self._num_channels = opts.num_channels self._surface_load_store = bool(opts.is_surface_load_store) - self._device_id = _get_current_device_id() + self._device_id = device_id return self diff --git a/cuda_core/cuda/core/texture/_mipmapped_array.pyi b/cuda_core/cuda/core/texture/_mipmapped_array.pyi index e4ec707b458..3d843abe70a 100644 --- a/cuda_core/cuda/core/texture/_mipmapped_array.pyi +++ b/cuda_core/cuda/core/texture/_mipmapped_array.pyi @@ -2,6 +2,8 @@ from dataclasses import dataclass +from cuda.core._context import Context + @dataclass class MipmappedArrayOptions: @@ -105,8 +107,8 @@ class MipmappedArray: def __exit__(self, exc_type, exc, tb): ... def __repr__(self): ... -def _create_mipmapped_array(options): - """Allocate a new :class:`MipmappedArray` on the current device. +def _create_mipmapped_array(options, ctx: Context, device_id: int): + """Allocate a new :class:`MipmappedArray` on the specified device. Backs :meth:`cuda.core.Device.create_mipmapped_array`. ``options`` is a :class:`MipmappedArrayOptions` (or a mapping accepted by it); its fields are diff --git a/cuda_core/cuda/core/texture/_mipmapped_array.pyx b/cuda_core/cuda/core/texture/_mipmapped_array.pyx index abc30c6b25c..8d6bf5a2589 100644 --- a/cuda_core/cuda/core/texture/_mipmapped_array.pyx +++ b/cuda_core/cuda/core/texture/_mipmapped_array.pyx @@ -5,6 +5,7 @@ from __future__ import annotations from cuda.bindings cimport cydriver +from cuda.core._context cimport Context from cuda.core.texture._array cimport _array_from_handle from cuda.core.texture._array import ( _ARRAYFORMAT_TO_CU, @@ -20,10 +21,7 @@ from cuda.core._resource_handles cimport ( create_mipmapped_array_handle, get_last_error, ) -from cuda.core._utils.cuda_utils cimport ( - HANDLE_RETURN, - _get_current_device_id, -) +from cuda.core._utils.cuda_utils cimport HANDLE_RETURN from dataclasses import dataclass @@ -191,8 +189,8 @@ cdef class MipmappedArray: f"num_levels={self._num_levels})" ) -def _create_mipmapped_array(options): - """Allocate a new :class:`MipmappedArray` on the current device. +def _create_mipmapped_array(options, Context ctx, int device_id): + """Allocate a new :class:`MipmappedArray` on the specified device. Backs :meth:`cuda.core.Device.create_mipmapped_array`. ``options`` is a :class:`MipmappedArrayOptions` (or a mapping accepted by it); its fields are @@ -221,7 +219,8 @@ def _create_mipmapped_array(options): Flags=flags, ) - cdef MipmappedArrayHandle h = create_mipmapped_array_handle(desc3d, c_levels) + cdef MipmappedArrayHandle h = create_mipmapped_array_handle( + ctx._h_context, desc3d, c_levels) if not h: HANDLE_RETURN(get_last_error()) @@ -232,5 +231,5 @@ def _create_mipmapped_array(options): self._num_channels = opts.num_channels self._num_levels = opts.num_levels self._surface_load_store = bool(opts.is_surface_load_store) - self._device_id = _get_current_device_id() + self._device_id = device_id return self diff --git a/cuda_core/cuda/core/texture/_surface.pyi b/cuda_core/cuda/core/texture/_surface.pyi index b153ff31fed..c363b63b05e 100644 --- a/cuda_core/cuda/core/texture/_surface.pyi +++ b/cuda_core/cuda/core/texture/_surface.pyi @@ -1,5 +1,8 @@ # This file was generated by stubgen-pyx v0.2.19 from cuda_core/cuda/core/texture/_surface.pyx +from cuda.core._context import Context + + class SurfaceObject: """A bindless surface handle for kernel-side typed load/store. @@ -39,8 +42,8 @@ class SurfaceObject: def __exit__(self, exc_type, exc, tb): ... def __repr__(self): ... -def _create_surface_object(resource): - """Create a :class:`SurfaceObject` on the current device. +def _create_surface_object(resource, ctx: Context, device_id: int): + """Create a :class:`SurfaceObject` on the specified device. Backs :meth:`cuda.core.Device.create_surface_object`. ``resource`` must be a :class:`ResourceDescriptor` wrapping an :class:`OpaqueArray` allocated with diff --git a/cuda_core/cuda/core/texture/_surface.pyx b/cuda_core/cuda/core/texture/_surface.pyx index 074f438ad47..790ce048ecd 100644 --- a/cuda_core/cuda/core/texture/_surface.pyx +++ b/cuda_core/cuda/core/texture/_surface.pyx @@ -7,19 +7,19 @@ from __future__ import annotations from libc.string cimport memset from cuda.bindings cimport cydriver +from cuda.core._context cimport Context from cuda.core.texture._array cimport OpaqueArray, OpaqueArray_check_open from cuda.core._resource_handles cimport ( + ContextHandle, SurfObjectHandle, as_cu, as_intptr, create_surf_object_handle, + get_array_context, get_last_error, ) from cuda.core.texture._texture import ResourceDescriptor -from cuda.core._utils.cuda_utils cimport ( - HANDLE_RETURN, - _get_current_device_id, -) +from cuda.core._utils.cuda_utils cimport HANDLE_RETURN cdef class SurfaceObject: @@ -85,8 +85,8 @@ cdef class SurfaceObject: return f"SurfaceObject(handle=0x{as_intptr(self._handle):x})" -def _create_surface_object(resource): - """Create a :class:`SurfaceObject` on the current device. +def _create_surface_object(resource, Context ctx, int device_id): + """Create a :class:`SurfaceObject` on the specified device. Backs :meth:`cuda.core.Device.create_surface_object`. ``resource`` must be a :class:`ResourceDescriptor` wrapping an :class:`OpaqueArray` allocated with @@ -106,6 +106,14 @@ def _create_surface_object(resource): cdef OpaqueArray arr = resource.source OpaqueArray_check_open(arr) + if arr._device_id != device_id: + raise ValueError( + f"resource belongs to device {arr._device_id}, " + f"but surface creation was requested on device {device_id}" + ) + cdef ContextHandle resource_context = get_array_context(arr._handle) + if resource_context and as_cu(resource_context) != as_cu(ctx._h_context): + raise ValueError("resource is not compatible with this Device object") if not arr.is_surface_load_store: raise ValueError( "OpaqueArray must be created with is_surface_load_store=True to be " @@ -117,12 +125,13 @@ def _create_surface_object(resource): res_desc.resType = cydriver.CU_RESOURCE_TYPE_ARRAY res_desc.res.array.hArray = as_cu(arr._handle) - cdef SurfObjectHandle h = create_surf_object_handle(res_desc, arr._handle) + cdef SurfObjectHandle h = create_surf_object_handle( + ctx._h_context, res_desc, arr._handle) if not h: HANDLE_RETURN(get_last_error()) cdef SurfaceObject self = SurfaceObject.__new__(SurfaceObject) self._handle = h self._source_ref = resource - self._device_id = _get_current_device_id() + self._device_id = device_id return self diff --git a/cuda_core/cuda/core/texture/_texture.pyi b/cuda_core/cuda/core/texture/_texture.pyi index 16508003091..f210be8ac89 100644 --- a/cuda_core/cuda/core/texture/_texture.pyi +++ b/cuda_core/cuda/core/texture/_texture.pyi @@ -3,6 +3,7 @@ from dataclasses import dataclass from cuda.bindings import cydriver +from cuda.core._context import Context from cuda.core.typing import AddressModeType, FilterModeType, ReadModeType _TRSF_READ_AS_INTEGER = 1 @@ -215,8 +216,8 @@ def _normalize_enum(name, value, enum_type): def _normalize_address_modes(address_mode): """Return a 3-tuple of :class:`AddressModeType` values from a scalar or 1-3 tuple. Individual entries may be plain strings.""" -def _create_texture_object(resource, options): - """Create a :class:`TextureObject` on the current device. +def _create_texture_object(resource, options, ctx: Context, device_id: int): + """Create a :class:`TextureObject` on the specified device. Backs :meth:`cuda.core.Device.create_texture_object`. ``resource`` is a :class:`ResourceDescriptor`; ``options`` is a :class:`TextureObjectOptions` diff --git a/cuda_core/cuda/core/texture/_texture.pyx b/cuda_core/cuda/core/texture/_texture.pyx index ef63c01972a..f1c29490b8a 100644 --- a/cuda_core/cuda/core/texture/_texture.pyx +++ b/cuda_core/cuda/core/texture/_texture.pyx @@ -8,6 +8,7 @@ from libc.stdint cimport intptr_t from libc.string cimport memset from cuda.bindings cimport cydriver +from cuda.core._context cimport Context from cuda.core.texture._array cimport OpaqueArray, OpaqueArray_check_open from cuda.core.texture._array import ( _ARRAYFORMAT_TO_CU, @@ -19,18 +20,18 @@ from cuda.core._memory._buffer cimport Buffer, Buffer_check_open from cuda.core.texture._mipmapped_array cimport MipmappedArray, MipmappedArray_check_open from cuda.core.texture._mipmapped_array import MipmappedArray as _PyMipmappedArray from cuda.core._resource_handles cimport ( + ContextHandle, TexObjectHandle, as_cu, as_intptr, create_tex_object_handle_array, create_tex_object_handle_linear, create_tex_object_handle_mipmap, + get_array_context, get_last_error, + get_mipmapped_array_context, ) -from cuda.core._utils.cuda_utils cimport ( - HANDLE_RETURN, - _get_current_device_id, -) +from cuda.core._utils.cuda_utils cimport HANDLE_RETURN from cuda.core.typing import AddressModeType, FilterModeType, ReadModeType @@ -474,8 +475,9 @@ cdef class TextureObject: return f"TextureObject(handle=0x{as_intptr(self._handle):x})" -def _create_texture_object(resource, options): - """Create a :class:`TextureObject` on the current device. +def _create_texture_object( + resource, options, Context ctx, int device_id): + """Create a :class:`TextureObject` on the specified device. Backs :meth:`cuda.core.Device.create_texture_object`. ``resource`` is a :class:`ResourceDescriptor`; ``options`` is a :class:`TextureObjectOptions` @@ -500,19 +502,26 @@ def _create_texture_object(resource, options): cdef MipmappedArray mip cdef Buffer buf cdef intptr_t devptr + cdef ContextHandle resource_context + cdef int resource_device_id if resource.kind == "array": arr = resource.source OpaqueArray_check_open(arr) + resource_context = get_array_context(arr._handle) + resource_device_id = arr._device_id res_desc.resType = cydriver.CU_RESOURCE_TYPE_ARRAY res_desc.res.array.hArray = as_cu(arr._handle) elif resource.kind == "mipmapped_array": mip = resource.source MipmappedArray_check_open(mip) + resource_context = get_mipmapped_array_context(mip._handle) + resource_device_id = mip._device_id res_desc.resType = cydriver.CU_RESOURCE_TYPE_MIPMAPPED_ARRAY res_desc.res.mipmap.hMipmappedArray = as_cu(mip._handle) elif resource.kind == "linear": buf = resource.source Buffer_check_open(buf) + resource_device_id = buf.device_id devptr = int(buf.handle) res_desc.resType = cydriver.CU_RESOURCE_TYPE_LINEAR res_desc.res.linear.devPtr = devptr @@ -522,6 +531,7 @@ def _create_texture_object(resource, options): elif resource.kind == "pitch2d": buf = resource.source Buffer_check_open(buf) + resource_device_id = buf.device_id devptr = int(buf.handle) res_desc.resType = cydriver.CU_RESOURCE_TYPE_PITCH2D res_desc.res.pitch2D.devPtr = devptr @@ -534,6 +544,13 @@ def _create_texture_object(resource, options): raise NotImplementedError( f"ResourceDescriptor kind {resource.kind!r} is not yet supported" ) + if resource_device_id >= 0 and resource_device_id != device_id: + raise ValueError( + f"resource belongs to device {resource_device_id}, " + f"but texture creation was requested on device {device_id}" + ) + if resource_context and as_cu(resource_context) != as_cu(ctx._h_context): + raise ValueError("resource is not compatible with this Device object") # --- Texture descriptor --- # filter_mode/read_mode/mipmap_filter_mode are normalized to their @@ -585,11 +602,14 @@ def _create_texture_object(resource, options): cdef TexObjectHandle h if resource.kind == "array": - h = create_tex_object_handle_array(res_desc, tex_desc, arr._handle) + h = create_tex_object_handle_array( + ctx._h_context, res_desc, tex_desc, arr._handle) elif resource.kind == "mipmapped_array": - h = create_tex_object_handle_mipmap(res_desc, tex_desc, mip._handle) + h = create_tex_object_handle_mipmap( + ctx._h_context, res_desc, tex_desc, mip._handle) else: # linear or pitch2d — both backed by a device Buffer - h = create_tex_object_handle_linear(res_desc, tex_desc, buf._h_ptr) + h = create_tex_object_handle_linear( + ctx._h_context, res_desc, tex_desc, buf._h_ptr) if not h: HANDLE_RETURN(get_last_error()) @@ -597,5 +617,5 @@ def _create_texture_object(resource, options): self._handle = h self._source_ref = resource self._options = opts - self._device_id = _get_current_device_id() + self._device_id = device_id return self diff --git a/cuda_core/docs/source/interoperability.rst b/cuda_core/docs/source/interoperability.rst index 87347eb9d25..33d11d540c7 100644 --- a/cuda_core/docs/source/interoperability.rst +++ b/cuda_core/docs/source/interoperability.rst @@ -26,6 +26,10 @@ Conversely, if any GPU library already sets a device (or context) to current, th method ensures that the same device/context is picked up by and shared with ``cuda.core``. +Other :class:`Device` methods do not change the current context. For example, +``dev1.sync()`` synchronizes device 1 and leaves the current context unchanged, +even when another device is current. + ``__cuda_stream__`` protocol ---------------------------- diff --git a/cuda_core/docs/source/release/1.2.0-notes.rst b/cuda_core/docs/source/release/1.2.0-notes.rst index f96a205d1e8..9591cc3a083 100644 --- a/cuda_core/docs/source/release/1.2.0-notes.rst +++ b/cuda_core/docs/source/release/1.2.0-notes.rst @@ -40,6 +40,13 @@ New features Fixes and enhancements ---------------------- +- :class:`Device` methods that create resources or synchronize now act on that + device, even when another device is current. They do not change which device + is current. For example, ``dev1.sync()`` synchronizes device 1 even when + device 0 is current. :meth:`Device.set_current` now always returns a + :class:`Context` with the correct device ID. + (`#2311 `__) + - A :class:`Buffer` is now freed correctly even when the CUDA context current at teardown is not the one it was allocated in, or when no context is current at all. This happens routinely when a buffer is released by the garbage diff --git a/cuda_core/tests/conftest.py b/cuda_core/tests/conftest.py index ff4cddcb28f..0466277bf6c 100644 --- a/cuda_core/tests/conftest.py +++ b/cuda_core/tests/conftest.py @@ -212,6 +212,15 @@ def deinit_cuda(): _ = _device_unset_current() +@pytest.fixture +def device_x2(deinit_cuda): + """Provide two CUDA devices, or skip when fewer are available.""" + devices = Device.get_all_devices() + if len(devices) < 2: + pytest.skip("Test requires at least 2 CUDA devices") + return devices[:2] + + @pytest.fixture def deinit_all_contexts_function(): def pop_all_contexts(): diff --git a/cuda_core/tests/helpers/contexts.py b/cuda_core/tests/helpers/contexts.py new file mode 100644 index 00000000000..9818875c4da --- /dev/null +++ b/cuda_core/tests/helpers/contexts.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from contextlib import contextmanager + +from cuda.core._utils.cuda_utils import driver, handle_return + +__all__ = [ + "assert_device_operations_use_bound_context", + "current_context_handle", + "no_current_context", + "use_context", +] + + +def current_context_handle(): + """Return the current CUDA context handle, or zero if none is current.""" + return int(handle_return(driver.cuCtxGetCurrent())) + + +def assert_device_operations_use_bound_context(device): + """Check that Device operations use its bound context and preserve the ambient context.""" + bound_context = device.context + ambient_context_handle = current_context_handle() + stream = event = builder = None + + try: + stream = device.create_stream() + assert stream.context == bound_context + assert current_context_handle() == ambient_context_handle + + event = device.create_event() + assert event.context == bound_context + assert current_context_handle() == ambient_context_handle + + builder = device.create_graph_builder() + assert builder.stream.context == bound_context + assert current_context_handle() == ambient_context_handle + + device.sync() + assert current_context_handle() == ambient_context_handle + + builder.close() + builder = None + assert current_context_handle() == ambient_context_handle + + event.close() + event = None + assert current_context_handle() == ambient_context_handle + + stream.close() + stream = None + assert current_context_handle() == ambient_context_handle + finally: + if builder is not None: + builder.close() + if event is not None: + event.close() + if stream is not None: + stream.close() + + +@contextmanager +def no_current_context(): + """Temporarily remove the calling thread's sole current CUDA context.""" + if current_context_handle() == 0: + raise RuntimeError("no_current_context requires a current CUDA context") + + previous = handle_return(driver.cuCtxPopCurrent()) + try: + if current_context_handle() != 0: + raise RuntimeError("no_current_context requires exactly one stacked CUDA context") + yield + finally: + handle_return(driver.cuCtxPushCurrent(previous)) + + +@contextmanager +def use_context(device, context): + """Temporarily make a context current and restore the previous context.""" + if current_context_handle() == 0: + raise RuntimeError("use_context requires a current CUDA context to restore") + + previous = device.set_current(context) + if previous is None: + raise RuntimeError("Device.set_current() did not return the previous CUDA context") + try: + yield + finally: + device.set_current(previous) diff --git a/cuda_core/tests/memory_ipc/test_peer_access.py b/cuda_core/tests/memory_ipc/test_peer_access.py index 992e01aa540..a82690d46b5 100644 --- a/cuda_core/tests/memory_ipc/test_peer_access.py +++ b/cuda_core/tests/memory_ipc/test_peer_access.py @@ -96,7 +96,11 @@ def test_main(self, ipc_mempool_device_x2, grant_access_in_parent): buffer.close() # TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()` + # Make dev0 current; Device.sync() must act on dev1 and leave dev0 current. + dev0.set_current() + assert Device().device_id == dev0.device_id dev1.sync() + assert Device().device_id == dev0.device_id mr.close() def child_main(self, mr, buffer): diff --git a/cuda_core/tests/test_device.py b/cuda_core/tests/test_device.py index 0d2e5e00952..48f5b3cf484 100644 --- a/cuda_core/tests/test_device.py +++ b/cuda_core/tests/test_device.py @@ -2,12 +2,18 @@ # SPDX-License-Identifier: Apache-2.0 import contextlib +from concurrent.futures import ThreadPoolExecutor import pytest +from helpers.contexts import ( + assert_device_operations_use_bound_context, + current_context_handle, + no_current_context, +) import cuda.core from cuda.bindings import driver, runtime -from cuda.core import Device +from cuda.core import Device, StreamOptions from cuda.core._utils.cuda_utils import ComputeCapability, handle_return from cuda.core._utils.version import driver_version @@ -100,6 +106,80 @@ def test_device_create_event(init_cuda): assert event.handle +@pytest.mark.agent_authored(model="gpt-5.6") +def test_device_operations_target_receiver_and_restore_current(device_x2): + dev0, dev1 = device_x2 + dev0.set_current() + dev1.set_current() + assert_device_operations_use_bound_context(dev0) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_device_operations_restore_no_current_context(deinit_cuda): + device = Device(0) + device.set_current() + + with no_current_context(): + assert_device_operations_use_bound_context(device) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_device_create_stream_restores_context_after_failure(device_x2): + dev0, dev1 = device_x2 + dev0.set_current() + dev1.set_current() + ctx1_handle = current_context_handle() + + with pytest.raises(ValueError, match="priority=.*out of range"): + dev0.create_stream(options=StreamOptions(priority=2**30)) + assert current_context_handle() == ctx1_handle + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_set_current_returns_previous_context_with_owning_device(device_x2): + dev0, dev1 = device_x2 + dev0.set_current() + ctx0 = dev0.context + dev1.set_current() + + previous = dev0.set_current(ctx0) + assert previous.handle == dev1.context.handle + dev1.set_current(previous) + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_device_receiver_switching_is_thread_local(device_x2): + dev0, dev1 = device_x2 + dev0.set_current() + main_context = current_context_handle() + + def worker(): + worker_dev0 = Device(dev0.device_id) + worker_dev0.set_current() + target_context = current_context_handle() + worker_dev1 = Device(dev1.device_id) + worker_dev1.set_current() + foreign_context = current_context_handle() + + stream = None + try: + stream = worker_dev0.create_stream() + resource_context = int(stream.context.handle) + finally: + if stream is not None: + stream.close() + + return resource_context, target_context, current_context_handle(), foreign_context + + with ThreadPoolExecutor(max_workers=1) as executor: + worker_result = executor.submit(worker).result() + + resource_context, target_context, restored_context, foreign_context = worker_result + assert resource_context == target_context + assert restored_context == foreign_context + assert current_context_handle() == main_context + + def test_pci_bus_id(): device = Device() bus_id = handle_return(runtime.cudaDeviceGetPCIBusId(13, device.device_id)) diff --git a/cuda_core/tests/test_green_context.py b/cuda_core/tests/test_green_context.py index 52fc372a27b..5b88083d22b 100644 --- a/cuda_core/tests/test_green_context.py +++ b/cuda_core/tests/test_green_context.py @@ -2,11 +2,9 @@ # # SPDX-License-Identifier: Apache-2.0 - -import contextlib - import numpy as np import pytest +from helpers.contexts import assert_device_operations_use_bound_context, use_context from cuda.core import ( ContextOptions, @@ -152,16 +150,6 @@ def _find_backfill_only_two_group_split(sm): return None -@contextlib.contextmanager -def _use_green_ctx(dev, ctx): - """Context manager: set green ctx current, restore previous on exit.""" - prev = dev.set_current(ctx) - try: - yield - finally: - dev.set_current(prev) - - @pytest.mark.agent_authored(model="gpt-5.6") def test_memory_node_updates_preserve_green_context( init_cuda, @@ -173,7 +161,7 @@ def test_memory_node_updates_preserve_green_context( memory_resource = LegacyPinnedMemoryResource() src = memory_resource.allocate(4) dst = memory_resource.allocate(4) - with _use_green_ctx(init_cuda, green_ctx): + with use_context(init_cuda, green_ctx): graph_def = GraphDefinition() memset_node = graph_def.memset(dst, 0, 4) memcpy_node = graph_def.memcpy(dst, src, 4) @@ -528,16 +516,48 @@ def test_stream_and_event_track_green_context(self, green_ctx): stream.sync() event.sync() + @pytest.mark.agent_authored(model="gpt-5.6") + def test_device_receiver_targets_stored_green_context(self, init_cuda, green_ctx): + primary_ctx = init_cuda.context + + with use_context(init_cuda, green_ctx): + handle_return(driver.cuCtxSetCurrent(primary_ctx.handle)) + assert_device_operations_use_bound_context(init_cuda) + + @pytest.mark.agent_authored(model="gpt-5.6") + def test_texture_rejects_resource_from_other_context(self, init_cuda, green_ctx): + from cuda.core.texture import ( + OpaqueArrayOptions, + ResourceDescriptor, + ) + from cuda.core.typing import ArrayFormatType + + array = init_cuda.create_opaque_array( + OpaqueArrayOptions( + shape=(8, 8), + format=ArrayFormatType.UINT8, + num_channels=4, + ) + ) + try: + with ( + use_context(init_cuda, green_ctx), + pytest.raises(ValueError, match="resource is not compatible with this Device object"), + ): + init_cuda.create_texture_object(resource=ResourceDescriptor.from_opaque_array(array)) + finally: + array.close() + def test_close_while_current_raises(self, init_cuda, green_ctx): """close() on a current context raises — test via set_current.""" dev = init_cuda - with _use_green_ctx(dev, green_ctx), pytest.raises(RuntimeError, match="while it is current"): + with use_context(dev, green_ctx), pytest.raises(RuntimeError, match="while it is current"): green_ctx.close() def test_set_current_swap_regression(self, init_cuda, green_ctx): """set_current still works (backward compat) and preserves identity.""" dev = init_cuda - with _use_green_ctx(dev, green_ctx): + with use_context(dev, green_ctx): pass # just verify push/pop works # Swap again and check identity round-trip prev = dev.set_current(green_ctx) diff --git a/cuda_core/tests/test_launcher.py b/cuda_core/tests/test_launcher.py index 2ab766cc2f0..2ee783016f2 100644 --- a/cuda_core/tests/test_launcher.py +++ b/cuda_core/tests/test_launcher.py @@ -24,7 +24,7 @@ StreamOptions, launch, ) -from cuda.core._memory._legacy import _SynchronousMemoryResource +from cuda.core._memory._device_memory_resource import _SynchronousMemoryResource from cuda.core._utils.cuda_utils import CUDAError from cuda.core.typing import ObjectCodeFormatType, SourceCodeType diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 44227d4b1c5..b6b219e002e 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -23,6 +23,7 @@ thread_unsafe_on_windows, ) from helpers.constants import POOL_SIZE +from helpers.contexts import current_context_handle, no_current_context from helpers.memory import ( create_managed_memory_resource_or_skip, create_pinned_memory_resource_or_xfail, @@ -739,14 +740,10 @@ def test_close_with_default_stream_requires_context(): # Use a real stream at creation so _init succeeds without a current context later. buf = Buffer.from_handle(1, 1024, mr=mr, stream=stream) - previous = handle_return(driver.cuCtxPopCurrent()) - assert int(previous) != 0 - try: - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + with no_current_context(): + assert current_context_handle() == 0 with pytest.raises(RuntimeError, match="no CUDA context is current"): buf.close(stream=default_stream()) - finally: - handle_return(driver.cuCtxSetCurrent(previous)) buf.close() # clean up using the recorded stream (which carries a context) @@ -758,14 +755,10 @@ def test_from_handle_mr_default_stream_requires_context(buffer_type): device = Device() device.set_current() mr = StubMemoryResource(device) - previous = handle_return(driver.cuCtxPopCurrent()) - assert int(previous) != 0 - try: - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + with no_current_context(): + assert current_context_handle() == 0 with pytest.raises(RuntimeError, match="no CUDA context is current"): buffer_type.from_handle(1, 1024, mr=mr) - finally: - handle_return(driver.cuCtxSetCurrent(previous)) @pytest.mark.agent_authored(model="gpt-5.6") @@ -777,15 +770,11 @@ def test_from_handle_mr_explicit_stream_without_current_context(buffer_type): stream = device.create_stream() CapturingMR, telemetry = make_instrumented_memory_resource(record_streams=True) mr = CapturingMR(device) - previous = handle_return(driver.cuCtxPopCurrent()) - assert int(previous) != 0 - try: - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + with no_current_context(): + assert current_context_handle() == 0 buf = buffer_type.from_handle(1, 1024, mr=mr, stream=stream) buf.close() - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 - finally: - handle_return(driver.cuCtxSetCurrent(previous)) + assert current_context_handle() == 0 assert telemetry["deallocations"][-1]["stream"].handle == stream.handle @@ -814,39 +803,31 @@ def test_mr_deallocation_without_current_context(init_cuda, capsys, replace_stre stream = init_cuda.create_stream() if replace_stream else None assert len(telemetry["active"]) == 1 - previous = handle_return(driver.cuCtxPopCurrent()) - assert int(previous) != 0 - try: - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + with no_current_context(): + assert current_context_handle() == 0 buf.close(stream) assert len(telemetry["active"]) == 0 - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + assert current_context_handle() == 0 assert "mr.deallocate() failed" not in capsys.readouterr().err - finally: - handle_return(driver.cuCtxSetCurrent(previous)) @pytest.mark.agent_authored(model="cursor-grok-4.5") @pytest.mark.parametrize("replace_stream", [False, True]) -def test_mr_deallocation_with_foreign_context(capsys, replace_stream): +def test_mr_deallocation_with_foreign_context(device_x2, capsys, replace_stream): """MR-backed Buffer teardown switches away from an unrelated current context.""" - if len(Device.get_all_devices()) < 2: - pytest.skip("Test requires at least 2 GPUs") - - alloc_dev = Device(0) + alloc_dev, foreign_dev = device_x2 alloc_dev.set_current() TrackingMR, telemetry = make_instrumented_memory_resource(DummyDeviceMemoryResource, track_active=True) mr = TrackingMR(alloc_dev) buf = mr.allocate(1024) stream = alloc_dev.create_stream() if replace_stream else None assert len(telemetry["active"]) == 1 - alloc_ctx = int(handle_return(driver.cuCtxGetCurrent())) + alloc_ctx = current_context_handle() - foreign_dev = Device(1) foreign_dev.set_current() - foreign_ctx = int(handle_return(driver.cuCtxGetCurrent())) + foreign_ctx = current_context_handle() assert foreign_ctx != 0 assert foreign_ctx != alloc_ctx @@ -854,7 +835,7 @@ def test_mr_deallocation_with_foreign_context(capsys, replace_stream): buf.close(stream) assert len(telemetry["active"]) == 0 - assert int(handle_return(driver.cuCtxGetCurrent())) == foreign_ctx + assert current_context_handle() == foreign_ctx assert "mr.deallocate() failed" not in capsys.readouterr().err finally: alloc_dev.set_current() @@ -886,21 +867,17 @@ def test_pool_buffer_deallocates_without_current_context(mempool_device, capfd): stream.sync() used_after_alloc = mr.attributes.used_mem_current - previous = handle_return(driver.cuCtxPopCurrent()) - assert int(previous) != 0 - try: - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + with no_current_context(): + assert current_context_handle() == 0 buf.close() stream.sync() assert mr.attributes.used_mem_current < used_after_alloc - assert int(handle_return(driver.cuCtxGetCurrent())) == 0 + assert current_context_handle() == 0 err = capfd.readouterr().err - assert "failed during resource destruction" not in err + assert "cuMemFreeAsync failed" not in err assert "mr.deallocate() failed" not in err - finally: - handle_return(driver.cuCtxSetCurrent(previous)) @pytest.mark.agent_authored(model="cursor-grok-4.5") @@ -914,16 +891,16 @@ def test_pool_buffer_deallocates_with_foreign_context(mempool_device_x2, capfd): buf = mr.allocate(size, stream=stream) stream.sync() used_after_alloc = mr.attributes.used_mem_current - alloc_ctx = int(handle_return(driver.cuCtxGetCurrent())) + alloc_ctx = current_context_handle() foreign_dev.set_current() - foreign_ctx = int(handle_return(driver.cuCtxGetCurrent())) + foreign_ctx = current_context_handle() assert foreign_ctx != 0 assert foreign_ctx != alloc_ctx try: buf.close() - assert int(handle_return(driver.cuCtxGetCurrent())) == foreign_ctx + assert current_context_handle() == foreign_ctx # Observe the free on the allocation device, then restore the foreign context. alloc_dev.set_current() @@ -932,7 +909,7 @@ def test_pool_buffer_deallocates_with_foreign_context(mempool_device_x2, capfd): foreign_dev.set_current() err = capfd.readouterr().err - assert "failed during resource destruction" not in err + assert "cuMemFreeAsync failed" not in err finally: alloc_dev.set_current() @@ -2189,7 +2166,7 @@ def test_legacy_pinned_device_id_raises(): def test_synchronous_memory_resource_basic(init_cuda): """_SynchronousMemoryResource exercises properties and allocate paths (zero, non-zero, with-stream).""" - from cuda.core._memory._legacy import _SynchronousMemoryResource + from cuda.core._memory._device_memory_resource import _SynchronousMemoryResource dev = Device() mr = _SynchronousMemoryResource(dev.device_id) @@ -2222,7 +2199,7 @@ def test_synchronous_memory_resource_basic(init_cuda): def test_synchronous_memory_resource_deallocate_accepts_stream(init_cuda): """_SynchronousMemoryResource.deallocate accepts an explicit stream.""" - from cuda.core._memory._legacy import _SynchronousMemoryResource + from cuda.core._memory._device_memory_resource import _SynchronousMemoryResource dev = Device() mr = _SynchronousMemoryResource(dev.device_id) @@ -2232,6 +2209,60 @@ def test_synchronous_memory_resource_deallocate_accepts_stream(init_cuda): stream.close() +@pytest.mark.agent_authored(model="gpt-5.6") +def test_synchronous_memory_resource_uses_its_context(device_x2): + """Synchronous allocation targets its stored context and restores the current one.""" + from cuda.core._memory._device_memory_resource import _SynchronousMemoryResource + + alloc_dev, current_dev = device_x2 + alloc_dev.set_current() + stream = alloc_dev.create_stream() + mr = _SynchronousMemoryResource(alloc_dev.device_id, alloc_dev.context) + + current_dev.set_current() + current_context = current_context_handle() + + buf = mr.allocate(64, stream=stream) + try: + pointer_context = handle_return( + driver.cuPointerGetAttribute( + driver.CUpointer_attribute.CU_POINTER_ATTRIBUTE_CONTEXT, + int(buf.handle), + ) + ) + pointer_device = handle_return( + driver.cuPointerGetAttribute( + driver.CUpointer_attribute.CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL, + int(buf.handle), + ) + ) + assert int(pointer_context) == int(alloc_dev.context.handle) + assert pointer_device == alloc_dev.device_id + assert current_context_handle() == current_context + finally: + buf.close(stream=stream) + stream.close() + + assert current_context_handle() == current_context + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_synchronous_memory_resource_restores_context_after_failure(device_x2): + """A failed synchronous allocation restores the context that was current.""" + from cuda.core._memory._device_memory_resource import _SynchronousMemoryResource + + alloc_dev, current_dev = device_x2 + alloc_dev.set_current() + mr = _SynchronousMemoryResource(alloc_dev.device_id, alloc_dev.context) + current_dev.set_current() + current_context = current_context_handle() + + with pytest.raises(CUDAError): + mr.allocate(sys.maxsize) + + assert current_context_handle() == current_context + + @pytest.mark.parametrize( ("method", "spec", "match"), [ diff --git a/cuda_core/tests/test_stream.py b/cuda_core/tests/test_stream.py index 6309c134ccc..bd909734e09 100644 --- a/cuda_core/tests/test_stream.py +++ b/cuda_core/tests/test_stream.py @@ -190,7 +190,7 @@ class MyStream(Stream): dev = Device() dev.set_current() - stream = MyStream._init(options=StreamOptions(), device_id=dev.device_id) + stream = MyStream._init(options=StreamOptions(), device_id=dev.device_id, ctx=dev.context) assert isinstance(stream, MyStream) diff --git a/cuda_core/tests/test_texture_surface.py b/cuda_core/tests/test_texture_surface.py index 7436b99647f..254eb2faf4d 100644 --- a/cuda_core/tests/test_texture_surface.py +++ b/cuda_core/tests/test_texture_surface.py @@ -5,6 +5,7 @@ import numpy as np import pytest +from helpers.contexts import current_context_handle, no_current_context import cuda.core from cuda.core import ( @@ -45,6 +46,110 @@ def test_resource_descriptor_init_disabled(): ResourceDescriptor() +@pytest.mark.agent_authored(model="gpt-5.6") +def test_texture_resources_target_receiver_context(device_x2): + dev0, dev1 = device_x2 + dev0.set_current() + ctx0 = dev0.context + dev1.set_current() + ctx1_handle = current_context_handle() + + array = dev0.create_opaque_array( + OpaqueArrayOptions( + shape=(8, 8), + format=ArrayFormatType.UINT8, + num_channels=4, + is_surface_load_store=True, + ) + ) + assert array.device == dev0 + assert current_context_handle() == ctx1_handle + + mipmap = dev0.create_mipmapped_array( + MipmappedArrayOptions( + shape=(8, 8), + format=ArrayFormatType.UINT8, + num_channels=4, + num_levels=2, + ) + ) + assert mipmap.device == dev0 + assert current_context_handle() == ctx1_handle + + level = mipmap.get_level(0) + assert level.device == dev0 + assert current_context_handle() == ctx1_handle + + resource = ResourceDescriptor.from_opaque_array(array) + texture = dev0.create_texture_object(resource=resource, options=TextureObjectOptions()) + surface = dev0.create_surface_object(resource=resource) + assert texture.device == dev0 + assert surface.device == dev0 + assert current_context_handle() == ctx1_handle + + surface.close() + texture.close() + level.close() + mipmap.close() + array.close() + assert current_context_handle() == ctx1_handle + assert int(ctx0.handle) != ctx1_handle + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_texture_resources_restore_no_current_context(deinit_cuda): + device = Device(0) + device.set_current() + + array = texture = None + try: + with no_current_context(): + array = device.create_opaque_array( + OpaqueArrayOptions( + shape=(8, 8), + format=ArrayFormatType.UINT8, + num_channels=4, + ) + ) + texture = device.create_texture_object(resource=ResourceDescriptor.from_opaque_array(array)) + assert current_context_handle() == 0 + + texture.close() + texture = None + array.close() + array = None + assert current_context_handle() == 0 + finally: + if texture is not None: + texture.close() + if array is not None: + array.close() + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_texture_creation_rejects_mismatched_receiver(device_x2): + dev0, dev1 = device_x2 + dev0.set_current() + array = dev0.create_opaque_array( + OpaqueArrayOptions( + shape=(8, 8), + format=ArrayFormatType.UINT8, + num_channels=4, + is_surface_load_store=True, + ) + ) + resource = ResourceDescriptor.from_opaque_array(array) + dev1.set_current() + ctx1_handle = current_context_handle() + + with pytest.raises(ValueError, match="resource belongs to device 0"): + dev1.create_texture_object(resource=resource) + with pytest.raises(ValueError, match="resource belongs to device 0"): + dev1.create_surface_object(resource=resource) + assert current_context_handle() == ctx1_handle + array.close() + + def test_array_2d_create_and_properties(init_cuda): arr = Device().create_opaque_array( OpaqueArrayOptions(shape=(32, 16), format=ArrayFormatType.FLOAT32, num_channels=1)