[drbh]

hey guys! congrats on consolidating the repos!

this pr adds CudaStream::borrow_raw, which wraps a raw CUstream without taking ownership, so a host can launch cuda-oxide kernels on a stream that something else owns (eg. the current torch stream)

the generated launchers on a #[cuda_module] take &Arc<CudaStream>, and today the only way to get one is CudaContext::new_stream / default_stream / CudaStream::fork. dropping any of those destroys the stream, so there is no safe way to hand a framework's stream to a LoadedModule. we hit this building cuda-oxide kernels as python extensions for https://github.com/huggingface/kernels, where the stream comes from TVMFFIEnvGetStream (torch via tvm-ffi)

this follows the pattern runtime.rs already uses for external handles. Device, Stream, Module and Function each have an unsafe fn borrow_raw taking the raw *mut c_void, backed by an owned: bool field that Drop checks before releasing anything (added in 6f3b354e7)

cuda-rust/cuda-core/src/runtime.rs

Lines 726 to 733 in e027ac3

	pub unsafe fn borrow_raw(cu_stream: *mut c_void, device: &Arc<Device>) -> Arc<Self> {
	Arc::new(Stream {
	cu_stream: cu_stream as cuda_bindings::CUstream,
	device: device.clone(),
	owned: false,
	_keep_alive: KeepAlive::none(),
	})
	}

cuda-rust/cuda-core/src/runtime.rs

Lines 684 to 688 in e027ac3

	impl Drop for Stream {
	fn drop(&mut self) {
	if !self.owned || self.cu_stream.is_null() {
	return;
	}

CudaStream::borrow_raw is the same shape on the simt side: same *mut c_void argument, same safety contract as Stream::borrow_raw, and a crate private owned field. the existing constructors set owned: true so nothing changes for them. a borrowed stream is not counted in num_streams and its Drop returns early. DeviceBuffer::from_raw_parts is the existing simt precedent for wrapping memory we do not own, this adds the stream equivalent

note: this does not add a borrow_with_owner variant like runtime::Stream has. happy to add it here if you want parity, we did not need it since the caller only holds the stream for the duration of one launch

cargo test -p cuda-core --test borrow_raw
cargo clippy -p cuda-core --all-targets -- -D warnings
RUSTDOCFLAGS=-Dwarnings cargo doc -p cuda-core --no-deps

running 6 tests
test borrowed_module_drop_does_not_unload ... ok
test borrowed_drop_leaves_source_stream_usable ... ok
test borrowed_drop_leaves_source_device_usable ... ok
test borrowed_function_uses_parent_module ... ok
test borrowed_cuda_stream_drop_leaves_source_stream_usable ... ok
test many_borrowed_drops_do_not_invalidate_source ... ok

test result: ok. 6 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.12s

a new test in cuda-core/tests/borrow_raw.rs borrows and drops a CudaStream 8 times and checks the source stream and context are still usable. also checked end to end on an rtx 4090 with a kernel loaded from python and launched on a torch side stream

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.