diff --git a/cuda-core/src/simt/context.rs b/cuda-core/src/simt/context.rs index 2a76c5a389..187d5567c4 100644 --- a/cuda-core/src/simt/context.rs +++ b/cuda-core/src/simt/context.rs @@ -410,6 +410,7 @@ impl CudaContext { Arc::new(CudaStream { cu_stream: std::ptr::null_mut(), ctx: self.clone(), + owned: true, }) } @@ -480,6 +481,7 @@ impl CudaContext { Ok(Arc::new(CudaStream { cu_stream, ctx: self.clone(), + owned: true, })) } diff --git a/cuda-core/src/simt/stream.rs b/cuda-core/src/simt/stream.rs index 8a97417cc4..677d6722d0 100644 --- a/cuda-core/src/simt/stream.rs +++ b/cuda-core/src/simt/stream.rs @@ -42,6 +42,9 @@ pub struct CudaStream { pub(crate) cu_stream: cuda_bindings::CUstream, /// Owning context. Kept alive for the lifetime of this stream. pub(crate) ctx: Arc, + /// `false` for a stream wrapped with [`borrow_raw`](CudaStream::borrow_raw), + /// which is never destroyed here. + pub(crate) owned: bool, } /// # Safety @@ -56,10 +59,13 @@ unsafe impl Sync for CudaStream {} /// Destroys the underlying `CUstream` on drop and decrements the context's /// live stream count. /// -/// The default stream (null handle) is never destroyed. Errors during -/// teardown are recorded on the context rather than panicking. +/// The default stream (null handle) and borrowed streams are never destroyed. +/// Errors during teardown are recorded on the context rather than panicking. impl Drop for CudaStream { fn drop(&mut self) { + if !self.owned { + return; + } self.ctx.record_err(self.ctx.bind_to_thread()); if !self.cu_stream.is_null() { self.ctx.num_streams.fetch_sub(1, Ordering::Relaxed); @@ -70,6 +76,27 @@ impl Drop for CudaStream { } impl CudaStream { + /// Wraps an externally-owned CUDA stream without taking ownership, such as + /// the current stream of a framework that shares this primary context. + /// + /// `cu_stream` is the raw `CUstream` opaque pointer, as for + /// [`Stream::borrow_raw`](crate::Stream::borrow_raw). + /// + /// # Safety + /// + /// The caller must ensure: + /// - `cu_stream` points to a valid CUDA stream in `ctx`, or is null for + /// the default stream + /// - The stream outlives the returned `CudaStream` + /// - No concurrent destruction of the stream + pub unsafe fn borrow_raw(cu_stream: *mut c_void, ctx: &Arc) -> Arc { + Arc::new(CudaStream { + cu_stream: cu_stream as cuda_bindings::CUstream, + ctx: ctx.clone(), + owned: false, + }) + } + /// Returns the raw `CUstream` handle (null for the default stream). pub fn cu_stream(&self) -> cuda_bindings::CUstream { self.cu_stream @@ -155,6 +182,7 @@ impl CudaStream { let stream = Arc::new(CudaStream { cu_stream, ctx: self.ctx.clone(), + owned: true, }); stream.join(self)?; Ok(stream) diff --git a/cuda-core/tests/borrow_raw.rs b/cuda-core/tests/borrow_raw.rs index b2713d4b3e..c102cdb055 100644 --- a/cuda-core/tests/borrow_raw.rs +++ b/cuda-core/tests/borrow_raw.rs @@ -3,13 +3,13 @@ * SPDX-License-Identifier: Apache-2.0 */ -//! Verifies the ownership contract of `Device::borrow_raw` and -//! `Stream::borrow_raw`: dropping a borrowed wrapper must not release the -//! primary context or destroy the stream, so the source handles keep -//! working afterward. +//! Verifies the ownership contract of `Device::borrow_raw`, +//! `Stream::borrow_raw` and `CudaStream::borrow_raw`: dropping a borrowed +//! wrapper must not release the primary context or destroy the stream, so +//! the source handles keep working afterward. use core::ffi::{c_int, c_void}; -use cuda_core::{Device, Function, Module, Stream}; +use cuda_core::{CudaContext, CudaStream, Device, Function, Module, Stream}; const NOOP_PTX: &str = "\ .version 7.0 @@ -128,3 +128,21 @@ fn many_borrowed_drops_do_not_invalidate_source() { unsafe { source_stream.synchronize() }.unwrap(); let _another = source_dev.new_stream().unwrap(); } + +#[test] +fn borrowed_cuda_stream_drop_leaves_source_stream_usable() { + if !has_gpu() { + return; + } + let ctx = CudaContext::new(0).unwrap(); + let source_stream = ctx.new_stream().unwrap(); + + for _ in 0..8 { + let borrowed = + unsafe { CudaStream::borrow_raw(source_stream.cu_stream() as *mut c_void, &ctx) }; + borrowed.synchronize().unwrap(); + } + + source_stream.synchronize().unwrap(); + let _another = ctx.new_stream().unwrap(); +}