From 08dc2b80034ef725a6d2df0e8f1b0b0dcd1172cd Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Fri, 23 Oct 2026 12:19:11 +0100 Subject: [PATCH] feat(cuda): track device memory accesses across streams Add typed device allocations and views, host/device copies, pointer guards, pooled access events, and kernel launches. Order conflicting accesses and frees across streams while avoiding redundant waits. Include round-trip, cross-stream read, free-ordering, and event-pool tests. Existing consumers still use cudarc. Original-commit: 5fd3f54eba Signed-off-by: Alexander Droste --- Cargo.lock | 3 + vortex-cuda/driver/Cargo.toml | 9 +- vortex-cuda/driver/src/context.rs | 45 +- vortex-cuda/driver/src/launch.rs | 200 ++++++ vortex-cuda/driver/src/lib.rs | 27 +- vortex-cuda/driver/src/memory.rs | 1099 +++++++++++++++++++++++++++++ vortex-cuda/driver/src/repr.rs | 42 ++ vortex-cuda/driver/src/stream.rs | 44 ++ vortex-cuda/driver/src/tests.rs | 95 +++ 9 files changed, 1559 insertions(+), 5 deletions(-) create mode 100644 vortex-cuda/driver/src/launch.rs create mode 100644 vortex-cuda/driver/src/memory.rs create mode 100644 vortex-cuda/driver/src/repr.rs create mode 100644 vortex-cuda/driver/src/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 1232da7d7ef..e388dde1f20 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11767,7 +11767,10 @@ dependencies = [ name = "vortex-cuda-driver" version = "0.1.0" dependencies = [ + "half", "libloading 0.8.9", + "parking_lot", + "smallvec", ] [[package]] diff --git a/vortex-cuda/driver/Cargo.toml b/vortex-cuda/driver/Cargo.toml index 1ae7ad19d96..7cb879b2d87 100644 --- a/vortex-cuda/driver/Cargo.toml +++ b/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 } diff --git a/vortex-cuda/driver/src/context.rs b/vortex-cuda/driver/src/context.rs index 1e13b8c899b..ad997616fed 100644 --- a/vortex-cuda/driver/src/context.rs +++ b/vortex-cuda/driver/src/context.rs @@ -1,7 +1,7 @@ // 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; @@ -9,6 +9,8 @@ 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>, } -// 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 { + 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) { diff --git a/vortex-cuda/driver/src/launch.rs b/vortex-cuda/driver/src/launch.rs new file mode 100644 index 00000000000..ca8f5f4fdd7 --- /dev/null +++ b/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>, + timing: Option, +} + +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 { + /// 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> for LaunchArgs<'a> { + #[inline] + fn arg(&mut self, arg: &'b CudaSlice) -> &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> for LaunchArgs<'a> { + #[inline] + fn arg(&mut self, arg: &'b mut CudaSlice) -> &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, 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?)) + } +} diff --git a/vortex-cuda/driver/src/lib.rs b/vortex-cuda/driver/src/lib.rs index 76026e47234..c22e126b040 100644 --- a/vortex-cuda/driver/src/lib.rs +++ b/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; diff --git a/vortex-cuda/driver/src/memory.rs b/vortex-cuda/driver/src/memory.rs new file mode 100644 index 00000000000..0da06c177e9 --- /dev/null +++ b/vortex-cuda/driver/src/memory.rs @@ -0,0 +1,1099 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Device allocations, their views, and stream ordering between their accesses. +//! +//! Every allocation tracks which streams accessed it, so work on one stream never uses memory +//! before another stream's earlier write to it completes, and memory is never freed or +//! rewritten before other streams' earlier accesses complete. Unlike tracking every access +//! with events, only accesses that cross streams cost driver calls: +//! +//! - A write records one event right after it is enqueued, taken from the context's event +//! pool, because later readers on other streams must wait for exactly that write rather than +//! for whatever their writer enqueued afterwards. +//! - A read waits for the last write only if that write was on another stream, and only once +//! per stream and write. +//! - A read on another stream than the allocation's records a per-stream event, reused for +//! that stream's later reads, so later writes and the free can wait for it. +//! - Reads on the allocation stream, and the allocation itself, record nothing. If a later +//! access on another stream must follow them, it records the allocation stream at that time +//! and waits for that, which also covers them because they were enqueued earlier. +//! - The free runs on the allocation stream after waiting only for other streams' accesses. +//! +//! Callers must use a pointer only on the stream they passed when obtaining it, and keep the +//! returned [`SyncOnDrop`] alive until the work using it is enqueued. + +#![expect( + clippy::same_name_method, + reason = "inherent accessors mirror the DeviceSlice ones so callers need not import the trait" +)] +#![expect( + clippy::expect_used, + reason = "out-of-range views and overflowing sizes are caller bugs, as with std slices" +)] + +use std::marker::PhantomData; +use std::ops::Bound; +use std::ops::RangeBounds; +use std::sync::Arc; + +use parking_lot::Mutex; +use smallvec::SmallVec; + +use crate::context::CudaContext; +use crate::repr::DeviceRepr; +use crate::repr::ValidAsZeroBits; +use crate::result::DriverError; +use crate::stream::CudaEvent; +use crate::stream::CudaStream; +use crate::stream::StreamId; +use crate::sys; + +/// An event recorded on a stream after an access, owned by the access state. +#[derive(Clone, Copy)] +struct Recorded { + stream: StreamId, + event: sys::CUevent, +} + +#[derive(Default)] +struct AccessState { + /// The last write and the event recorded right after it. + write: Option, + /// Streams that already waited for `write`, besides its own. + waited: SmallVec<[StreamId; 4]>, + /// The latest read on each stream other than the allocation stream. + reads: SmallVec<[Recorded; 2]>, + /// Whether the allocation stream holds the allocation or reads newer than `snapshot`. + allocation_stream_pending: bool, + /// The allocation stream, recorded when another stream had to follow its pending work. + snapshot: Option, +} + +// SAFETY: Event handles are plain driver handles, usable from any thread. +unsafe impl Send for AccessState {} + +/// Stream ordering of one allocation's accesses, shared by its views. +pub(crate) struct Access { + state: Mutex, +} + +/// Whether an access reads or writes. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum AccessKind { + Read, + Write, +} + +impl Access { + fn new() -> Self { + Self { + state: Mutex::new(AccessState { + allocation_stream_pending: true, + ..AccessState::default() + }), + } + } + + /// Makes `stream` wait for earlier accesses that an access of `kind` must follow. + pub(crate) fn before( + &self, + owner: &CudaStream, + stream: &CudaStream, + kind: AccessKind, + ) -> Result<(), DriverError> { + let context = stream.context(); + stream.bind()?; + let mut state = self.state.lock(); + let id = stream.id(); + if let Some(write) = state.write + && write.stream != id + && !state.waited.contains(&id) + { + stream.wait_raw(write.event)?; + state.waited.push(id); + } + // Unrecorded allocation-stream work: the allocation, which every access follows, and + // reads, which writes follow. A recorded write already follows the allocation. + let follow_owner = id != owner.id() + && match kind { + AccessKind::Read => state.write.is_none(), + AccessKind::Write => true, + }; + if follow_owner { + if state.allocation_stream_pending { + // Record everything enqueued there so far, which includes that work. + let event = match state.snapshot { + Some(event) => event, + None => { + let event = context.take_event()?; + state.snapshot = Some(event); + event + } + }; + owner.record_raw(event)?; + state.allocation_stream_pending = false; + } + if let Some(snapshot) = state.snapshot { + stream.wait_raw(snapshot)?; + } + } + if kind == AccessKind::Write { + for read in &state.reads { + if read.stream != id { + stream.wait_raw(read.event)?; + } + } + } + Ok(()) + } + + /// Records an access of `kind` on `stream`, after it was enqueued. + pub(crate) fn after( + &self, + owner: &CudaStream, + stream: &CudaStream, + kind: AccessKind, + ) -> Result<(), DriverError> { + let context = stream.context(); + stream.bind()?; + let mut state = self.state.lock(); + let id = stream.id(); + match kind { + AccessKind::Write => { + // Readers and older writes are ordered before this write by `before`, so only + // this write remains to be followed. + let event = match state.write.take() { + Some(previous) => previous.event, + None => context.take_event()?, + }; + stream.record_raw(event)?; + state.write = Some(Recorded { stream: id, event }); + state.waited.clear(); + for read in state.reads.drain(..) { + context.give_event(read.event); + } + if id == owner.id() { + state.allocation_stream_pending = false; + } + } + AccessKind::Read if id == owner.id() => { + state.allocation_stream_pending = true; + } + AccessKind::Read => { + let event = match state.reads.iter().position(|read| read.stream == id) { + Some(index) => state.reads[index].event, + None => { + let event = context.take_event()?; + state.reads.push(Recorded { stream: id, event }); + event + } + }; + stream.record_raw(event)?; + } + } + Ok(()) + } + + /// Makes the allocation stream wait for every other stream's access before a free. + fn before_free(&self, owner: &CudaStream) -> Result<(), DriverError> { + owner.bind()?; + let state = self.state.lock(); + let id = owner.id(); + if let Some(write) = state.write + && write.stream != id + && !state.waited.contains(&id) + { + owner.wait_raw(write.event)?; + } + for read in &state.reads { + owner.wait_raw(read.event)?; + } + Ok(()) + } + + /// Returns the access events to the context's pool. + fn release(&mut self, context: &CudaContext) { + let state = self.state.get_mut(); + if let Some(write) = state.write.take() { + context.give_event(write.event); + } + for read in state.reads.drain(..) { + context.give_event(read.event); + } + if let Some(snapshot) = state.snapshot.take() { + context.give_event(snapshot); + } + } +} + +/// Records an access when dropped, after the work using its pointer was enqueued. +#[must_use] +pub struct SyncOnDrop<'a> { + pending: Pending<'a>, +} + +enum Pending<'a> { + None, + Access(&'a Access, &'a CudaStream, &'a CudaStream, AccessKind), + Event(&'a CudaEvent, &'a CudaStream), +} + +impl<'a> SyncOnDrop<'a> { + pub(crate) fn new( + access: &'a Access, + owner: &'a CudaStream, + stream: &'a CudaStream, + kind: AccessKind, + ) -> Self { + Self { + pending: Pending::Access(access, owner, stream, kind), + } + } + + /// A guard that records nothing. + pub fn none() -> Self { + Self { + pending: Pending::None, + } + } + + /// A guard that records `stream`'s work into `event` when dropped. + pub fn record_event(event: &'a CudaEvent, stream: &'a CudaStream) -> Self { + Self { + pending: Pending::Event(event, stream), + } + } +} + +impl Drop for SyncOnDrop<'_> { + fn drop(&mut self) { + match std::mem::replace(&mut self.pending, Pending::None) { + Pending::None => {} + Pending::Access(access, owner, stream, kind) => { + stream + .context() + .record_err(access.after(owner, stream, kind)); + } + Pending::Event(event, stream) => { + stream.context().record_err(event.record(stream)); + } + } + } +} + +impl std::fmt::Debug for SyncOnDrop<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let pending = match &self.pending { + Pending::None => "none", + Pending::Access(..) => "access", + Pending::Event(..) => "event", + }; + f.debug_struct("SyncOnDrop") + .field("pending", &pending) + .finish() + } +} + +/// A typed device allocation, freed in stream order on its allocation stream. +pub struct CudaSlice { + pointer: sys::CUdeviceptr, + len: usize, + access: Access, + stream: Arc, + marker: PhantomData<*const T>, +} + +// SAFETY: The allocation is device memory reached only through the driver, and the access +// state is synchronized. +unsafe impl Send for CudaSlice {} +// SAFETY: As above; shared references only read the handle and lock the access state. +unsafe impl Sync for CudaSlice {} + +impl std::fmt::Debug for CudaSlice { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaSlice") + .field("pointer", &self.pointer) + .field("len", &self.len) + .field("stream", &self.stream.id()) + .finish() + } +} + +impl Drop for CudaSlice { + fn drop(&mut self) { + let context = self.stream.context(); + context.record_err(self.access.before_free(&self.stream)); + context.record_err(self.stream.free_raw(self.pointer)); + self.access.release(context); + } +} + +impl CudaSlice { + /// Wraps an allocation made on `stream`. + /// + /// # Safety + /// + /// `pointer` must be a stream-ordered allocation of at least `len` elements made on + /// `stream`, owned by the returned slice. + pub(crate) unsafe fn from_raw( + pointer: sys::CUdeviceptr, + len: usize, + stream: Arc, + ) -> Self { + Self { + pointer, + len, + access: Access::new(), + stream, + marker: PhantomData, + } + } + + /// The number of elements. + pub fn len(&self) -> usize { + self.len + } + + /// The number of bytes. + pub fn num_bytes(&self) -> usize { + self.len * size_of::() + } + + /// Whether the slice holds no elements. + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// The device ordinal. + pub fn ordinal(&self) -> usize { + self.stream.context().ordinal() + } + + /// The context the memory belongs to. + pub fn context(&self) -> &Arc { + self.stream.context() + } + + /// The stream that allocated the memory and will free it. + pub fn stream(&self) -> &Arc { + &self.stream + } + + /// The whole slice as a view. + pub fn as_view(&self) -> CudaView<'_, T> { + CudaView { + pointer: self.pointer, + len: self.len, + access: &self.access, + stream: &self.stream, + marker: PhantomData, + } + } + + /// The whole slice as a mutable view. + pub fn as_view_mut(&mut self) -> CudaViewMut<'_, T> { + CudaViewMut { + pointer: self.pointer, + len: self.len, + access: &self.access, + stream: &self.stream, + marker: PhantomData, + } + } + + /// A view of the elements in `bounds`. + /// + /// # Panics + /// + /// Panics if `bounds` is out of range. + pub fn slice(&self, bounds: impl RangeBounds) -> CudaView<'_, T> { + self.as_view().slice(bounds) + } + + /// A view of the elements in `bounds`, if in range. + pub fn try_slice(&self, bounds: impl RangeBounds) -> Option> { + self.as_view().try_slice(bounds) + } + + /// A mutable view of the elements in `bounds`. + /// + /// # Panics + /// + /// Panics if `bounds` is out of range. + pub fn slice_mut(&mut self, bounds: impl RangeBounds) -> CudaViewMut<'_, T> { + self.as_view_mut().into_slice_mut(bounds) + } + + /// Reinterprets the memory as `len` elements of `S`, if they fit. + /// + /// # Safety + /// + /// The bytes must be valid values of `S`. + pub unsafe fn transmute(&self, len: usize) -> Option> { + (len * size_of::() <= self.num_bytes()).then_some(CudaView { + pointer: self.pointer, + len, + access: &self.access, + stream: &self.stream, + marker: PhantomData, + }) + } + + /// Reinterprets the memory as `len` mutable elements of `S`, if they fit. + /// + /// # Safety + /// + /// The bytes must be valid values of `S`, and `T` must accept any bytes written as `S`. + pub unsafe fn transmute_mut(&mut self, len: usize) -> Option> { + (len * size_of::() <= self.num_bytes()).then_some(CudaViewMut { + pointer: self.pointer, + len, + access: &self.access, + stream: &self.stream, + marker: PhantomData, + }) + } +} + +impl CudaSlice { + /// A device copy of the slice on its stream. + pub fn try_clone(&self) -> Result { + self.stream.clone_dtod(self) + } +} + +/// Resolves `bounds` against `len`, if in range. +fn range(bounds: &impl RangeBounds, len: usize) -> Option<(usize, usize)> { + let start = match bounds.start_bound() { + Bound::Included(&start) => start, + Bound::Excluded(&start) => start.checked_add(1)?, + Bound::Unbounded => 0, + }; + let end = match bounds.end_bound() { + Bound::Included(&end) => end.checked_add(1)?, + Bound::Excluded(&end) => end, + Bound::Unbounded => len, + }; + (start <= end && end <= len).then_some((start, end)) +} + +/// A shared view of part of a [`CudaSlice`]. +pub struct CudaView<'a, T> { + pointer: sys::CUdeviceptr, + len: usize, + access: &'a Access, + stream: &'a Arc, + marker: PhantomData<&'a [T]>, +} + +impl Clone for CudaView<'_, T> { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for CudaView<'_, T> {} + +impl std::fmt::Debug for CudaView<'_, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaView") + .field("pointer", &self.pointer) + .field("len", &self.len) + .finish() + } +} + +impl<'a, T> CudaView<'a, T> { + /// The number of elements. + pub fn len(&self) -> usize { + self.len + } + + /// Whether the view holds no elements. + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// The number of bytes. + pub fn num_bytes(&self) -> usize { + self.len * size_of::() + } + + /// A view of the elements in `bounds`. + /// + /// # Panics + /// + /// Panics if `bounds` is out of range. + pub fn slice(&self, bounds: impl RangeBounds) -> CudaView<'a, T> { + self.try_slice(bounds) + .expect("CUDA view range out of bounds") + } + + /// A view of the elements in `bounds`, if in range. + pub fn try_slice(&self, bounds: impl RangeBounds) -> Option> { + let (start, end) = range(&bounds, self.len)?; + Some(CudaView { + pointer: self.pointer + (start * size_of::()) as u64, + len: end - start, + ..*self + }) + } + + /// Reinterprets the view as `len` elements of `S`, if they fit. + /// + /// # Safety + /// + /// The bytes must be valid values of `S`. + pub unsafe fn transmute(&self, len: usize) -> Option> { + (len * size_of::() <= self.num_bytes()).then_some(CudaView { + pointer: self.pointer, + len, + access: self.access, + stream: self.stream, + marker: PhantomData, + }) + } + + /// Like [`DevicePtr::device_ptr`], with the guard living as long as the view. + pub fn view_ptr(self, stream: &'a CudaStream) -> (sys::CUdeviceptr, SyncOnDrop<'a>) { + let guard = access(self.access, self.stream, stream, AccessKind::Read); + (self.pointer, guard) + } +} + +/// A mutable view of part of a [`CudaSlice`]. +pub struct CudaViewMut<'a, T> { + pointer: sys::CUdeviceptr, + len: usize, + access: &'a Access, + stream: &'a Arc, + marker: PhantomData<&'a mut [T]>, +} + +impl std::fmt::Debug for CudaViewMut<'_, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaViewMut") + .field("pointer", &self.pointer) + .field("len", &self.len) + .finish() + } +} + +impl<'a, T> CudaViewMut<'a, T> { + /// The number of elements. + pub fn len(&self) -> usize { + self.len + } + + /// Whether the view holds no elements. + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// The number of bytes. + pub fn num_bytes(&self) -> usize { + self.len * size_of::() + } + + /// The view as a shared view. + pub fn as_view(&self) -> CudaView<'_, T> { + CudaView { + pointer: self.pointer, + len: self.len, + access: self.access, + stream: self.stream, + marker: PhantomData, + } + } + + /// Converts into a shared view with the original lifetime. + pub fn into_view(self) -> CudaView<'a, T> { + CudaView { + pointer: self.pointer, + len: self.len, + access: self.access, + stream: self.stream, + marker: PhantomData, + } + } + + /// A shared view of the elements in `bounds`. + /// + /// # Panics + /// + /// Panics if `bounds` is out of range. + pub fn slice(&self, bounds: impl RangeBounds) -> CudaView<'_, T> { + self.as_view().slice(bounds) + } + + /// A mutable view of the elements in `bounds`. + /// + /// # Panics + /// + /// Panics if `bounds` is out of range. + pub fn slice_mut(&mut self, bounds: impl RangeBounds) -> CudaViewMut<'_, T> { + let (start, end) = range(&bounds, self.len).expect("CUDA view range out of bounds"); + CudaViewMut { + pointer: self.pointer + (start * size_of::()) as u64, + len: end - start, + access: self.access, + stream: self.stream, + marker: PhantomData, + } + } + + fn into_slice_mut(self, bounds: impl RangeBounds) -> CudaViewMut<'a, T> { + let (start, end) = range(&bounds, self.len).expect("CUDA view range out of bounds"); + CudaViewMut { + pointer: self.pointer + (start * size_of::()) as u64, + len: end - start, + ..self + } + } + + /// Reinterprets the view as `len` mutable elements of `S`, if they fit. + /// + /// # Safety + /// + /// The bytes must be valid values of `S`, and `T` must accept any bytes written as `S`. + pub unsafe fn transmute_mut(&mut self, len: usize) -> Option> { + (len * size_of::() <= self.num_bytes()).then_some(CudaViewMut { + pointer: self.pointer, + len, + access: self.access, + stream: self.stream, + marker: PhantomData, + }) + } + + /// Like [`DevicePtrMut::device_ptr_mut`], with the guard living as long as the view. + pub fn view_ptr_mut(self, stream: &'a CudaStream) -> (sys::CUdeviceptr, SyncOnDrop<'a>) { + let guard = access(self.access, self.stream, stream, AccessKind::Write); + (self.pointer, guard) + } +} + +/// Orders `stream` after earlier conflicting accesses and returns the guard recording this one. +fn access<'a>( + access: &'a Access, + owner: &'a CudaStream, + stream: &'a CudaStream, + kind: AccessKind, +) -> SyncOnDrop<'a> { + stream + .context() + .record_err(access.before(owner, stream, kind)); + SyncOnDrop::new(access, owner, stream, kind) +} + +/// Device memory with a length and an allocation stream. +pub trait DeviceSlice { + /// The number of elements. + fn len(&self) -> usize; + + /// The number of bytes. + fn num_bytes(&self) -> usize { + self.len() * size_of::() + } + + /// Whether there are no elements. + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// The stream that allocated the memory. + fn stream(&self) -> &Arc; +} + +/// Device memory that can be read on a stream. +pub trait DevicePtr: DeviceSlice { + /// The device pointer for reading on `stream`, which first waits for earlier writes. + /// + /// Drop the guard after the read is enqueued. + fn device_ptr<'a>(&'a self, stream: &'a CudaStream) -> (sys::CUdeviceptr, SyncOnDrop<'a>); +} + +/// Device memory that can be written on a stream. +pub trait DevicePtrMut: DeviceSlice { + /// The device pointer for writing on `stream`, which first waits for earlier accesses. + /// + /// Drop the guard after the write is enqueued. + fn device_ptr_mut<'a>( + &'a mut self, + stream: &'a CudaStream, + ) -> (sys::CUdeviceptr, SyncOnDrop<'a>); +} + +impl DeviceSlice for CudaSlice { + fn len(&self) -> usize { + self.len + } + + fn stream(&self) -> &Arc { + &self.stream + } +} + +impl DeviceSlice for CudaView<'_, T> { + fn len(&self) -> usize { + self.len + } + + fn stream(&self) -> &Arc { + self.stream + } +} + +impl DeviceSlice for CudaViewMut<'_, T> { + fn len(&self) -> usize { + self.len + } + + fn stream(&self) -> &Arc { + self.stream + } +} + +impl DevicePtr for CudaSlice { + fn device_ptr<'a>(&'a self, stream: &'a CudaStream) -> (sys::CUdeviceptr, SyncOnDrop<'a>) { + let guard = access(&self.access, &self.stream, stream, AccessKind::Read); + (self.pointer, guard) + } +} + +impl DevicePtr for CudaView<'_, T> { + fn device_ptr<'a>(&'a self, stream: &'a CudaStream) -> (sys::CUdeviceptr, SyncOnDrop<'a>) { + let guard = access(self.access, self.stream, stream, AccessKind::Read); + (self.pointer, guard) + } +} + +impl DevicePtr for CudaViewMut<'_, T> { + fn device_ptr<'a>(&'a self, stream: &'a CudaStream) -> (sys::CUdeviceptr, SyncOnDrop<'a>) { + let guard = access(self.access, self.stream, stream, AccessKind::Read); + (self.pointer, guard) + } +} + +impl DevicePtrMut for CudaSlice { + fn device_ptr_mut<'a>( + &'a mut self, + stream: &'a CudaStream, + ) -> (sys::CUdeviceptr, SyncOnDrop<'a>) { + let guard = access(&self.access, &self.stream, stream, AccessKind::Write); + (self.pointer, guard) + } +} + +impl DevicePtrMut for CudaViewMut<'_, T> { + fn device_ptr_mut<'a>( + &'a mut self, + stream: &'a CudaStream, + ) -> (sys::CUdeviceptr, SyncOnDrop<'a>) { + let guard = access(self.access, self.stream, stream, AccessKind::Write); + (self.pointer, guard) + } +} + +/// The raw parts a kernel launch needs to order and record an argument's access. +pub(crate) struct ArgAccess<'a> { + pub(crate) access: &'a Access, + pub(crate) owner: &'a CudaStream, + pub(crate) kind: AccessKind, +} + +impl CudaSlice { + pub(crate) fn arg_access(&self, kind: AccessKind) -> (ArgAccess<'_>, &sys::CUdeviceptr) { + ( + ArgAccess { + access: &self.access, + owner: &self.stream, + kind, + }, + &self.pointer, + ) + } +} + +impl<'a, T> CudaView<'a, T> { + pub(crate) fn arg_access(&self) -> (ArgAccess<'a>, &sys::CUdeviceptr) { + ( + ArgAccess { + access: self.access, + owner: self.stream, + kind: AccessKind::Read, + }, + &self.pointer, + ) + } +} + +impl<'a, T> CudaViewMut<'a, T> { + pub(crate) fn arg_access(&self) -> (ArgAccess<'a>, &sys::CUdeviceptr) { + ( + ArgAccess { + access: self.access, + owner: self.stream, + kind: AccessKind::Write, + }, + &self.pointer, + ) + } +} + +/// Host memory that copies can read from or write into. +pub trait HostSlice { + /// The number of elements. + fn len(&self) -> usize; + + /// Whether there are no elements. + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// The elements for a copy on `stream`, with a guard to drop after the copy is enqueued. + /// + /// # Safety + /// + /// The slice may be used only by work enqueued on `stream` before the guard drops. + unsafe fn stream_synced_slice<'a>( + &'a self, + stream: &'a CudaStream, + ) -> (&'a [T], SyncOnDrop<'a>); + + /// The elements for a copy into them on `stream`, with a guard as above. + /// + /// # Safety + /// + /// As for [`stream_synced_slice`](Self::stream_synced_slice). + unsafe fn stream_synced_mut_slice<'a>( + &'a mut self, + stream: &'a CudaStream, + ) -> (&'a mut [T], SyncOnDrop<'a>); +} + +impl HostSlice for [T] { + fn len(&self) -> usize { + <[T]>::len(self) + } + + unsafe fn stream_synced_slice<'a>( + &'a self, + _stream: &'a CudaStream, + ) -> (&'a [T], SyncOnDrop<'a>) { + (self, SyncOnDrop::none()) + } + + unsafe fn stream_synced_mut_slice<'a>( + &'a mut self, + _stream: &'a CudaStream, + ) -> (&'a mut [T], SyncOnDrop<'a>) { + (self, SyncOnDrop::none()) + } +} + +impl HostSlice for [T; N] { + fn len(&self) -> usize { + N + } + + unsafe fn stream_synced_slice<'a>( + &'a self, + _stream: &'a CudaStream, + ) -> (&'a [T], SyncOnDrop<'a>) { + (self, SyncOnDrop::none()) + } + + unsafe fn stream_synced_mut_slice<'a>( + &'a mut self, + _stream: &'a CudaStream, + ) -> (&'a mut [T], SyncOnDrop<'a>) { + (self, SyncOnDrop::none()) + } +} + +impl HostSlice for Vec { + fn len(&self) -> usize { + Vec::len(self) + } + + unsafe fn stream_synced_slice<'a>( + &'a self, + _stream: &'a CudaStream, + ) -> (&'a [T], SyncOnDrop<'a>) { + (self, SyncOnDrop::none()) + } + + unsafe fn stream_synced_mut_slice<'a>( + &'a mut self, + _stream: &'a CudaStream, + ) -> (&'a mut [T], SyncOnDrop<'a>) { + (self, SyncOnDrop::none()) + } +} + +impl CudaStream { + /// Allocates `len` uninitialized elements on this stream. + /// + /// # Safety + /// + /// The memory is uninitialized. + pub unsafe fn alloc( + self: &Arc, + len: usize, + ) -> Result, DriverError> { + let bytes = len + .checked_mul(size_of::()) + .expect("CUDA allocation size overflow"); + let pointer = self.alloc_raw(bytes)?; + // SAFETY: A fresh stream-ordered allocation of `len` elements on this stream. + Ok(unsafe { CudaSlice::from_raw(pointer, len, Arc::clone(self)) }) + } + + /// Allocates `len` zeroed elements on this stream. + pub fn alloc_zeros( + self: &Arc, + len: usize, + ) -> Result, DriverError> { + // SAFETY: Zeroed below; zero bits are valid values of `T`. + let mut slice = unsafe { self.alloc(len) }?; + self.memset_zeros(&mut slice)?; + Ok(slice) + } + + /// Zeroes `target` on this stream. + pub fn memset_zeros>( + self: &Arc, + target: &mut Target, + ) -> Result<(), DriverError> { + let bytes = target.num_bytes(); + let (pointer, _record) = target.device_ptr_mut(self); + self.memset_d8_raw(pointer, 0, bytes) + } + + /// Copies `source` into a new allocation on this stream. + pub fn clone_htod + ?Sized>( + self: &Arc, + source: &Source, + ) -> Result, DriverError> { + // SAFETY: Fully written by the copy below. + let mut target = unsafe { self.alloc(source.len()) }?; + self.memcpy_htod(source, &mut target)?; + Ok(target) + } + + /// Copies `source` into `target` on this stream. + /// + /// Pageable sources are staged by the driver before this returns; page-locked sources must + /// stay alive and unmodified until the copy completes. + /// + /// # Panics + /// + /// Panics if `target` is shorter than `source`. + pub fn memcpy_htod + ?Sized, Target: DevicePtrMut>( + self: &Arc, + source: &Source, + target: &mut Target, + ) -> Result<(), DriverError> { + assert!(target.len() >= source.len()); + // SAFETY: The slice is used only by the copy below, enqueued on this stream. + let (source, _record_source) = unsafe { source.stream_synced_slice(self) }; + let (pointer, _record) = target.device_ptr_mut(self); + self.bind()?; + // SAFETY: The target holds at least `size_of_val(source)` bytes, ordered for writing. + unsafe { + sys::cuMemcpyHtoDAsync_v2( + pointer, + source.as_ptr().cast(), + size_of_val(source), + self.cu_stream(), + ) + } + .result() + } + + /// Copies `source` into a new host vector, returning once the copy completed. + pub fn clone_dtoh>( + self: &Arc, + source: &Source, + ) -> Result, DriverError> { + let mut target: Vec = Vec::with_capacity(source.len()); + let (pointer, _record) = source.device_ptr(self); + self.bind()?; + // SAFETY: The vector has capacity for `source.len()` elements. A copy into pageable + // memory returns only after it completes, which the synchronization below also + // guarantees for page-locked memory. + unsafe { + sys::cuMemcpyDtoHAsync_v2( + target.as_mut_ptr().cast(), + pointer, + source.num_bytes(), + self.cu_stream(), + ) + .result()?; + sys::cuStreamSynchronize(self.cu_stream()).result()?; + target.set_len(source.len()); + } + Ok(target) + } + + /// Copies `source` into `target`, returning once the copy completed. + /// + /// # Panics + /// + /// Panics if `target` is shorter than `source`. + pub fn memcpy_dtoh, Target: HostSlice + ?Sized>( + self: &Arc, + source: &Source, + target: &mut Target, + ) -> Result<(), DriverError> { + assert!(target.len() >= source.len()); + // SAFETY: The slice is used only by the copy below, enqueued on this stream. + let (target, _record_target) = unsafe { target.stream_synced_mut_slice(self) }; + let (pointer, _record) = source.device_ptr(self); + self.bind()?; + // SAFETY: The target holds at least `source.num_bytes()` bytes; the copy completes + // before this returns. + unsafe { + sys::cuMemcpyDtoHAsync_v2( + target.as_mut_ptr().cast(), + pointer, + source.num_bytes(), + self.cu_stream(), + ) + .result()?; + sys::cuStreamSynchronize(self.cu_stream()).result() + } + } + + /// Copies `source` into `target` on this stream. + /// + /// # Panics + /// + /// Panics if `target` is shorter than `source`. + pub fn memcpy_dtod, Target: DevicePtrMut>( + self: &Arc, + source: &Source, + target: &mut Target, + ) -> Result<(), DriverError> { + assert!(target.len() >= source.len()); + let bytes = source.num_bytes(); + let (source, _record_source) = source.device_ptr(self); + let (target, _record_target) = target.device_ptr_mut(self); + self.bind()?; + // SAFETY: Both ranges hold `bytes` bytes, ordered for this stream. + unsafe { sys::cuMemcpyDtoDAsync_v2(target, source, bytes, self.cu_stream()) }.result() + } + + /// Copies `source` into a new allocation on this stream. + pub fn clone_dtod>( + self: &Arc, + source: &Source, + ) -> Result, DriverError> { + // SAFETY: Fully written by the copy below. + let mut target = unsafe { self.alloc(source.len()) }?; + self.memcpy_dtod(source, &mut target)?; + Ok(target) + } +} diff --git a/vortex-cuda/driver/src/repr.rs b/vortex-cuda/driver/src/repr.rs new file mode 100644 index 00000000000..39280b2f7b9 --- /dev/null +++ b/vortex-cuda/driver/src/repr.rs @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Marker traits for values stored in device memory or passed to kernels. + +/// A type with the same representation on the device, usable in device memory and as a kernel +/// argument. +/// +/// # Safety +/// +/// The type must be plain data with a layout the device code agrees on, usually `#[repr(C)]`. +pub unsafe trait DeviceRepr {} + +/// A type whose all-zero bit pattern is a valid value. +/// +/// # Safety +/// +/// Zeroed memory must be a valid value of the type. +pub unsafe trait ValidAsZeroBits {} + +macro_rules! impl_plain { + ($($t:ty),*) => { + $( + // SAFETY: Primitive plain data. + unsafe impl DeviceRepr for $t {} + // SAFETY: Zero is a valid primitive value. + unsafe impl ValidAsZeroBits for $t {} + )* + }; +} + +impl_plain!( + bool, i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize, f32, f64 +); + +#[cfg(feature = "f16")] +impl_plain!(half::f16, half::bf16); + +// SAFETY: Arrays of device values are laid out contiguously. +unsafe impl DeviceRepr for [T; N] {} +// SAFETY: Arrays of values valid as zeros are valid as zeros. +unsafe impl ValidAsZeroBits for [T; N] {} diff --git a/vortex-cuda/driver/src/stream.rs b/vortex-cuda/driver/src/stream.rs index b2748773a5b..89bdc498223 100644 --- a/vortex-cuda/driver/src/stream.rs +++ b/vortex-cuda/driver/src/stream.rs @@ -100,6 +100,10 @@ impl CudaStream { }) } + pub(crate) fn id(&self) -> StreamId { + self.id + } + /// The raw stream handle; do not destroy it. pub fn cu_stream(&self) -> sys::CUstream { self.stream @@ -165,6 +169,46 @@ impl CudaStream { // SAFETY: Both handles belong to this stream's context and are valid. unsafe { sys::cuEventRecord(event, self.stream) }.result() } + + pub(crate) fn alloc_raw(&self, bytes: usize) -> Result { + self.bind()?; + let mut pointer = 0; + // SAFETY: The context is current; the driver writes the allocation. + unsafe { + if self.context.has_async_alloc() { + sys::cuMemAllocAsync(&raw mut pointer, bytes, self.stream) + } else { + sys::cuMemAlloc_v2(&raw mut pointer, bytes) + } + } + .result()?; + Ok(pointer) + } + + pub(crate) fn free_raw(&self, pointer: sys::CUdeviceptr) -> Result<(), DriverError> { + self.bind()?; + // SAFETY: Callers free an allocation of this context once, after its accesses were + // ordered before this stream's queued work. + unsafe { + if self.context.has_async_alloc() { + sys::cuMemFreeAsync(pointer, self.stream).result() + } else { + sys::cuStreamSynchronize(self.stream).result()?; + sys::cuMemFree_v2(pointer).result() + } + } + } + + pub(crate) fn memset_d8_raw( + &self, + pointer: sys::CUdeviceptr, + value: u8, + bytes: usize, + ) -> Result<(), DriverError> { + self.bind()?; + // SAFETY: Callers pass a range of `bytes` writable device bytes ordered for this stream. + unsafe { sys::cuMemsetD8Async(pointer, value, bytes, self.stream) }.result() + } } /// A marker in a stream's work that other streams and the host can wait for. diff --git a/vortex-cuda/driver/src/tests.rs b/vortex-cuda/driver/src/tests.rs new file mode 100644 index 00000000000..890768fe6e0 --- /dev/null +++ b/vortex-cuda/driver/src/tests.rs @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use crate::CudaContext; +use crate::CudaStream; +use crate::DevicePtr; +use crate::DevicePtrMut; +use crate::DriverError; +use crate::sys; + +/// A context and two streams, or `None` without a CUDA device. +fn streams() -> Option<(Arc, Arc, Arc)> { + // SAFETY: Only probes for the driver library. + if !unsafe { sys::is_culib_present() } { + return None; + } + let context = CudaContext::new(0).ok()?; + let first = context.new_stream().ok()?; + let second = context.new_stream().ok()?; + Some((context, first, second)) +} + +#[test] +fn host_device_round_trip() -> Result<(), DriverError> { + let Some((context, stream, _)) = streams() else { + return Ok(()); + }; + let values: Vec = (0..100_000).collect(); + let device = stream.clone_htod(&values)?; + assert_eq!(stream.clone_dtoh(&device)?, values); + assert_eq!(stream.clone_dtoh(&device.slice(10..20))?, &values[10..20]); + context.check_err() +} + +/// A write on one stream is visible to reads on another, through views as well. +#[test] +fn cross_stream_reads_follow_writes() -> Result<(), DriverError> { + let Some((context, writer, reader)) = streams() else { + return Ok(()); + }; + for round in 0..50u8 { + let mut device = writer.alloc_zeros::(64 << 20)?; + { + let bytes = device.num_bytes(); + let (pointer, _record) = device.device_ptr_mut(&writer); + // SAFETY: The allocation holds `bytes` bytes. + unsafe { sys::cuMemsetD8Async(pointer, round, bytes, writer.cu_stream()) }.result()?; + } + let tail = reader.clone_dtoh(&device.slice((64 << 20) - 16..))?; + assert!(tail.iter().all(|value| *value == round)); + } + context.check_err() +} + +/// A free after a read on another stream waits for that read, and the memory a later +/// allocation reuses is not overwritten under it. +#[test] +fn frees_follow_reads_on_other_streams() -> Result<(), DriverError> { + let Some((context, owner, reader)) = streams() else { + return Ok(()); + }; + let source = owner.clone_htod(&vec![7u8; 32 << 20])?; + let mut copies = Vec::new(); + for _ in 0..8 { + // SAFETY: Fully written by the copy below. + let mut target = unsafe { reader.alloc::(32 << 20) }?; + reader.memcpy_dtod(&source, &mut target)?; + copies.push(target); + } + drop(source); + let reused = owner.clone_htod(&vec![9u8; 32 << 20])?; + for copy in &copies { + assert!(reader.clone_dtoh(copy)?.iter().all(|value| *value == 7)); + } + assert!(owner.clone_dtoh(&reused)?.iter().all(|value| *value == 9)); + context.check_err() +} + +/// Same-stream accesses take no events; cross-stream ones reuse returned events. +#[test] +fn events_are_pooled() -> Result<(), DriverError> { + let Some((context, first, second)) = streams() else { + return Ok(()); + }; + for _ in 0..100 { + let device = first.clone_htod(&[1u32; 1024])?; + let (_pointer, record) = device.device_ptr(&second); + drop(record); + } + // One write event and one read event, reused for every buffer. + assert!(context.pooled_events() <= 2); + context.check_err() +}