diff --git a/cuda_core/cuda/core/_cpp/resource_handles.cpp b/cuda_core/cuda/core/_cpp/resource_handles.cpp index ee116a9f353..46f7b019379 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; @@ -46,6 +54,7 @@ decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate = nullptr; decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority = nullptr; decltype(&cuStreamDestroy) p_cuStreamDestroy = nullptr; +decltype(&cuStreamGetCtx) p_cuStreamGetCtx = nullptr; decltype(&cuEventCreate) p_cuEventCreate = nullptr; decltype(&cuEventDestroy) p_cuEventDestroy = nullptr; @@ -128,46 +137,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 +197,235 @@ 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&&...); + if (!h_context) { + return CUDA_ERROR_INVALID_CONTEXT; + } + 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&&); + if (!h_context) { + return CUDA_ERROR_INVALID_CONTEXT; + } + 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)); + } else { + warn_on_cuda_error( + "cuCtxSetCurrent (restoring the caller's context)", composite, + "failed; cleanup of the new resource skipped because its context " + "is no longer current (resource leaked)"); } - 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; +} - CUresult status() const noexcept { return status_; } +#undef ASSERT_NOTHROW_INVOCABLE + +// 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 { + GILReleaseGuard gil; + 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 { + GILReleaseGuard gil; + return invoke_in_context(h_context, [&]() noexcept { + return p_cuCtxGetStreamPriorityRange(least_priority, greatest_priority); + }); +} + // ============================================================================ // CUDA user-object deferred cleanup // @@ -478,12 +655,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 +855,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 +877,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 +947,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{}; } @@ -783,6 +962,16 @@ StreamHandle get_per_thread_stream() { return handle; } +StreamHandle create_context_bound_legacy_stream(const ContextHandle& h_context) { + if (!h_context) { + return {}; + } + // Default deleter: this handle never owns CU_STREAM_LEGACY, so nothing + // needs to run when the last reference is released. + auto box = std::make_shared(StreamBox{CU_STREAM_LEGACY, h_context}); + return StreamHandle(box, &box->resource); +} + // ============================================================================ // Deallocation streams // @@ -797,12 +986,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 +997,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 +1026,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 +1064,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 +1076,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; } ); @@ -942,8 +1100,23 @@ EventHandle create_event_handle(const ContextHandle& h_ctx, unsigned int flags, return h; } -EventHandle create_event_handle_noctx(unsigned int flags) { - return create_event_handle(ContextHandle{}, flags, false, false, false, -1); +EventHandle create_event_handle_for_stream(CUstream stream, unsigned int flags) { + // Resolve the stream's owning context (for default-stream tokens this is + // the current context, per cuStreamGetCtx) and create the event there, so + // it can be recorded on `stream` no matter which context is current. + CUcontext ctx = nullptr; + { + GILReleaseGuard gil; + err = p_cuStreamGetCtx(stream, &ctx); + } + if (err != CUDA_SUCCESS) { + return {}; + } + if (!ctx) { + err = CUDA_ERROR_INVALID_CONTEXT; + return {}; + } + return create_event_handle(create_context_handle_ref(ctx), flags, false, false, false, -1); } EventHandle create_event_handle_ref(CUevent event) { @@ -967,7 +1140,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 +1253,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 +1280,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 +1288,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 +1310,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 +1318,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 +1331,15 @@ DevicePtrHandle deviceptr_alloc_async(size_t size, const StreamHandle& h_stream) return DevicePtrHandle(box, &box->resource); } -DevicePtrHandle deviceptr_alloc(size_t size) { +// Allocate device memory synchronously with the provided context current. +CUresult deviceptr_alloc_raw(CUdeviceptr* ptr, size_t size, + const ContextHandle& h_context) noexcept { GILReleaseGuard gil; - CUdeviceptr ptr; - if (CUDA_SUCCESS != (err = p_cuMemAlloc(&ptr, size))) { - return {}; - } - - auto box = std::shared_ptr( - new DevicePtrBox{ptr, DeallocationStream{}}, - [](DevicePtrBox* b) { - GILReleaseGuard gil; - p_cuMemFree(b->resource); - delete b; - } - ); - return DevicePtrHandle(box, &box->resource); + 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 +1403,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 +1442,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 +1538,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 +1547,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 +1570,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 +1578,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 +2837,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 +2857,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_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_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 +2922,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_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_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_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 +3004,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..419710ea0d9 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; @@ -77,6 +82,7 @@ extern decltype(&cuGreenCtxStreamCreate) p_cuGreenCtxStreamCreate; extern decltype(&cuStreamCreateWithPriority) p_cuStreamCreateWithPriority; extern decltype(&cuStreamDestroy) p_cuStreamDestroy; +extern decltype(&cuStreamGetCtx) p_cuStreamGetCtx; extern decltype(&cuEventCreate) p_cuEventCreate; extern decltype(&cuEventDestroy) p_cuEventDestroy; @@ -246,6 +252,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. Releases the GIL around the driver call. +// 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 // ============================================================================ @@ -287,6 +304,14 @@ StreamHandle get_legacy_stream(); // Note: Per-thread stream has no specific context dependency. StreamHandle get_per_thread_stream(); +// Wrap CU_STREAM_LEGACY with an explicit context, bypassing the "bind to +// whatever is current" resolution that a bare default-stream token uses (see +// make_deallocation_stream). Lets a resource that always operates in one +// known context (e.g. a synchronous, non-pooled allocator) record a correct +// deallocation context without requiring that context to be current when the +// token is created. Returns an empty handle for an empty h_context. +StreamHandle create_context_bound_legacy_stream(const ContextHandle& h_context); + // ============================================================================ // Event handle functions // ============================================================================ @@ -300,11 +325,14 @@ EventHandle create_event_handle(const ContextHandle& h_ctx, unsigned int flags, bool timing_enabled, bool is_blocking_sync, bool ipc_enabled, int device_id); -// Create an owning event handle without context dependency. -// Use for temporary events that are created and destroyed in the same scope. +// Create an owning event in the context that owns `stream`, so it can be +// recorded on that stream regardless of which context is current. Default- +// stream tokens resolve to the current context (cuStreamGetCtx semantics). +// Use for temporary ordering events that are created and destroyed in the +// same scope; the handle carries no device id. // When the last reference is released, cuEventDestroy is called automatically. // Returns empty handle on error (caller must check). -EventHandle create_event_handle_noctx(unsigned int flags); +EventHandle create_event_handle_for_stream(CUstream stream, unsigned int flags); // Create an owning event handle from an IPC handle. // The originating process owns the event and its context. @@ -371,10 +399,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 +768,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 +778,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 +790,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 ab6bef5cbc1..8c2b273a6cd 100644 --- a/cuda_core/cuda/core/_device.pyi +++ b/cuda_core/cuda/core/_device.pyi @@ -583,7 +583,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. @@ -591,6 +591,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.""" @@ -607,6 +610,12 @@ class Device: Providing a `ctx` causes the previous set context to be popped and returned. + If `ctx` was created on a different device than this receiver, the call + is delegated to that device's own :meth:`set_current`. This keeps the + owning device's bookkeeping consistent and lets a context this method + handed out for a foreign device be pushed back through any ``Device`` + object, matching the CUDA context stack's own thread-wide semantics. + Parameters ---------- ctx : :obj:`~_context.Context`, optional @@ -615,7 +624,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 -------- @@ -647,7 +658,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: @@ -659,7 +670,7 @@ class Device: Note ---- - Device must be initialized. + Device must be initialized. New streams are created on this device. Parameters ---------- @@ -675,7 +686,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 ---- @@ -718,7 +729,12 @@ class Device: """ def sync(self) -> None: - """Synchronize the device. + """Synchronize this device's bound context. + + Waits for all preceding work in this device's bound :obj:`~_context.Context` + to complete. Only that context is synchronized, not the device as a + whole; work queued in a different context on the same device (e.g. a + green context) is unaffected. Note ---- @@ -726,7 +742,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 ------- @@ -735,12 +751,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 ---- @@ -759,12 +773,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 ---- @@ -783,15 +795,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 ---- @@ -812,15 +822,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..a7e7d59e04a 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._synchronous_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() @@ -1253,6 +1263,12 @@ class Device: Providing a `ctx` causes the previous set context to be popped and returned. + If `ctx` was created on a different device than this receiver, the call + is delegated to that device's own :meth:`set_current`. This keeps the + owning device's bookkeeping consistent and lets a context this method + handed out for a foreign device be pushed back through any ``Device`` + object, matching the CUDA context stack's own thread-wide semantics. + Parameters ---------- ctx : :obj:`~_context.Context`, optional @@ -1261,7 +1277,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 +1294,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: @@ -1283,16 +1302,19 @@ class Device: assert_type(ctx, Context) Context_check_open(ctx) if ctx._device_id != self._device_id: - raise RuntimeError( - "the provided context was created on the device with" - f" id={ctx._device_id}, which is different from the target id={self._device_id}" - ) + # The CUDA context stack is per-thread, not per-Device-object, + # so pushing/popping a foreign-device context is delegated to + # the device that owns it; its own bookkeeping (_context, + # _has_inited) is what should track this push, not ours. + return Device(ctx._device_id).set_current(ctx) 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 +1322,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 +1404,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 +1416,7 @@ class Device: Note ---- - Device must be initialized. + Device must be initialized. New streams are created on this device. Parameters ---------- @@ -1412,7 +1435,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 +1485,12 @@ class Device: return self.memory_resource.allocate(size, stream=stream) def sync(self) -> None: - """Synchronize the device. + """Synchronize this device's bound context. + + Waits for all preceding work in this device's bound :obj:`~_context.Context` + to complete. Only that context is synchronized, not the device as a + whole; work queued in a different context on the same device (e.g. a + green context) is unaffected. Note ---- @@ -1470,10 +1498,11 @@ class Device: """ self._check_context_initialized() - handle_return(runtime.cudaDeviceSynchronize()) + cdef Context ctx = self._context + HANDLE_RETURN(context_synchronize(ctx._h_context)) 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 +1516,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 +1540,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 +1567,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 +1601,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 +1632,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/_buffer.pyi b/cuda_core/cuda/core/_memory/_buffer.pyi index 7b7a58e9558..114e15a8de3 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyi +++ b/cuda_core/cuda/core/_memory/_buffer.pyi @@ -250,7 +250,7 @@ class Buffer: def __release_buffer__(self, buffer: memoryview, /) -> None: ... @property def device_id(self) -> int: - """Return the device ordinal of this buffer.""" + """Return the device ordinal of this buffer, or -1 for memory not bound to a device.""" @property def handle(self) -> int: """Return the buffer handle object. diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index 316f07f2912..2484ad82b00 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -666,7 +666,7 @@ cdef class Buffer: @property def device_id(self) -> int: - """Return the device ordinal of this buffer.""" + """Return the device ordinal of this buffer, or -1 for memory not bound to a device.""" Buffer_check_open(self) if self._memory_resource is not None: return self._memory_resource.device_id diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx index d72b0e45ebc..62dc4f9e747 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx @@ -11,11 +11,7 @@ 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 ( - as_cu, - get_device_mempool, - get_last_error, -) +from cuda.core._resource_handles cimport as_cu, get_device_mempool, get_last_error from cuda.core._utils.cuda_utils cimport ( check_or_create_options, HANDLE_RETURN, diff --git a/cuda_core/cuda/core/_memory/_legacy.py b/cuda_core/cuda/core/_memory/_legacy.py index 4acbcb54e3a..f3dff33a133 100644 --- a/cuda_core/cuda/core/_memory/_legacy.py +++ b/cuda_core/cuda/core/_memory/_legacy.py @@ -94,49 +94,5 @@ def is_host_accessible(self) -> bool: @property 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 + """Return -1. Pinned memory is host memory and is not bound to a specific device.""" + return -1 diff --git a/cuda_core/cuda/core/_memory/_synchronous_memory_resource.pyi b/cuda_core/cuda/core/_memory/_synchronous_memory_resource.pyi new file mode 100644 index 00000000000..73136b896b9 --- /dev/null +++ b/cuda_core/cuda/core/_memory/_synchronous_memory_resource.pyi @@ -0,0 +1,23 @@ +# This file was generated by stubgen-pyx v0.2.22 from cuda_core/cuda/core/_memory/_synchronous_memory_resource.pyx + +from cuda.core._context import Context +from cuda.core._memory._buffer import Buffer, MemoryResource +from cuda.core._stream import Stream +from cuda.core.graph import GraphBuilder +from cuda.core.typing import DevicePointerType + +__all__ = [] + +class _SynchronousMemoryResource(MemoryResource): + __slots__ = ('_context', '_device_id') + + def __init__(self, device_id: int, context=None) -> None: ... + def _resolve_context(self) -> Context: ... + 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: ... diff --git a/cuda_core/cuda/core/_memory/_synchronous_memory_resource.pyx b/cuda_core/cuda/core/_memory/_synchronous_memory_resource.pyx new file mode 100644 index 00000000000..f02f38f69b1 --- /dev/null +++ b/cuda_core/cuda/core/_memory/_synchronous_memory_resource.pyx @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +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._resource_handles cimport ( + ContextHandle, + create_context_bound_legacy_stream, + deviceptr_alloc_raw, + get_last_error, + get_primary_context, +) +from cuda.core._stream cimport Stream, Stream_accept, Stream_is_default_token +from cuda.core._utils.cuda_utils cimport HANDLE_RETURN + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from cuda.core.graph import GraphBuilder + from cuda.core.typing import DevicePointerType + +__all__ = [] + + +class _SynchronousMemoryResource(MemoryResource): + __slots__ = ("_context", "_device_id") + + def __init__(self, device_id: int, context=None) -> None: + from .._device import Device + + self._device_id = Device(device_id).device_id + # Resolved lazily (in _resolve_context) so that construction with + # context=None does no CUDA work; the primary context is retained + # only once actually needed, on the first allocate()/deallocate(). + self._context = context + + def _resolve_context(self) -> Context: + cdef ContextHandle h_context + if self._context is None: + h_context = get_primary_context(self._device_id) + if not h_context: + HANDLE_RETURN(get_last_error()) + self._context = Context._from_handle( + Context, h_context, self._device_id) + return self._context + + def allocate( + self, + size_t size, + *, + stream: Stream | GraphBuilder | None = None, + ) -> Buffer: + # cuMemAlloc/cuMemFree are synchronous; a caller-supplied stream is + # accepted (and validated) for interface conformance and, if it is a + # real stream, recorded as the stream that orders deallocation. + cdef Context context = self._resolve_context() + cdef Stream dealloc_stream = None + if stream is not None: + dealloc_stream = Stream_accept(stream) + if dealloc_stream is None or Stream_is_default_token(dealloc_stream): + # A default-stream token carries no context of its own; Buffer._init + # would bind it to whichever context is current when it records the + # deallocation stream (and fail if none is). Bind it to this + # resource's context instead, so Buffer teardown frees in the right + # context no matter what is current then. Always the legacy token: + # a per-thread token would also arm the cross-thread PTDS warning, + # which is noise for a synchronous resource. + dealloc_stream = Stream._from_handle( + Stream, create_context_bound_legacy_stream(context._h_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, stream=dealloc_stream) + + def deallocate( + self, + ptr: DevicePointerType, + size_t size, + *, + stream: Stream | GraphBuilder | None = None, + ) -> None: + if stream is not None: + Stream_accept(stream).sync() + # No context switch here, by design (settled in the review of #2750): + # cuMemFree does not need a current context. The driver resolves the + # allocation's owning context from the pointer through unified + # addressing and frees it there (cuapiMemFree_common: "a current + # context is not required to free the device memory"). On the Buffer + # teardown path the C++ deleter has additionally already made the + # recorded deallocation context, bound by allocate() above, current. + 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 diff --git a/cuda_core/cuda/core/_memoryview.pyx b/cuda_core/cuda/core/_memoryview.pyx index 6b287f68c1a..c4a49d76946 100644 --- a/cuda_core/cuda/core/_memoryview.pyx +++ b/cuda_core/cuda/core/_memoryview.pyx @@ -26,8 +26,9 @@ import numpy from cuda.bindings cimport cydriver from cuda.core._resource_handles cimport ( EventHandle, - create_event_handle_noctx, + create_event_handle_for_stream, as_cu, + get_last_error, ) from cuda.core._utils.cuda_utils import handle_return, driver @@ -1227,7 +1228,12 @@ cpdef StridedMemoryView view_as_cai(obj, stream_ptr, view=None): # establish stream order if producer_s != consumer_s: with nogil: - h_event = create_event_handle_noctx(cydriver.CUevent_flags.CU_EVENT_DISABLE_TIMING) + # The event must belong to the producer stream's context to + # be recorded on it, whatever context is current here. + h_event = create_event_handle_for_stream( + producer_s, cydriver.CUevent_flags.CU_EVENT_DISABLE_TIMING) + if not h_event: + HANDLE_RETURN(get_last_error()) HANDLE_RETURN(cydriver.cuEventRecord( as_cu(h_event), producer_s)) HANDLE_RETURN(cydriver.cuStreamWaitEvent( diff --git a/cuda_core/cuda/core/_resource_handles.pxd b/cuda_core/cuda/core/_resource_handles.pxd index 568af27ac2e..fe075b6414b 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( @@ -189,13 +195,16 @@ cdef void retry_deferred_cleanup() noexcept cdef ContextHandle get_stream_context(const StreamHandle& h) noexcept nogil cdef StreamHandle get_legacy_stream() except+ nogil cdef StreamHandle get_per_thread_stream() except+ nogil +cdef StreamHandle create_context_bound_legacy_stream( + const ContextHandle& h_context) except+ nogil # Event handles cdef EventHandle create_event_handle( const ContextHandle& h_ctx, unsigned int flags, bint timing_enabled, bint is_blocking_sync, bint ipc_enabled, int device_id) except+ nogil -cdef EventHandle create_event_handle_noctx(unsigned int flags) except+ nogil +cdef EventHandle create_event_handle_for_stream( + cydriver.CUstream stream, unsigned int flags) except+ nogil cdef EventHandle create_event_handle_ref(cydriver.CUevent event) except+ nogil cdef EventHandle create_event_handle_ipc( const cydriver.CUipcEventHandle& ipc_handle, bint is_blocking_sync) except+ nogil @@ -219,7 +228,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 +335,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..0f8d6e15cde 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" ( @@ -67,14 +73,16 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": const StreamHandle& h) noexcept nogil StreamHandle get_legacy_stream "cuda_core::get_legacy_stream" () except+ nogil StreamHandle get_per_thread_stream "cuda_core::get_per_thread_stream" () except+ nogil + StreamHandle create_context_bound_legacy_stream "cuda_core::create_context_bound_legacy_stream" ( + const ContextHandle& h_context) except+ nogil # Event handles (note: _create_event_handle* are internal due to C++ overloading) EventHandle create_event_handle "cuda_core::create_event_handle" ( const ContextHandle& h_ctx, unsigned int flags, bint timing_enabled, bint is_blocking_sync, bint ipc_enabled, int device_id) except+ nogil - EventHandle create_event_handle_noctx "cuda_core::create_event_handle_noctx" ( - unsigned int flags) except+ nogil + EventHandle create_event_handle_for_stream "cuda_core::create_event_handle_for_stream" ( + cydriver.CUstream stream, unsigned int flags) except+ nogil EventHandle create_event_handle_ref "cuda_core::create_event_handle_ref" ( cydriver.CUevent event) except+ nogil EventHandle create_event_handle_ipc "cuda_core::create_event_handle_ipc" ( @@ -107,7 +115,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 +262,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 +312,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)" @@ -311,6 +332,7 @@ cdef extern from "_cpp/resource_handles.hpp" namespace "cuda_core": # Stream void* p_cuStreamCreateWithPriority "reinterpret_cast(cuda_core::p_cuStreamCreateWithPriority)" void* p_cuStreamDestroy "reinterpret_cast(cuda_core::p_cuStreamDestroy)" + void* p_cuStreamGetCtx "reinterpret_cast(cuda_core::p_cuStreamGetCtx)" # Event void* p_cuEventCreate "reinterpret_cast(cuda_core::p_cuEventCreate)" @@ -408,11 +430,12 @@ 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 + global p_cuStreamCreateWithPriority, p_cuStreamDestroy, p_cuStreamGetCtx global p_cuEventCreate, p_cuEventDestroy, p_cuIpcOpenEventHandle global p_cuDeviceGetCount global p_cuMemPoolSetAccess, p_cuMemPoolDestroy, p_cuMemPoolCreate @@ -435,11 +458,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") @@ -449,6 +478,7 @@ cdef void _init_driver_fn_pointers() noexcept: # Stream p_cuStreamCreateWithPriority = _get_driver_fn("cuStreamCreateWithPriority") p_cuStreamDestroy = _get_driver_fn("cuStreamDestroy") + p_cuStreamGetCtx = _get_driver_fn("cuStreamGetCtx") # Event p_cuEventCreate = _get_driver_fn("cuEventCreate") diff --git a/cuda_core/cuda/core/_stream.pyx b/cuda_core/cuda/core/_stream.pyx index 7035aa0f59c..e662d67c87f 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 @@ -29,9 +32,10 @@ from cuda.core._resource_handles cimport ( EventHandle, StreamHandle, create_context_handle_ref, - create_event_handle_noctx, + create_event_handle_for_stream, 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 @@ -153,14 +160,8 @@ cdef class Stream: # TODO: we might want to consider memoizing high/low per CUDA context and avoid this call 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 @@ -366,9 +367,14 @@ cdef class Stream: f" got {type(event_or_stream)}" ) from e - # Wait on stream via temporary event + # Wait on stream via a temporary event created in that stream's own + # context; an event from the current context would be rejected by + # cuEventRecord when the streams live on different devices. with nogil: - h_event = create_event_handle_noctx(cydriver.CUevent_flags.CU_EVENT_DISABLE_TIMING) + h_event = create_event_handle_for_stream( + as_cu(stream._h_stream), cydriver.CUevent_flags.CU_EVENT_DISABLE_TIMING) + if not h_event: + HANDLE_RETURN(get_last_error()) HANDLE_RETURN(cydriver.cuEventRecord(as_cu(h_event), as_cu(stream._h_stream))) # TODO: support flags other than 0? HANDLE_RETURN(cydriver.cuStreamWaitEvent(as_cu(self._h_stream), as_cu(h_event), 0)) diff --git a/cuda_core/cuda/core/_tensor_bridge.pyx b/cuda_core/cuda/core/_tensor_bridge.pyx index c7a6743213c..ae7a6794507 100644 --- a/cuda_core/cuda/core/_tensor_bridge.pyx +++ b/cuda_core/cuda/core/_tensor_bridge.pyx @@ -56,8 +56,9 @@ from cuda.core._layout cimport _StridedLayout from cuda.bindings cimport cydriver from cuda.core._resource_handles cimport ( EventHandle, - create_event_handle_noctx, + create_event_handle_for_stream, as_cu, + get_last_error, ) from cuda.core._utils.cuda_utils cimport HANDLE_RETURN @@ -318,8 +319,12 @@ cpdef int sync_torch_stream(int32_t device_index, b"aoti_torch_get_current_cuda_stream") if producer_s != consumer_s: with nogil: - h_event = create_event_handle_noctx( - cydriver.CUevent_flags.CU_EVENT_DISABLE_TIMING) + # The event must belong to the producer stream's context to be + # recorded on it, whatever context is current here. + h_event = create_event_handle_for_stream( + producer_s, cydriver.CUevent_flags.CU_EVENT_DISABLE_TIMING) + if not h_event: + HANDLE_RETURN(get_last_error()) HANDLE_RETURN(cydriver.cuEventRecord( as_cu(h_event), producer_s)) HANDLE_RETURN(cydriver.cuStreamWaitEvent( diff --git a/cuda_core/cuda/core/texture/_array.pyi b/cuda_core/cuda/core/texture/_array.pyi index 17cebc48963..e18adde1393 100644 --- a/cuda_core/cuda/core/texture/_array.pyi +++ b/cuda_core/cuda/core/texture/_array.pyi @@ -4,6 +4,7 @@ from dataclasses import dataclass import numpy 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)} @@ -162,8 +163,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 d4fc1be1888..72788bcd272 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 7b509771213..ff92f1edce2 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.22 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 865f982916b..50c9bb50816 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..28ddf2d6aa8 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 # -1 for memory not bound to a device 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 # -1 for memory not bound to a device 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.3.0-notes.rst b/cuda_core/docs/source/release/1.3.0-notes.rst new file mode 100644 index 00000000000..37ff06b34e5 --- /dev/null +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -0,0 +1,37 @@ +.. SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +.. SPDX-License-Identifier: Apache-2.0 + +.. currentmodule:: cuda.core + +``cuda.core`` 1.3.0 Release Notes +================================== + +Fixes and enhancements +---------------------- + +- :class:`Device` methods that create resources or synchronize now act on that + device's bound context, even when another device is current. They do not + change which device is current. For example, ``dev1.sync()`` synchronizes + device 1's bound context even when device 0 is current, and no longer + touches other contexts on device 1 (such as a green context). :meth:`Device.set_current` + now always returns a :class:`Context` with the correct device ID; passing a + context created on a different device than the receiver now delegates to + that device's own :meth:`~Device.set_current` instead of raising, so a + context this method returns can always be pushed back through any + :class:`Device` object. + (`#2311 `__) + +- :meth:`Stream.wait` given a stream now works when that stream belongs to a + device that is not current. The temporary ordering event is created in the + waited-on stream's context rather than the current one, which + ``cuEventRecord`` rejects when the two differ. The stream ordering applied + when importing foreign arrays and tensors uses the producer stream's context + the same way. + (`#2311 `__) + +- :attr:`LegacyPinnedMemoryResource.device_id` now returns ``-1``, as + documented for memory that is not bound to a device and as + :class:`PinnedMemoryResource` already does, instead of raising + ``RuntimeError``. :attr:`Buffer.device_id` on such a buffer returns ``-1`` + as well, which also lets a pinned buffer back a linear or pitched texture + resource. diff --git a/cuda_core/tests/conftest.py b/cuda_core/tests/conftest.py index 435e6761898..a8678292422 100644 --- a/cuda_core/tests/conftest.py +++ b/cuda_core/tests/conftest.py @@ -87,6 +87,8 @@ def wrapper(*args, **kwargs): kwargs["mempool_device_x2"] = _mempool_device_impl(2) if "mempool_device_x3" in kwargs: kwargs["mempool_device_x3"] = _mempool_device_impl(3) + if "device_x2" in kwargs: + kwargs["device_x2"] = _device_x2_impl() # These are used by test_green_context.py. The original fixtures include # pytest.skip() but that should have correctly fired by this time. @@ -214,6 +216,24 @@ def deinit_cuda(): _ = _device_unset_current() +def _device_x2_impl(): + devices = Device.get_all_devices() + if len(devices) < 2: + pytest.skip("Test requires at least 2 CUDA devices") + return devices[:2] + + +@pytest.fixture +def device_x2(init_cuda): + """Provide two CUDA devices, or skip when fewer are available. + + Depends on ``init_cuda`` so that, under pytest-run-parallel, the test is + wrapped by ``_wrap_worker_cuda_test`` and the devices are re-fetched on the + worker thread (Device objects are thread-local). + """ + return _device_x2_impl() + + @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..7ee01bb255f --- /dev/null +++ b/cuda_core/tests/helpers/contexts.py @@ -0,0 +1,127 @@ +# 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_event_record_rejected_from_ambient_context(event): + """Assert that recording ``event`` into a stream from the ambient context fails. + + An event and the stream it records must belong to the same context; + otherwise cuEventRecord fails with CUDA_ERROR_INVALID_HANDLE. cuStreamCreate + creates the probe stream in whatever context is currently ambient, so this + is a live check that ``event`` was not actually recorded from there. + """ + ambient_stream = handle_return(driver.cuStreamCreate(0)) + try: + (record_status,) = driver.cuEventRecord(event.handle, ambient_stream) + assert record_status == driver.CUresult.CUDA_ERROR_INVALID_HANDLE, ( + "Recording an event into a stream from a different context should fail with " + f"CUDA_ERROR_INVALID_HANDLE, got {record_status!r}" + ) + finally: + handle_return(driver.cuStreamDestroy(ambient_stream)) + + +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() + assert int(bound_context.handle) != ambient_context_handle, ( + "Precondition failed: the device's bound context must not be the current (ambient) context." + ) + stream = event = builder = None + + try: + stream = device.create_stream() + assert stream.context == bound_context + assert current_context_handle() == ambient_context_handle + # Live check: query the driver directly, rather than comparing cached + # metadata, to confirm the stream was actually created in the bound + # context rather than whatever was ambient. + driver_stream_ctx = handle_return(driver.cuStreamGetCtx(stream.handle)) + assert int(driver_stream_ctx) == int(bound_context.handle), ( + "cuStreamGetCtx reports a context other than the one the stream was created in." + ) + + event = device.create_event() + assert event.context == bound_context + assert current_context_handle() == ambient_context_handle + # Only exercised when the ambient context belongs to a different + # physical device: cuEventRecord's cross-context rejection is + # guaranteed distinct there. Two contexts on the *same* device (e.g. + # a green context vs. the primary context) may resolve to the same + # underlying device context for this check, so skip rather than + # assert unverified driver behavior. + if ambient_context_handle and int(handle_return(driver.cuCtxGetDevice())) != device.device_id: + _assert_event_record_rejected_from_ambient_context(event) + + 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..dbfb9fed5a9 100644 --- a/cuda_core/tests/test_device.py +++ b/cuda_core/tests/test_device.py @@ -2,12 +2,19 @@ # 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, +) +from helpers.nanosleep_kernel import NanosleepKernel 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 +107,131 @@ 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="claude-sonnet-5") +def test_set_current_round_trips_through_a_different_device(device_x2): + """The pre-#2311 idiom `prev = dev.set_current(ctx); ...; dev.set_current(prev)` + must keep working even when `prev` belongs to a different device than + `dev`: set_current() delegates to the context's owning device instead of + raising, so restoring through the original Device handle round-trips.""" + dev0, dev1 = device_x2 + dev0.set_current() + ctx0 = dev0.context + + dev1.set_current() + ctx1 = dev1.context + + dev0.set_current() # dev0 current again; dev1.context is still ctx1 + + prev = dev1.set_current(ctx1) + assert prev.handle == ctx0.handle + assert current_context_handle() == int(ctx1.handle) + + restored = dev1.set_current(prev) + assert restored.handle == ctx1.handle + assert current_context_handle() == int(ctx0.handle) + + +@pytest.mark.agent_authored(model="claude-sonnet-5") +def test_device_sync_waits_for_bound_context_work(device_x2): + """dev0.sync() must wait for work queued on dev0's bound context even + while dev1 is ambient, not just preserve the ambient context (#2311).""" + dev0, dev1 = device_x2 + if dev0.compute_capability.major < 7: + pytest.skip("__nanosleep is only available starting Volta (sm70)") + dev0.set_current() + stream = dev0.create_stream() + nanosleep = NanosleepKernel(dev0, sleep_duration_ms=20) + event = None + try: + nanosleep.launch(stream) + event = stream.record() + + dev1.set_current() + ambient_context_handle = current_context_handle() + + dev0.sync() + assert event.is_done + assert current_context_handle() == ambient_context_handle + finally: + if event is not None: + event.close() + stream.close() + + +@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..5f1954c6b58 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,45 @@ 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 + + with ( + init_cuda.create_opaque_array( + OpaqueArrayOptions( + shape=(8, 8), + format=ArrayFormatType.UINT8, + num_channels=4, + ) + ) as array, + 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)) + 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..083bdcbee8a 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._synchronous_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..769d780b36b 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() @@ -2180,16 +2157,15 @@ def test_legacy_pinned_allocate_zero_size(init_cuda): assert int(buf.handle) == 0 -def test_legacy_pinned_device_id_raises(): - """LegacyPinnedMemoryResource.device_id raises; pinned memory is not bound to a GPU.""" +def test_legacy_pinned_device_id_is_not_applicable(): + """LegacyPinnedMemoryResource.device_id is -1, as documented for memory not bound to a device.""" mr = LegacyPinnedMemoryResource() - with pytest.raises(RuntimeError, match="not bound to any GPU"): - _ = mr.device_id + assert mr.device_id == -1 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._synchronous_memory_resource import _SynchronousMemoryResource dev = Device() mr = _SynchronousMemoryResource(dev.device_id) @@ -2222,7 +2198,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._synchronous_memory_resource import _SynchronousMemoryResource dev = Device() mr = _SynchronousMemoryResource(dev.device_id) @@ -2232,6 +2208,101 @@ 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._synchronous_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._synchronous_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.agent_authored(model="claude-sonnet-5") +def test_synchronous_memory_resource_default_stream_deallocates_in_own_context(device_x2, capsys): + """Buffer teardown with no explicit stream frees in the resource's own + context, not whatever context happens to be current at close() time.""" + from cuda.core._memory._synchronous_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() + + buf = mr.allocate(64) # no explicit stream: records a context-bound default token + assert current_context_handle() == current_context + + buf.close() # no explicit stream: reuses the recorded token + assert current_context_handle() == current_context + assert capsys.readouterr().err == "" + + +@pytest.mark.agent_authored(model="claude-sonnet-5") +def test_synchronous_memory_resource_allocate_without_current_context(device_x2, capsys): + """allocate()/close() with no explicit stream succeed with no context + current, instead of raising or leaking the allocation (#2311).""" + from cuda.core._memory._synchronous_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() + + with no_current_context(): + buf = mr.allocate(64) + assert current_context_handle() == 0 + buf.close() + assert current_context_handle() == 0 + + assert capsys.readouterr().err == "" + + @pytest.mark.parametrize( ("method", "spec", "match"), [ diff --git a/cuda_core/tests/test_stream.py b/cuda_core/tests/test_stream.py index 6309c134ccc..02e817cf85b 100644 --- a/cuda_core/tests/test_stream.py +++ b/cuda_core/tests/test_stream.py @@ -74,6 +74,31 @@ def test_stream_wait_event(init_cuda): s2.sync() +@pytest.mark.agent_authored(model="claude-fable-5-1") +def test_stream_wait_stream_on_other_device(device_x2): + """Stream.wait(other_stream) must work when the streams live on different + devices and neither device is necessarily current: the temporary ordering + event has to be created in the *recorded* stream's context, since + cuEventRecord rejects an event from another context (#2311).""" + from helpers.contexts import current_context_handle + + dev0, dev1 = device_x2 + dev0.set_current() + s0 = dev0.create_stream() + dev1.set_current() + s1 = dev1.create_stream() + ambient = current_context_handle() + try: + s1.wait(s0) # dev1 current: s0's device is not current + s0.wait(s1) # dev1 current: self's device is not current + s0.sync() + s1.sync() + assert current_context_handle() == ambient + finally: + s0.close() + s1.close() + + def test_stream_wait_invalid_event(init_cuda): stream = Device().create_stream(options=StreamOptions()) with pytest.raises(ValueError): @@ -190,7 +215,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..42ae716e8f3 100644 --- a/cuda_core/tests/test_texture_surface.py +++ b/cuda_core/tests/test_texture_surface.py @@ -5,10 +5,12 @@ import numpy as np import pytest +from helpers.contexts import current_context_handle, no_current_context import cuda.core from cuda.core import ( Device, + LegacyPinnedMemoryResource, ) from cuda.core.texture import ( MipmappedArrayOptions, @@ -45,6 +47,107 @@ 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() + + with dev0.create_opaque_array( + OpaqueArrayOptions( + shape=(8, 8), + format=ArrayFormatType.UINT8, + num_channels=4, + is_surface_load_store=True, + ) + ) as array: + assert array.device == dev0 + assert current_context_handle() == ctx1_handle + with dev0.create_mipmapped_array( + MipmappedArrayOptions( + shape=(8, 8), + format=ArrayFormatType.UINT8, + num_channels=4, + num_levels=2, + ) + ) as mipmap: + assert mipmap.device == dev0 + assert current_context_handle() == ctx1_handle + with mipmap.get_level(0) as level: + assert level.device == dev0 + assert current_context_handle() == ctx1_handle + resource = ResourceDescriptor.from_opaque_array(array) + with ( + dev0.create_texture_object(resource=resource, options=TextureObjectOptions()) as texture, + dev0.create_surface_object(resource=resource) as surface, + ): + assert texture.device == dev0 + assert surface.device == dev0 + assert current_context_handle() == ctx1_handle + 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() + + with no_current_context(): + with ( + device.create_opaque_array( + OpaqueArrayOptions( + shape=(8, 8), + format=ArrayFormatType.UINT8, + num_channels=4, + ) + ) as array, + device.create_texture_object(resource=ResourceDescriptor.from_opaque_array(array)) as texture, + ): + assert array.device == device + assert texture.device == device + assert current_context_handle() == 0 + assert current_context_handle() == 0 + + +@pytest.mark.agent_authored(model="gpt-5.6") +def test_texture_creation_rejects_mismatched_receiver(device_x2): + dev0, dev1 = device_x2 + dev0.set_current() + with dev0.create_opaque_array( + OpaqueArrayOptions( + shape=(8, 8), + format=ArrayFormatType.UINT8, + num_channels=4, + is_surface_load_store=True, + ) + ) as array: + 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 + + +@pytest.mark.agent_authored(model="claude-sonnet-5") +def test_texture_linear_accepts_pinned_buffer(init_cuda): + """Pinned memory is device-accessible and not bound to any device, so a + pinned buffer is a valid linear texture backing whose device check is + skipped (device_id == -1) rather than failed (#2311).""" + mr = LegacyPinnedMemoryResource() + with mr.allocate(256) as buf: + assert buf.device_id == -1 + + resource = ResourceDescriptor.from_linear(buf, format=ArrayFormatType.UINT8, num_channels=1) + with init_cuda.create_texture_object(resource=resource) as texture: + assert texture.device == init_cuda + + def test_array_2d_create_and_properties(init_cuda): arr = Device().create_opaque_array( OpaqueArrayOptions(shape=(32, 16), format=ArrayFormatType.FLOAT32, num_channels=1)