feat(cuda-core): add CudaStream::borrow_raw for externally owned streams

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,
         }))
     }
 

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<CudaContext>,
+    /// `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<CudaContext>) -> Arc<Self> {
+        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)

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();
+}