feat(cuda): add driver context and resource ownership

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<Arc<Self>, 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<Arc<Self>, 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<sys::CUdevice, DriverError> {
+        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<Arc<Self>, 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<i32, DriverError> {
+        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<String, DriverError> {
+        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<i32, DriverError> {
+        // 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<Self>) -> Result<Arc<CudaStream>, 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<Self>,
+        priority: i32,
+    ) -> Result<Arc<CudaStream>, DriverError> {
+        CudaStream::create(self, Some(priority))
+    }
+
+    /// The legacy default stream.
+    pub fn default_stream(self: &Arc<Self>) -> Arc<CudaStream> {
+        // 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<Self>,
+        flags: Option<sys::CUevent_flags>,
+    ) -> Result<CudaEvent, DriverError> {
+        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<Self>, image: &[u8]) -> Result<Arc<CudaModule>, 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<T>(&self, result: Result<T, DriverError>) {
+        if let Err(error) = result {
+            self.error_state.store(error.0.0, Ordering::Relaxed);
+        }
+    }
+}

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;

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<CudaContext>,
+}
+
+// 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<CudaContext>, image: &[u8]) -> Result<Arc<Self>, 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<Self>, name: &str) -> Result<CudaFunction, DriverError> {
+        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<CudaModule>,
+}
+
+// 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
+    }
+}