From b896d5ddfd3925f7c2002b2cdaed861c9d9ded4b Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Fri, 23 Oct 2026 12:19:10 +0100 Subject: [PATCH] feat(cuda): add driver context and resource ownership Add owned and borrowed CUDA resources, explicit thread-context binding, stream and event synchronization, module loading, and deferred driver errors. Stage allocation helpers and event pooling with the memory-tracking change that consumes them. Original-commit: 5fd3f54eba Signed-off-by: Alexander Droste --- vortex-cuda/driver/src/context.rs | 249 ++++++++++++++++++++++++++++ vortex-cuda/driver/src/lib.rs | 12 +- vortex-cuda/driver/src/module.rs | 95 +++++++++++ vortex-cuda/driver/src/stream.rs | 261 ++++++++++++++++++++++++++++++ 4 files changed, 615 insertions(+), 2 deletions(-) create mode 100644 vortex-cuda/driver/src/context.rs create mode 100644 vortex-cuda/driver/src/module.rs create mode 100644 vortex-cuda/driver/src/stream.rs diff --git a/vortex-cuda/driver/src/context.rs b/vortex-cuda/driver/src/context.rs new file mode 100644 index 00000000000..1e13b8c899b --- /dev/null +++ b/vortex-cuda/driver/src/context.rs @@ -0,0 +1,249 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The device's primary context and deferred errors. + +use std::ffi::CStr; +use std::ffi::c_char; +use std::sync::Arc; +use std::sync::atomic::AtomicU32; +use std::sync::atomic::Ordering; + +use crate::module::CudaModule; +use crate::result; +use crate::result::DriverError; +use crate::stream::CudaEvent; +use crate::stream::CudaStream; +use crate::sys; +use crate::sys::CUresult; + +/// The primary context of one device. +/// +/// Every operation binds the context to the calling thread first, so a context can be used +/// from any thread. +pub struct CudaContext { + device: sys::CUdevice, + context: sys::CUcontext, + ordinal: usize, + has_async_alloc: bool, + /// Whether the context is the device's primary context, released rather than destroyed. + primary: bool, + error_state: AtomicU32, +} + +// SAFETY: The handles are driver handles usable from any thread. +unsafe impl Send for CudaContext {} +// SAFETY: As above. +unsafe impl Sync for CudaContext {} + +impl std::fmt::Debug for CudaContext { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaContext") + .field("ordinal", &self.ordinal) + .finish() + } +} + +impl PartialEq for CudaContext { + fn eq(&self, other: &Self) -> bool { + self.device == other.device && self.context == other.context + } +} + +impl Eq for CudaContext {} + +impl Drop for CudaContext { + fn drop(&mut self) { + self.record_err(self.bind_to_thread()); + // SAFETY: Releases the retain made by `new`, or destroys the context created by + // `new_non_primary`, once. + self.record_err( + unsafe { + if self.primary { + sys::cuDevicePrimaryCtxRelease_v2(self.device) + } else { + sys::cuCtxDestroy_v2(self.context) + } + } + .result(), + ); + } +} + +impl CudaContext { + /// Retains the primary context of the device with `ordinal`. + pub fn new(ordinal: usize) -> Result, DriverError> { + let device = Self::device(ordinal)?; + let mut context = std::ptr::null_mut(); + // SAFETY: `device` is valid; the driver writes the retained context. + unsafe { sys::cuDevicePrimaryCtxRetain(&raw mut context, device) }.result()?; + Self::wrap(device, context, ordinal, true) + } + + /// Creates a context separate from the device's primary context, with `cuCtxCreate` flags. + pub fn new_non_primary(ordinal: usize, flags: u32) -> Result, DriverError> { + let device = Self::device(ordinal)?; + let mut context = std::ptr::null_mut(); + // SAFETY: `device` is valid; the driver writes the created context. + unsafe { sys::cuCtxCreate_v2(&raw mut context, flags, device) }.result()?; + Self::wrap(device, context, ordinal, false) + } + + fn device(ordinal: usize) -> Result { + result::init()?; + result::device::get( + i32::try_from(ordinal).map_err(|_| DriverError(CUresult::CUDA_ERROR_INVALID_VALUE))?, + ) + } + + fn wrap( + device: sys::CUdevice, + context: sys::CUcontext, + ordinal: usize, + primary: bool, + ) -> Result, DriverError> { + // SAFETY: `device` is valid. + let pools = unsafe { + result::device::get_attribute( + device, + sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED, + ) + }?; + let context = Arc::new(Self { + device, + context, + ordinal, + has_async_alloc: pools > 0, + primary, + error_state: AtomicU32::new(0), + }); + context.bind_to_thread()?; + Ok(context) + } + + /// The number of devices. + pub fn device_count() -> Result { + result::init()?; + result::device::get_count() + } + + /// The device ordinal. + pub fn ordinal(&self) -> usize { + self.ordinal + } + + /// The raw device handle. + pub fn cu_device(&self) -> sys::CUdevice { + self.device + } + + /// The raw context handle; do not release it. + pub fn cu_ctx(&self) -> sys::CUcontext { + self.context + } + + /// Whether allocations are stream-ordered through the device's memory pool. + pub fn has_async_alloc(&self) -> bool { + self.has_async_alloc + } + + /// The device's name. + pub fn name(&self) -> Result { + let mut name = [0 as c_char; 256]; + // SAFETY: The buffer holds 256 bytes, which the driver NUL-terminates. + unsafe { sys::cuDeviceGetName(name.as_mut_ptr(), 256, self.device) }.result()?; + // SAFETY: NUL-terminated by the driver. + Ok(unsafe { CStr::from_ptr(name.as_ptr()) } + .to_string_lossy() + .into_owned()) + } + + /// A device attribute. + pub fn attribute(&self, attribute: sys::CUdevice_attribute) -> Result { + // SAFETY: The device handle is valid for the context's lifetime. + unsafe { result::device::get_attribute(self.device, attribute) } + } + + /// The device's free and total memory in bytes. + pub fn mem_get_info(&self) -> Result<(usize, usize), DriverError> { + self.bind_to_thread()?; + let (mut free, mut total) = (0, 0); + // SAFETY: The driver writes both values. + unsafe { sys::cuMemGetInfo_v2(&raw mut free, &raw mut total) }.result()?; + Ok((free, total)) + } + + /// Makes this context current on the calling thread, if it is not already. + /// + /// First returns, and clears, any error recorded where none could be returned. + pub fn bind_to_thread(&self) -> Result<(), DriverError> { + if self.error_state.load(Ordering::Relaxed) != 0 { + self.check_err()?; + } + let mut current = std::ptr::null_mut(); + // SAFETY: The driver writes the current context. + unsafe { sys::cuCtxGetCurrent(&raw mut current) }.result()?; + if current != self.context { + // SAFETY: The context is retained for this object's lifetime. + unsafe { sys::cuCtxSetCurrent(self.context) }.result()?; + } + Ok(()) + } + + /// Blocks until all work in the context completes. + pub fn synchronize(&self) -> Result<(), DriverError> { + self.bind_to_thread()?; + // SAFETY: The context is current. + unsafe { sys::cuCtxSynchronize() }.result() + } + + /// A new non-blocking stream. + pub fn new_stream(self: &Arc) -> Result, DriverError> { + CudaStream::create(self, None) + } + + /// A new non-blocking stream with `priority`; lower values are higher priorities. + pub fn new_stream_with_priority( + self: &Arc, + priority: i32, + ) -> Result, DriverError> { + CudaStream::create(self, Some(priority)) + } + + /// The legacy default stream. + pub fn default_stream(self: &Arc) -> Arc { + // SAFETY: The null stream is valid for the context's lifetime. + unsafe { CudaStream::borrowed(self, std::ptr::null_mut()) } + } + + /// A new event with `flags`, by default timing-disabled. + pub fn new_event( + self: &Arc, + flags: Option, + ) -> Result { + CudaEvent::create( + self, + flags.unwrap_or(sys::CUevent_flags::CU_EVENT_DISABLE_TIMING), + ) + } + + /// Loads a module from a cubin, fatbin or PTX image. + pub fn load_module(self: &Arc, image: &[u8]) -> Result, DriverError> { + CudaModule::load(self, image) + } + + /// The first error recorded where none could be returned, clearing it. + pub fn check_err(&self) -> Result<(), DriverError> { + match self.error_state.swap(0, Ordering::Relaxed) { + 0 => Ok(()), + error => Err(DriverError(CUresult(error))), + } + } + + /// Keeps an error from a call that cannot return it, such as in `Drop`. + pub fn record_err(&self, result: Result) { + if let Err(error) = result { + self.error_state.store(error.0.0, Ordering::Relaxed); + } + } +} diff --git a/vortex-cuda/driver/src/lib.rs b/vortex-cuda/driver/src/lib.rs index 64cbe2ce72f..76026e47234 100644 --- a/vortex-cuda/driver/src/lib.rs +++ b/vortex-cuda/driver/src/lib.rs @@ -3,10 +3,18 @@ //! Minimal CUDA driver API bindings for Vortex. //! -//! The driver library is loaded lazily, so dependents build and run without CUDA until -//! they first use it. The raw bindings and result wrappers operate on CUDA driver handles. +//! Owns CUDA contexts, streams, events, and loaded modules. The driver library is loaded +//! lazily, so dependents build and run without CUDA until they first use it. +mod context; +mod module; pub mod result; +mod stream; pub mod sys; +pub use context::CudaContext; +pub use module::CudaFunction; +pub use module::CudaModule; pub use result::DriverError; +pub use stream::CudaEvent; +pub use stream::CudaStream; diff --git a/vortex-cuda/driver/src/module.rs b/vortex-cuda/driver/src/module.rs new file mode 100644 index 00000000000..94d920c1d08 --- /dev/null +++ b/vortex-cuda/driver/src/module.rs @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Loaded device code. + +use std::ffi::CString; +use std::sync::Arc; + +use crate::context::CudaContext; +use crate::result::DriverError; +use crate::sys; +use crate::sys::CUresult; + +/// A loaded cubin, fatbin or PTX module, unloaded when the last function is dropped. +pub struct CudaModule { + module: sys::CUmodule, + context: Arc, +} + +// SAFETY: Module handles are usable from any thread. +unsafe impl Send for CudaModule {} +// SAFETY: As above. +unsafe impl Sync for CudaModule {} + +impl std::fmt::Debug for CudaModule { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaModule") + .field("module", &self.module) + .finish() + } +} + +impl Drop for CudaModule { + fn drop(&mut self) { + self.context.record_err(self.context.bind_to_thread()); + // SAFETY: Unloaded once, after every function referencing it was dropped. + self.context + .record_err(unsafe { sys::cuModuleUnload(self.module) }.result()); + } +} + +impl CudaModule { + pub(crate) fn load(context: &Arc, image: &[u8]) -> Result, DriverError> { + context.bind_to_thread()?; + let mut module = std::ptr::null_mut(); + // SAFETY: The driver reads a complete image; fatbins and cubins carry their own size. + unsafe { sys::cuModuleLoadData(&raw mut module, image.as_ptr().cast()) }.result()?; + Ok(Arc::new(Self { + module, + context: Arc::clone(context), + })) + } + + /// The kernel named `name`. + pub fn load_function(self: &Arc, name: &str) -> Result { + let name = + CString::new(name).map_err(|_| DriverError(CUresult::CUDA_ERROR_INVALID_VALUE))?; + self.context.bind_to_thread()?; + let mut function = std::ptr::null_mut(); + // SAFETY: The module is loaded; the driver writes the function handle. + unsafe { sys::cuModuleGetFunction(&raw mut function, self.module, name.as_ptr()) } + .result()?; + Ok(CudaFunction { + function, + _module: Arc::clone(self), + }) + } +} + +/// A kernel in a loaded module, which it keeps loaded. +#[derive(Clone)] +pub struct CudaFunction { + function: sys::CUfunction, + _module: Arc, +} + +// SAFETY: Function handles are usable from any thread. +unsafe impl Send for CudaFunction {} +// SAFETY: As above. +unsafe impl Sync for CudaFunction {} + +impl std::fmt::Debug for CudaFunction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaFunction") + .field("function", &self.function) + .finish() + } +} + +impl CudaFunction { + /// The raw function handle. + pub fn cu_function(&self) -> sys::CUfunction { + self.function + } +} diff --git a/vortex-cuda/driver/src/stream.rs b/vortex-cuda/driver/src/stream.rs new file mode 100644 index 00000000000..b2748773a5b --- /dev/null +++ b/vortex-cuda/driver/src/stream.rs @@ -0,0 +1,261 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Streams and events. + +use std::sync::Arc; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use crate::context::CudaContext; +use crate::result::DriverError; +use crate::sys; + +/// A process-unique stream identity: stream handles can be reused after a stream is destroyed, +/// so access tracking never compares raw handles. +pub(crate) type StreamId = u64; + +/// The next [`StreamId`]. +static NEXT_STREAM_ID: AtomicU64 = AtomicU64::new(1); + +/// A queue of device work, executed in order. +pub struct CudaStream { + id: StreamId, + stream: sys::CUstream, + context: Arc, + owned: bool, +} + +// SAFETY: Stream handles are usable from any thread. +unsafe impl Send for CudaStream {} +// SAFETY: As above. +unsafe impl Sync for CudaStream {} + +impl std::fmt::Debug for CudaStream { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaStream") + .field("id", &self.id) + .field("stream", &self.stream) + .finish() + } +} + +impl PartialEq for CudaStream { + fn eq(&self, other: &Self) -> bool { + self.id == other.id + } +} + +impl Eq for CudaStream {} + +impl Drop for CudaStream { + fn drop(&mut self) { + if self.owned { + self.context.record_err(self.context.bind_to_thread()); + // SAFETY: Owned streams are destroyed exactly once; queued work still completes. + self.context + .record_err(unsafe { sys::cuStreamDestroy_v2(self.stream) }.result()); + } + } +} + +impl CudaStream { + pub(crate) fn create( + context: &Arc, + priority: Option, + ) -> Result, DriverError> { + context.bind_to_thread()?; + let mut stream = std::ptr::null_mut(); + // SAFETY: The context is current; the driver writes the new stream. + unsafe { + match priority { + Some(priority) => sys::cuStreamCreateWithPriority( + &raw mut stream, + sys::CU_STREAM_NON_BLOCKING, + priority, + ), + None => sys::cuStreamCreate(&raw mut stream, sys::CU_STREAM_NON_BLOCKING), + } + } + .result()?; + Ok(Arc::new(Self { + id: NEXT_STREAM_ID.fetch_add(1, Ordering::Relaxed), + stream, + context: Arc::clone(context), + owned: true, + })) + } + + /// Wraps a stream owned by someone else, which is not destroyed on drop. + /// + /// # Safety + /// + /// `stream` must belong to `context` and outlive the returned wrapper. + pub unsafe fn borrowed(context: &Arc, stream: sys::CUstream) -> Arc { + Arc::new(Self { + id: NEXT_STREAM_ID.fetch_add(1, Ordering::Relaxed), + stream, + context: Arc::clone(context), + owned: false, + }) + } + + /// The raw stream handle; do not destroy it. + pub fn cu_stream(&self) -> sys::CUstream { + self.stream + } + + /// The stream's context. + pub fn context(&self) -> &Arc { + &self.context + } + + pub(crate) fn bind(&self) -> Result<(), DriverError> { + self.context.bind_to_thread() + } + + /// Blocks until all work queued so far completes. + pub fn synchronize(&self) -> Result<(), DriverError> { + self.bind()?; + // SAFETY: The stream is valid for this object's lifetime. + unsafe { sys::cuStreamSynchronize(self.stream) }.result() + } + + /// A new event recording the work queued so far. + pub fn record_event( + &self, + flags: Option, + ) -> Result { + let event = self.context.new_event(flags)?; + event.record(self)?; + Ok(event) + } + + /// Makes later work wait for the work recorded in `event`. + pub fn wait(&self, event: &CudaEvent) -> Result<(), DriverError> { + self.bind()?; + self.wait_raw(event.cu_event()) + } + + /// Makes later work wait for the work queued on `other` so far. + pub fn join(&self, other: &CudaStream) -> Result<(), DriverError> { + self.wait(&other.record_event(None)?) + } + + /// A new stream whose work follows the work queued here so far. + pub fn fork(&self) -> Result, DriverError> { + let stream = Self::create(&self.context, None)?; + stream.join(self)?; + Ok(stream) + } + + pub(crate) fn wait_raw(&self, event: sys::CUevent) -> Result<(), DriverError> { + // SAFETY: Both handles belong to this stream's context and are valid. + unsafe { + sys::cuStreamWaitEvent( + self.stream, + event, + sys::CUevent_wait_flags::CU_EVENT_WAIT_DEFAULT.0, + ) + } + .result() + } + + pub(crate) fn record_raw(&self, event: sys::CUevent) -> Result<(), DriverError> { + // SAFETY: Both handles belong to this stream's context and are valid. + unsafe { sys::cuEventRecord(event, self.stream) }.result() + } +} + +/// A marker in a stream's work that other streams and the host can wait for. +pub struct CudaEvent { + event: sys::CUevent, + context: Arc, +} + +// SAFETY: Event handles are usable from any thread. +unsafe impl Send for CudaEvent {} +// SAFETY: As above. +unsafe impl Sync for CudaEvent {} + +impl std::fmt::Debug for CudaEvent { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaEvent") + .field("event", &self.event) + .finish() + } +} + +impl Drop for CudaEvent { + fn drop(&mut self) { + self.context.record_err(self.context.bind_to_thread()); + // SAFETY: Owned events are destroyed once; pending work completes regardless. + self.context + .record_err(unsafe { sys::cuEventDestroy_v2(self.event) }.result()); + } +} + +impl CudaEvent { + pub(crate) fn create( + context: &Arc, + flags: sys::CUevent_flags, + ) -> Result { + context.bind_to_thread()?; + let mut event = std::ptr::null_mut(); + // SAFETY: The context is current; the driver writes the new event. + unsafe { sys::cuEventCreate(&raw mut event, flags.0) }.result()?; + Ok(Self { + event, + context: Arc::clone(context), + }) + } + + /// The raw event handle; do not destroy it. + pub fn cu_event(&self) -> sys::CUevent { + self.event + } + + /// The event's context. + pub fn context(&self) -> &Arc { + &self.context + } + + /// Records the work queued on `stream` so far; earlier waits are unaffected. + pub fn record(&self, stream: &CudaStream) -> Result<(), DriverError> { + stream.bind()?; + stream.record_raw(self.event) + } + + /// Blocks until the recorded work completes. + pub fn synchronize(&self) -> Result<(), DriverError> { + self.context.bind_to_thread()?; + // SAFETY: The event is valid for this object's lifetime. + unsafe { sys::cuEventSynchronize(self.event) }.result() + } + + /// Whether the recorded work completed; errors count as incomplete. + pub fn is_complete(&self) -> bool { + matches!(self.try_is_complete(), Ok(true)) + } + + /// Whether the recorded work completed, without waiting. + pub fn try_is_complete(&self) -> Result { + // SAFETY: The event is valid for this object's lifetime. + match unsafe { sys::cuEventQuery(self.event) } { + sys::CUresult::CUDA_SUCCESS => Ok(true), + sys::CUresult::CUDA_ERROR_NOT_READY => Ok(false), + error => Err(DriverError(error)), + } + } + + /// Milliseconds from this event to `end`, after both complete. + pub fn elapsed_ms(&self, end: &Self) -> Result { + self.synchronize()?; + end.synchronize()?; + let mut milliseconds = 0.0; + // SAFETY: Both events completed and were created with timing. + unsafe { sys::cuEventElapsedTime(&raw mut milliseconds, self.event, end.event) } + .result()?; + Ok(milliseconds) + } +}