feat(cuda): track device memory accesses across streams

vortex-cuda/driver/Cargo.toml

@@ -4,7 +4,7 @@
 [package]
 name = "vortex-cuda-driver"
 authors.workspace = true
-description = "Minimal CUDA driver API bindings for Vortex"
+description = "Minimal CUDA driver API bindings with cross-stream-only synchronization for Vortex"
 edition = { workspace = true }
 categories = { workspace = true }
 include = { workspace = true }
@@ -19,5 +19,12 @@ publish = false
 [lints]
 workspace = true
 
+[features]
+default = []
+f16 = ["dep:half"]
+
 [dependencies]
+half = { workspace = true, optional = true }
 libloading = { workspace = true }
+parking_lot = { workspace = true }
+smallvec = { workspace = true }

vortex-cuda/driver/src/context.rs

@@ -1,14 +1,16 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-//! The device's primary context and deferred errors.
+//! The device's primary context, its event pool, 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 parking_lot::Mutex;
+
 use crate::module::CudaModule;
 use crate::result;
 use crate::result::DriverError;
@@ -29,9 +31,12 @@ pub struct CudaContext {
     /// Whether the context is the device's primary context, released rather than destroyed.
     primary: bool,
     error_state: AtomicU32,
+    /// Timing-disabled events returned by finished accesses, reused to avoid creating and
+    /// destroying an event per access.
+    events: Mutex<Vec<sys::CUevent>>,
 }
 
-// SAFETY: The handles are driver handles usable from any thread.
+// SAFETY: The handles are driver handles usable from any thread; the pool is synchronized.
 unsafe impl Send for CudaContext {}
 // SAFETY: As above.
 unsafe impl Sync for CudaContext {}
@@ -55,6 +60,10 @@ impl Eq for CudaContext {}
 impl Drop for CudaContext {
     fn drop(&mut self) {
         self.record_err(self.bind_to_thread());
+        for event in std::mem::take(self.events.get_mut()) {
+            // SAFETY: Pooled events are owned by the pool and no longer in use.
+            self.record_err(unsafe { sys::cuEventDestroy_v2(event) }.result());
+        }
         // SAFETY: Releases the retain made by `new`, or destroys the context created by
         // `new_non_primary`, once.
         self.record_err(
@@ -116,6 +125,7 @@ impl CudaContext {
             has_async_alloc: pools > 0,
             primary,
             error_state: AtomicU32::new(0),
+            events: Mutex::new(Vec::new()),
         });
         context.bind_to_thread()?;
         Ok(context)
@@ -232,6 +242,37 @@ impl CudaContext {
         CudaModule::load(self, image)
     }
 
+    /// A timing-disabled event from the pool, or a new one.
+    pub(crate) fn take_event(&self) -> Result<sys::CUevent, DriverError> {
+        if let Some(event) = self.events.lock().pop() {
+            return Ok(event);
+        }
+        self.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,
+                sys::CUevent_flags::CU_EVENT_DISABLE_TIMING.0,
+            )
+        }
+        .result()?;
+        Ok(event)
+    }
+
+    /// Returns an event to the pool once no access will wait for it again.
+    ///
+    /// Waits already enqueued on it are unaffected by later records.
+    pub(crate) fn give_event(&self, event: sys::CUevent) {
+        self.events.lock().push(event);
+    }
+
+    /// The number of events waiting in the pool.
+    #[cfg(test)]
+    pub(crate) fn pooled_events(&self) -> usize {
+        self.events.lock().len()
+    }
+
     /// 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) {

vortex-cuda/driver/src/launch.rs

@@ -0,0 +1,200 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Kernel launches with stream-ordered device memory arguments.
+
+use std::ffi::c_void;
+
+use crate::memory::AccessKind;
+use crate::memory::ArgAccess;
+use crate::memory::CudaSlice;
+use crate::memory::CudaView;
+use crate::memory::CudaViewMut;
+use crate::module::CudaFunction;
+use crate::repr::DeviceRepr;
+use crate::result::DriverError;
+use crate::stream::CudaEvent;
+use crate::stream::CudaStream;
+use crate::sys;
+
+/// Grid, block and dynamic shared memory of a launch.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct LaunchConfig {
+    /// Blocks per grid dimension.
+    pub grid_dim: (u32, u32, u32),
+    /// Threads per block dimension.
+    pub block_dim: (u32, u32, u32),
+    /// Dynamic shared memory per block, in bytes.
+    pub shared_mem_bytes: u32,
+}
+
+impl LaunchConfig {
+    /// One-dimensional blocks of 1024 threads covering `n` elements.
+    pub fn for_num_elems(n: u32) -> Self {
+        Self {
+            grid_dim: (n.div_ceil(1024), 1, 1),
+            block_dim: (1024, 1, 1),
+            shared_mem_bytes: 0,
+        }
+    }
+}
+
+/// A kernel launch being assembled on a stream; create with [`CudaStream::launch_builder`].
+///
+/// Device memory arguments are ordered after earlier conflicting accesses and recorded as
+/// accesses of this launch.
+pub struct LaunchArgs<'a> {
+    stream: &'a CudaStream,
+    function: &'a CudaFunction,
+    // `Vec`, unlike `SmallVec`, lets argument borrows end before the builder drops.
+    args: Vec<*mut c_void>,
+    accesses: Vec<ArgAccess<'a>>,
+    timing: Option<sys::CUevent_flags>,
+}
+
+impl std::fmt::Debug for LaunchArgs<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LaunchArgs")
+            .field("args", &self.args.len())
+            .finish()
+    }
+}
+
+impl CudaStream {
+    /// Starts a launch of `function` on this stream.
+    pub fn launch_builder<'a>(&'a self, function: &'a CudaFunction) -> LaunchArgs<'a> {
+        LaunchArgs {
+            stream: self,
+            function,
+            args: Vec::with_capacity(8),
+            accesses: Vec::with_capacity(4),
+            timing: None,
+        }
+    }
+}
+
+/// A value that can be passed as a kernel argument.
+///
+/// # Safety
+///
+/// Implementations must push a pointer to storage that lives as long as the [`LaunchArgs`] and
+/// holds the argument in the kernel's expected representation.
+pub unsafe trait PushKernelArg<T> {
+    /// Appends the argument.
+    fn arg(&mut self, arg: T) -> &mut Self;
+}
+
+// SAFETY: The reference outlives the launch and `T` has a device representation.
+unsafe impl<'a, 'b: 'a, T: DeviceRepr> PushKernelArg<&'b T> for LaunchArgs<'a> {
+    #[inline]
+    fn arg(&mut self, arg: &'b T) -> &mut Self {
+        self.args.push((arg as *const T).cast_mut().cast());
+        self
+    }
+}
+
+// SAFETY: The device pointer lives in the slice, which outlives the launch.
+unsafe impl<'a, 'b: 'a, T> PushKernelArg<&'b CudaSlice<T>> for LaunchArgs<'a> {
+    #[inline]
+    fn arg(&mut self, arg: &'b CudaSlice<T>) -> &mut Self {
+        let (access, pointer) = arg.arg_access(AccessKind::Read);
+        self.accesses.push(access);
+        self.args
+            .push((pointer as *const sys::CUdeviceptr).cast_mut().cast());
+        self
+    }
+}
+
+// SAFETY: The device pointer lives in the slice, which outlives the launch.
+unsafe impl<'a, 'b: 'a, T> PushKernelArg<&'b mut CudaSlice<T>> for LaunchArgs<'a> {
+    #[inline]
+    fn arg(&mut self, arg: &'b mut CudaSlice<T>) -> &mut Self {
+        let (access, pointer) = arg.arg_access(AccessKind::Write);
+        self.accesses.push(access);
+        self.args
+            .push((pointer as *const sys::CUdeviceptr).cast_mut().cast());
+        self
+    }
+}
+
+// SAFETY: The device pointer lives in the view, which outlives the launch.
+unsafe impl<'a, 'b: 'a, 'c: 'b, T> PushKernelArg<&'b CudaView<'c, T>> for LaunchArgs<'a> {
+    #[inline]
+    fn arg(&mut self, arg: &'b CudaView<'c, T>) -> &mut Self {
+        let (access, pointer) = arg.arg_access();
+        self.accesses.push(access);
+        self.args
+            .push((pointer as *const sys::CUdeviceptr).cast_mut().cast());
+        self
+    }
+}
+
+// SAFETY: The device pointer lives in the view, which outlives the launch.
+unsafe impl<'a, 'b: 'a, 'c: 'b, T> PushKernelArg<&'b mut CudaViewMut<'c, T>> for LaunchArgs<'a> {
+    #[inline]
+    fn arg(&mut self, arg: &'b mut CudaViewMut<'c, T>) -> &mut Self {
+        let (access, pointer) = arg.arg_access();
+        self.accesses.push(access);
+        self.args
+            .push((pointer as *const sys::CUdeviceptr).cast_mut().cast());
+        self
+    }
+}
+
+impl LaunchArgs<'_> {
+    /// Makes [`launch`](Self::launch) return events recorded before and after the kernel.
+    pub fn record_kernel_launch(&mut self, flags: sys::CUevent_flags) -> &mut Self {
+        self.timing = Some(flags);
+        self
+    }
+
+    /// Enqueues the kernel after the arguments' earlier conflicting accesses.
+    ///
+    /// # Safety
+    ///
+    /// The arguments must match the kernel's parameters, and the kernel must access device
+    /// memory only within the arguments, writing only those passed mutably.
+    pub unsafe fn launch(
+        &mut self,
+        config: LaunchConfig,
+    ) -> Result<Option<(CudaEvent, CudaEvent)>, DriverError> {
+        self.stream.bind()?;
+        for access in &self.accesses {
+            access
+                .access
+                .before(access.owner, self.stream, access.kind)?;
+        }
+        let start = self
+            .timing
+            .map(|flags| self.stream.record_event(Some(flags)))
+            .transpose()?;
+        // SAFETY: Under the caller's guarantees; the argument pointers outlive the call.
+        unsafe {
+            sys::cuLaunchKernel(
+                self.function.cu_function(),
+                config.grid_dim.0,
+                config.grid_dim.1,
+                config.grid_dim.2,
+                config.block_dim.0,
+                config.block_dim.1,
+                config.block_dim.2,
+                config.shared_mem_bytes,
+                self.stream.cu_stream(),
+                self.args.as_mut_ptr(),
+                std::ptr::null_mut(),
+            )
+        }
+        .result()?;
+        // The kernel is already queued: a trailing-event failure must not skip access tracking.
+        let end = self
+            .timing
+            .map(|flags| self.stream.record_event(Some(flags)))
+            .transpose();
+        for access in &self.accesses {
+            access
+                .access
+                .after(access.owner, self.stream, access.kind)?;
+        }
+        Ok(start.zip(end?))
+    }
+}

vortex-cuda/driver/src/lib.rs

@@ -3,18 +3,41 @@
 
 //! Minimal CUDA driver API bindings for Vortex.
 //!
-//! 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.
+//! Covers the driver surface Vortex uses: the primary context, streams, events, modules,
+//! kernel launches, and stream-ordered device memory. The driver library is loaded lazily, so
+//! dependents build and run without CUDA until they first use it.
+//!
+//! Device memory orders its accesses across streams automatically, paying for driver calls
+//! only when an access crosses streams; see [`CudaSlice`].
 
 mod context;
+mod launch;
+mod memory;
 mod module;
+mod repr;
 pub mod result;
 mod stream;
 pub mod sys;
 
 pub use context::CudaContext;
+pub use launch::LaunchArgs;
+pub use launch::LaunchConfig;
+pub use launch::PushKernelArg;
+pub use memory::CudaSlice;
+pub use memory::CudaView;
+pub use memory::CudaViewMut;
+pub use memory::DevicePtr;
+pub use memory::DevicePtrMut;
+pub use memory::DeviceSlice;
+pub use memory::HostSlice;
+pub use memory::SyncOnDrop;
 pub use module::CudaFunction;
 pub use module::CudaModule;
+pub use repr::DeviceRepr;
+pub use repr::ValidAsZeroBits;
 pub use result::DriverError;
 pub use stream::CudaEvent;
 pub use stream::CudaStream;
+
+#[cfg(test)]
+mod tests;