From f74daee3a4f13997d8b7818beb4029a1b95376c6 Mon Sep 17 00:00:00 2001 From: Melih Elibol Date: Wed, 2 Sep 2026 13:30:47 -0700 Subject: [PATCH] fix(simt): soundness and correctness fixes for the cuda-oxide bridge and VMM Audit follow-up for the `simt` compatibility surface in cuda-core and cuda-async and for both VMM modules. Kept separate from the host-runtime patch because it changes cuda-oxide-facing APIs and is not part of 0.3.1. Soundness - Synchronize wrappers reported the recorded (sticky) context error before the driver wait, so DeviceBuffer copies returned Err while the DMA still wrote the caller's host memory (reproduced). CudaStream, CudaContext and CudaEvent::synchronize now wait first and report the sticky error after; DeviceBuffer binds and drains the sticky state before enqueueing, and leaks host memory rather than free it when the driver wait itself fails. - CudaStream::launch_host_function requires F: 'static (a closure borrowing a local compiled and ran after its frame was gone), reclaims the box when the enqueue fails, and documents the callback rules. - cuda_async::simt::Empty and Value derive Send from their fields; the unconditional `unsafe impl Send` laundered !Send closures. - The execute-once unzip halves (SelectLeft/SelectRight) replace their unsynchronized UnsafeCell check-then-act with a mutex/condvar ExecuteOnce; a consumer on another stream waits on an event the producer records. - ConstantHandle carries its symbol size and an Arc; oversize writes and foreign-context streams are rejected; it is no longer Copy. - DeviceCopy for raw pointers excludes unsized pointees (a wide pointer read back from device bytes carried an invalid vtable). - Mapping<'a> borrows its VirtualReservation and its PhysicalAllocation or MulticastObject, and MulticastBinding<'a> borrows its MulticastObject, in both cuda-core/src/vmm.rs and simt/vmm.rs; unmap is documented as not stream-ordered. Correctness - DeviceBuffer constructors bind the stream's context; they failed on an unbound thread and could allocate in a foreign context. - Drop paths release resources regardless of a sticky error; async buffers leaked while an error was recorded. - The thread-local device map survives failed lookups and re-entrant sync; SingleStream reports use-before-init instead of panicking; a colliding kernel-cache insert keeps the existing entry. - simt::vmm::align_size handles any granularity with checked arithmetic; streams are counted only after cuStreamCreate succeeds; PinnedHostBuffer checks the host pointer's alignment for T. Docs: stale "opaque CUlaunchAttribute_st" comments, the simt collision list, embedded-module docs, the CoordinateRequirement re-export, and the cuda-core README's description of the two host models. Breaking for cuda-oxide: ConstantHandle::from_raw(module, dptr, size) and write_blocking(src, len); Mapping::new(&reservation, va_offset, &phys, phys_offset, size) and new_multicast(...) with the new lifetimes; launch_host_function F: 'static; Empty F: Send. Tests: new GPU suites simt_module_globals and simt_execute_once, and new cases in simt_device_buffer, simt_device_buffer_leaks, simt_device_copy_impls, vmm and simt_vmm_p2p; the sticky-error and leak tests fail without the fixes. The simt_device_buffer file is serialized to avoid a pre-existing driver-level deadlock between a gated host callback and concurrent cuMemFree/cuStreamDestroy. Signed-off-by: Melih Elibol --- cuda-async/src/simt/device_context.rs | 165 ++++++++-- cuda-async/src/simt/device_operation.rs | 346 ++++++++++++++------ cuda-async/src/simt/scheduling_policies.rs | 43 ++- cuda-async/tests/simt_execute_once.rs | 178 ++++++++++ cuda-core/README.md | 56 +++- cuda-core/src/lib.rs | 10 +- cuda-core/src/simt/context.rs | 135 ++++++-- cuda-core/src/simt/device_buffer.rs | 140 ++++++-- cuda-core/src/simt/embedded.rs | 46 +++ cuda-core/src/simt/event.rs | 18 +- cuda-core/src/simt/mod.rs | 36 +- cuda-core/src/simt/module.rs | 177 ++++++++-- cuda-core/src/simt/pinned_host_buffer.rs | 21 +- cuda-core/src/simt/stream.rs | 90 ++++- cuda-core/src/simt/vmm.rs | 282 +++++++++++++--- cuda-core/src/vmm.rs | 218 +++++++++--- cuda-core/tests/simt_device_buffer.rs | 280 +++++++++++++++- cuda-core/tests/simt_device_buffer_leaks.rs | 113 ++++++- cuda-core/tests/simt_device_copy_impls.rs | 12 + cuda-core/tests/simt_module_globals.rs | 171 ++++++++++ cuda-core/tests/simt_vmm_multicast.rs | 52 +-- cuda-core/tests/simt_vmm_p2p.rs | 50 ++- cuda-core/tests/vmm.rs | 64 +++- cuda-core/tests/vmm_multicast.rs | 52 +-- cutile-rs/scripts/run_gpu_tests.sh | 4 +- 25 files changed, 2340 insertions(+), 419 deletions(-) create mode 100644 cuda-async/tests/simt_execute_once.rs create mode 100644 cuda-core/tests/simt_module_globals.rs diff --git a/cuda-async/src/simt/device_context.rs b/cuda-async/src/simt/device_context.rs index 29bf86b615..432772084e 100644 --- a/cuda-async/src/simt/device_context.rs +++ b/cuda-async/src/simt/device_context.rs @@ -26,6 +26,7 @@ use crate::simt::scheduling_policies::{ use cuda_core::{CudaContext, CudaFunction, CudaModule, CudaStream}; use rustc_hash::FxHashMap; use std::cell::Cell; +use std::collections::hash_map::Entry; use std::hash::{DefaultHasher, Hash, Hasher}; use std::sync::Arc; @@ -176,32 +177,51 @@ fn init_with_default_policy( ) } -/// Borrows the thread-local [`AsyncDeviceContext`] for `device_id` immutably. +/// Checks the thread's device map out of the thread-local, runs `f` on it, +/// and puts it back **whether or not `f` succeeded**. /// -/// Lazily initializes the context map and the specific device if needed. -fn with_global_device_context(device_id: usize, f: F) -> Result -where - F: FnOnce(&AsyncDeviceContext) -> R, -{ +/// The map lives in a `Cell>`, so it is moved out for the duration +/// of the call. An early `?` between the take and the put-back used to drop +/// the map on the floor, destroying every device context this thread had +/// initialized (contexts, stream pools, kernel cache) because one lookup of +/// an unrelated device failed; the next access then rebuilt everything from +/// scratch. +fn with_device_map( + device_id: usize, + f: impl FnOnce(&mut FxHashMap) -> Result, +) -> Result { DEVICE_CONTEXTS.with(|ctx| { let mut hashmap = match ctx.devices.take() { Some(hashmap) => hashmap, None => { + // Nothing is checked out yet, so a failure here loses nothing. init_device_contexts_default()?; ctx.devices .take() .ok_or_else(|| device_error(device_id, "Failed to initialize context"))? } }; + let result = f(&mut hashmap); + ctx.devices.replace(Some(hashmap)); + result + }) +} + +/// Borrows the thread-local [`AsyncDeviceContext`] for `device_id` immutably. +/// +/// Lazily initializes the context map and the specific device if needed. +fn with_global_device_context(device_id: usize, f: F) -> Result +where + F: FnOnce(&AsyncDeviceContext) -> R, +{ + with_device_map(device_id, |hashmap| { if !hashmap.contains_key(&device_id) { - init_with_default_policy(&mut hashmap, device_id)?; + init_with_default_policy(hashmap, device_id)?; } let device_context = hashmap .get(&device_id) .ok_or_else(|| device_error(device_id, "Failed to get context"))?; - let r = f(device_context); - ctx.devices.replace(Some(hashmap)); - Ok(r) + Ok(f(device_context)) }) } @@ -212,25 +232,14 @@ fn with_global_device_context_mut(device_id: usize, f: F) -> Result R, { - DEVICE_CONTEXTS.with(|ctx| { - let mut hashmap = match ctx.devices.take() { - Some(hashmap) => hashmap, - None => { - init_device_contexts_default()?; - ctx.devices - .take() - .ok_or_else(|| device_error(device_id, "Failed to initialize context"))? - } - }; + with_device_map(device_id, |hashmap| { if !hashmap.contains_key(&device_id) { - init_with_default_policy(&mut hashmap, device_id)?; + init_with_default_policy(hashmap, device_id)?; } let device_context = hashmap .get_mut(&device_id) .ok_or_else(|| device_error(device_id, "Failed to get context"))?; - let r = f(device_context); - ctx.devices.replace(Some(hashmap)); - Ok(r) + Ok(f(device_context)) }) } @@ -302,17 +311,23 @@ pub fn load_module_from_ptx( /// Inserts a compiled function into the per-device kernel cache. /// /// The function is keyed by the hash of `func_key`. Returns an error if a -/// different function was already registered under the same key (hash -/// collision). +/// function is already registered under the same key (a duplicate insert or +/// a hash collision); the existing entry is left in place rather than +/// replaced, so a caller that ignores the error still gets the function it +/// registered first from [`get_cuda_function`]. pub fn insert_cuda_function( device_id: usize, func_key: &impl FunctionKey, value: (Arc, Arc), ) -> Result<(), DeviceError> { with_global_device_context_mut(device_id, |dc| { - let key = func_key.get_hash_string(); - let res = dc.functions.insert(key, value); - device_assert(device_id, res.is_none(), "Unexpected cache key collision.") + match dc.functions.entry(func_key.get_hash_string()) { + Entry::Occupied(_) => Err(device_error(device_id, "Unexpected cache key collision.")), + Entry::Vacant(slot) => { + slot.insert(value); + Ok(()) + } + } })? } @@ -358,4 +373,96 @@ mod tests { .join() .expect("test thread should not panic"); } + + /// Runs `body` on a fresh thread (its own thread-local device map) if a + /// GPU is available; otherwise reports a skip. + fn with_gpu_thread(body: impl FnOnce() + Send + 'static) { + if let Err(error) = CudaContext::new(0) { + eprintln!("SKIPPED: CUDA unavailable ({error:?})"); + return; + } + std::thread::spawn(body) + .join() + .expect("test thread should not panic"); + } + + /// A failed lookup of one device must not destroy the contexts already + /// initialized for other devices on this thread. Device 0's context is + /// the same `Arc` before and after a request for a device that does not + /// exist. + #[test] + fn failed_device_lookup_keeps_the_other_contexts() { + with_gpu_thread(|| { + let before = with_cuda_context(0, Arc::clone).expect("device 0 initializes"); + + let bogus = with_cuda_context(usize::MAX, Arc::clone); + assert!(bogus.is_err(), "a nonexistent device must fail: {bogus:?}"); + + let after = with_cuda_context(0, Arc::clone).expect("device 0 is still there"); + assert!( + Arc::ptr_eq(&before, &after), + "device 0's context must survive the failed lookup" + ); + }); + } + + /// A second insert under an existing key must fail *without* replacing + /// the entry: the first function registered stays the one the cache + /// returns. + #[test] + fn insert_cuda_function_keeps_the_existing_entry_on_collision() { + const PTX: &str = r#" +.version 7.8 +.target sm_75 +.address_size 64 +.visible .entry first() { ret; } +.visible .entry second() { ret; } +"#; + + #[derive(Hash)] + struct Key; + impl FunctionKey for Key {} + + with_gpu_thread(|| { + let module = load_module_from_ptx(PTX, 0).expect("JIT the test PTX"); + let first = Arc::new(module.load_function("first").expect("first")); + let second = Arc::new(module.load_function("second").expect("second")); + + insert_cuda_function(0, &Key, (Arc::clone(&module), Arc::clone(&first))) + .expect("first insert"); + let collision = insert_cuda_function(0, &Key, (Arc::clone(&module), second)); + assert!(collision.is_err(), "duplicate key must be reported"); + + let cached = get_cuda_function(0, &Key).expect("the key is still registered"); + assert!( + Arc::ptr_eq(&cached, &first), + "the colliding insert must not have replaced the first function" + ); + }); + } + + /// `.sync()` must run the operation outside the thread-local borrow: an + /// operation that itself touches the device context (here, reads device + /// 0's `CudaContext`) must see the same context as the caller, not a + /// freshly rebuilt one. + #[test] + fn sync_does_not_rebuild_the_device_map_when_reentered() { + use crate::simt::device_operation::{empty, value, DeviceOperation}; + + with_gpu_thread(|| { + let outer = with_cuda_context(0, |ctx| Arc::as_ptr(ctx) as usize).expect("device 0"); + let inner = empty(|| { + value( + with_cuda_context(0, |ctx| Arc::as_ptr(ctx) as usize) + .expect("device 0 from inside the operation"), + ) + }) + .sync() + .expect("nested operation runs"); + assert_eq!( + outer, inner, + "the operation must observe the caller's device context" + ); + }); + } } diff --git a/cuda-async/src/simt/device_operation.rs b/cuda-async/src/simt/device_operation.rs index 21d336c6de..c4bb8dda0b 100644 --- a/cuda-async/src/simt/device_operation.rs +++ b/cuda-async/src/simt/device_operation.rs @@ -29,12 +29,11 @@ use crate::simt::device_context::with_default_device_policy; use crate::simt::device_future::DeviceFuture; use crate::simt::error::DeviceError; use crate::simt::scheduling_policies::SchedulingPolicy; -use cuda_core::{CudaContext, CudaStream}; -use std::cell::UnsafeCell; +use cuda_core::{CudaContext, CudaEvent, CudaStream}; use std::future::IntoFuture; use std::marker::PhantomData; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, Condvar, Mutex, MutexGuard}; +use std::thread::ThreadId; /// CUDA device ordinal. Type alias for readability. pub type Device = usize; @@ -220,7 +219,13 @@ pub trait DeviceOperation: /// Executes the operation synchronously on the default device using the /// thread-local scheduling policy. Blocks until the stream is idle. fn sync(self) -> Result<::Output, DeviceError> { - with_default_device_policy(|policy| policy.sync(self))? + // Take the policy out of the thread-local borrow before running the + // operation. `execute` runs user code, and a nested `.sync()` (or any + // device-context access) in there would find the thread's device map + // checked out and rebuild it from scratch: a second `CudaContext`, a + // second stream pool, and the originals dropped on the way back out. + let policy = with_default_device_policy(Arc::clone)?; + policy.sync(self) } /// Executes the operation on `stream` **without** synchronizing. @@ -430,13 +435,15 @@ where /// A [`DeviceOperation`] that immediately returns a pre-computed value without /// touching the GPU. -pub struct Value(T); - -/// # Safety /// -/// `Value` holds only the inner `T`. If `T` is `Send` the wrapper is too; -/// the bound is enforced by the `DeviceOperation` impl. -unsafe impl Send for Value {} +/// `Value` is `Send` exactly when `T` is, by the auto trait; there is no +/// manual impl to keep in step with the field. +/// +/// ```compile_fail,E0277 +/// fn assert_send() {} +/// assert_send::>>(); +/// ``` +pub struct Value(T); /// Returns the wrapped value directly -- no GPU work is performed. impl DeviceOperation for Value { @@ -482,26 +489,29 @@ impl IntoDeviceOperation for T { /// Useful when building the inner operation requires state only available /// after scheduling (though it does not receive the [`ExecutionContext`] -- /// see [`StreamOperation`] for that). -pub struct Empty, F: FnOnce() -> DO> { +/// +/// The closure must be `Send`: the operation is scheduled across threads, +/// and `Empty` is `Send` by the auto trait because that is its only field. +/// +/// ```compile_fail,E0277 +/// use cuda_async::simt::device_operation::{empty, value}; +/// let rc = std::rc::Rc::new(1_u8); +/// let _ = empty(move || value(*rc)); +/// ``` +pub struct Empty, F: FnOnce() -> DO + Send> { /// Closure that produces the inner operation. closure: F, } /// Wraps a closure in an [`Empty`] deferred operation. -pub fn empty, F: FnOnce() -> DO>( +pub fn empty, F: FnOnce() -> DO + Send>( closure: F, ) -> Empty { Empty { closure } } -/// # Safety -/// -/// The closure `F` and its produced operation `DO` are both `Send`. The -/// struct owns the closure exclusively. -unsafe impl, F: FnOnce() -> DO> Send for Empty {} - /// Invokes the closure to produce the inner operation, then executes it. -impl, F: FnOnce() -> DO> DeviceOperation +impl, F: FnOnce() -> DO + Send> DeviceOperation for Empty { type Output = O; @@ -515,7 +525,7 @@ impl, F: FnOnce() -> DO> Devi } /// Schedules via the thread-local default policy. -impl, F: FnOnce() -> DO> IntoFuture +impl, F: FnOnce() -> DO + Send> IntoFuture for Empty { type Output = Result; @@ -722,71 +732,224 @@ impl< } } -/// Shared state backing the [`SelectLeft`] / [`SelectRight`] split. +// ── Execute-once memoization (shared by the two halves of `unzip`) ───────── + +/// Where a memoized result was produced: the stream the operation ran on and +/// an event recorded on that stream immediately after its work. /// -/// Holds the source operation and its memoized results. The first selector -/// to execute triggers the source; subsequent selectors read the cached -/// values. Interior mutability is used via [`UnsafeCell`] because execution -/// is always sequential within a single stream. -pub struct Select> { - /// Guards one-shot execution of the source operation. - computed: AtomicBool, - /// Source operation. Consumed on the first call to [`execute`](Self::execute). - input: UnsafeCell>, - /// Cached left result. - left: UnsafeCell>, - /// Cached right result. - right: UnsafeCell>, +/// A consumer that picks up the cached value on a *different* stream has no +/// stream-order relationship to the producing work, so before the value is +/// handed out its stream is made to wait on the event (`cuStreamWaitEvent`). +/// Consumers on the producing stream are already ordered and skip the wait. +struct Producer { + stream: Arc, + event: CudaEvent, } -impl> Select { - /// Executes the source operation if it has not been executed yet, caching - /// both halves of the tuple. - /// - /// # Safety - /// - /// Must only be called from within a single-stream execution context. - /// Concurrent calls from different threads would race on the `UnsafeCell`s. - unsafe fn execute(self: &Arc, context: &ExecutionContext) -> Result<(), DeviceError> { - unsafe { - if !self.computed.load(Ordering::Acquire) { - let input = self.input.get(); - let input = input.as_mut(); - let input = input.unwrap().take().ok_or_else(|| { - crate::simt::error::device_error( - context.get_device_id(), - "Select operation failed.", - ) - })?; - let (left, right) = input.execute(context)?; - *self.left.get() = Some(left); - *self.right.get() = Some(right); - self.computed.store(true, Ordering::Release); +impl Producer { + /// Records the completion event for the work just enqueued on the + /// context's stream. + fn record(ctx: &ExecutionContext) -> Result { + let stream = Arc::clone(ctx.get_cuda_stream()); + let event = stream.record_event(None)?; + Ok(Self { stream, event }) + } + + /// Orders all future work on `consumer` after the producing work. + fn join(&self, consumer: &Arc) -> Result<(), DeviceError> { + if consumer.cu_stream() == self.stream.cu_stream() { + return Ok(()); + } + consumer.wait(&self.event)?; + Ok(()) + } +} + +enum OnceState { + /// Not yet executed; holds the operation. + Pending(Op), + /// Being executed by `thread`. Others wait on the condvar; the executing + /// thread itself re-entering is a bug reported as an error, not a + /// deadlock. + Running { thread: ThreadId }, + /// Executed; `producer` orders consumers on other streams after the + /// work. + Done { value: Out, producer: Producer }, + /// Execution failed (or the executor panicked); every consumer gets the + /// error. The operation was consumed and cannot be retried. + Failed(DeviceError), +} + +/// Runs an operation exactly once across any number of concurrent callers +/// and memoizes the outcome. +/// +/// This replaces the former check-then-act on `UnsafeCell`s behind +/// hand-written `Send` impls: two executors that both observed "not +/// computed" both took the operation (the second got "already taken") and +/// raced on the result cells. Here the state lives under a `Mutex`; the +/// first caller to see `Pending` takes the operation and runs it *outside* +/// the lock, later callers block on the condvar until it is `Done` or +/// `Failed`. `Send`/`Sync` are derived from the field types. +struct ExecuteOnce { + state: Mutex>, + settled: Condvar, +} + +impl ExecuteOnce { + fn pending(op: Op) -> Self { + Self { + state: Mutex::new(OnceState::Pending(op)), + settled: Condvar::new(), + } + } + + fn lock(&self) -> MutexGuard<'_, OnceState> { + // A poisoned lock only means a holder panicked; the state machine is + // never left mid-transition while the lock is held. + self.state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + /// Stores a value produced on `ctx`'s stream. + fn settle_done(&self, value: Out, ctx: &ExecutionContext) { + let outcome = match Producer::record(ctx) { + Ok(producer) => OnceState::Done { value, producer }, + Err(error) => { + // Without an event, consumers on other streams could not be + // ordered after the producing work. The value may own memory + // that work still writes; releasing it now would be the + // in-flight-free hazard, so leak it loudly instead. + eprintln!( + "cuda-async: leaking a memoized result after its completion \ + event could not be recorded: {error}" + ); + std::mem::forget(value); + OnceState::Failed(error) + } + }; + *self.lock() = outcome; + self.settled.notify_all(); + } + + /// Executes `run(op, ctx)` if this is the first caller (waiting for it + /// if another caller is mid-execution), then hands `take` the cached + /// value after ordering `ctx`'s stream behind the producing work. + fn execute( + &self, + ctx: &ExecutionContext, + run: impl FnOnce(Op, &ExecutionContext) -> Result, + take: impl FnOnce(&mut Out) -> Result, + ) -> Result { + // `Pending` is observed at most once per `ExecuteOnce`, so `run` is + // called at most once; the `Option` makes that visible to the borrow + // checker across the loop. + let mut run = Some(run); + let mut state = self.lock(); + loop { + match &*state { + OnceState::Pending(_) => { + let running = OnceState::Running { + thread: std::thread::current().id(), + }; + let OnceState::Pending(op) = std::mem::replace(&mut *state, running) else { + unreachable!("matched Pending above"); + }; + drop(state); + let run = run.take().expect("Pending is observed at most once"); + // If `run` unwinds, settle as Failed so waiters are + // released instead of blocking on `Running` forever. + let mut settle_on_unwind = SettleOnUnwind { once: Some(self) }; + let outcome = run(op, ctx); + settle_on_unwind.once = None; + match outcome { + Ok(value) => self.settle_done(value, ctx), + Err(error) => { + *self.lock() = OnceState::Failed(error); + self.settled.notify_all(); + } + } + state = self.lock(); + } + OnceState::Running { thread } => { + if *thread == std::thread::current().id() { + return Err(DeviceError::Internal( + "execute-once operation re-entered from inside its own \ + execution (an unzipped op executing its other half)" + .into(), + )); + } + state = self + .settled + .wait(state) + .unwrap_or_else(|poisoned| poisoned.into_inner()); + } + OnceState::Done { .. } => break, + OnceState::Failed(error) => return Err(error.clone()), } - Ok(()) } + let OnceState::Done { value, producer } = &mut *state else { + unreachable!("loop exits only on Done"); + }; + producer.join(ctx.get_cuda_stream())?; + take(value) } +} - /// Takes the cached left value. Must be called after [`execute`](Self::execute). - /// - /// # Safety - /// - /// Caller must ensure `execute` has completed and this is called at most once. - unsafe fn left(&self) -> T1 { - let cell = self.left.get(); - let cell = unsafe { cell.as_mut() }; - cell.unwrap().take().unwrap() +/// Marks an `ExecuteOnce` as failed if the executing closure unwinds. +struct SettleOnUnwind<'a, Op, Out> { + once: Option<&'a ExecuteOnce>, +} + +impl Drop for SettleOnUnwind<'_, Op, Out> { + fn drop(&mut self) { + if let Some(once) = self.once { + *once.lock() = OnceState::Failed(DeviceError::Internal( + "execute-once operation panicked while executing".into(), + )); + once.settled.notify_all(); + } } +} - /// Takes the cached right value. Must be called after [`execute`](Self::execute). +/// Execute-once state shared by the two halves of an +/// [`unzip`](Unzippable2::unzip). +/// +/// Whichever half executes first runs the input; the other half waits for +/// it if it is mid-execution (on another thread), then takes its side of the +/// result. A half executed on a different stream than the producing one has +/// its stream wait on the producer's completion event first. `Send`/`Sync` +/// are derived: the input is `DeviceOperation: Send`, and both halves are +/// `Send`, so no manual impl is needed. +pub struct Select> { + once: ExecuteOnce, Option)>, +} + +impl> Select { + /// Runs the input if needed and hands `take` the stored halves. /// /// # Safety /// - /// Caller must ensure `execute` has completed and this is called at most once. - unsafe fn right(&self) -> T2 { - let cell = self.right.get(); - let cell = unsafe { cell.as_mut() }; - cell.unwrap().take().unwrap() + /// Same contract as [`DeviceOperation::execute`]: GPU work may still be + /// in flight when this returns. + unsafe fn execute( + &self, + context: &ExecutionContext, + take: impl FnOnce(&mut (Option, Option)) -> Option, + side: &'static str, + ) -> Result { + self.once.execute( + context, + |input, ctx| { + let (left, right) = unsafe { input.execute(ctx) }?; + Ok((Some(left), Some(right))) + }, + |halves| { + take(halves).ok_or_else(|| { + DeviceError::Internal(format!("unzip: the {side} result was already taken")) + }) + }, + ) } } @@ -799,15 +962,6 @@ pub struct SelectLeft select: Arc>, } -/// # Safety -/// -/// The `Arc>` is safe to send because `Select` is only mutated -/// through `UnsafeCell` during single-stream execution, never concurrently. -unsafe impl> Send - for SelectLeft -{ -} - /// Triggers the shared source operation (if not yet done) and returns the /// left element. impl> DeviceOperation @@ -817,8 +971,8 @@ impl Result { unsafe { - self.select.execute(context)?; - Ok(self.select.left()) + self.select + .execute(context, |(left, _)| left.take(), "left") } } } @@ -846,14 +1000,6 @@ pub struct SelectRight>, } -/// # Safety -/// -/// See [`SelectLeft`]'s `Send` impl. -unsafe impl> Send - for SelectRight -{ -} - /// Triggers the shared source operation (if not yet done) and returns the /// right element. impl> DeviceOperation @@ -863,8 +1009,8 @@ impl Result { unsafe { - self.select.execute(context)?; - Ok(self.select.right()) + self.select + .execute(context, |(_, right)| right.take(), "right") } } } @@ -889,13 +1035,9 @@ impl>( input: DI, ) -> (SelectLeft, SelectRight) { - let select = Select { - computed: AtomicBool::new(false), - input: UnsafeCell::new(Some(input)), - left: UnsafeCell::new(None), - right: UnsafeCell::new(None), - }; - let select = Arc::new(select); + let select = Arc::new(Select { + once: ExecuteOnce::pending(input), + }); let out1 = SelectLeft { select: Arc::clone(&select), }; diff --git a/cuda-async/src/simt/scheduling_policies.rs b/cuda-async/src/simt/scheduling_policies.rs index a32d1ee885..ea0710dea8 100644 --- a/cuda-async/src/simt/scheduling_policies.rs +++ b/cuda-async/src/simt/scheduling_policies.rs @@ -212,10 +212,19 @@ impl SingleStream { /// /// # Safety /// - /// The caller must call [`SchedulingPolicy::init`] before use. + /// The caller must call [`SchedulingPolicy::init`] before use. Using the + /// policy before initialization produces an error, not undefined + /// behavior; the `unsafe` mirrors [`StreamPoolRoundRobin::new`]. pub unsafe fn new() -> Self { Self { stream: None } } + + /// The stream, or a [`DeviceError::Scheduling`] if `init` has not run. + fn initialized_stream(&self) -> Result<&Arc, DeviceError> { + self.stream.as_ref().ok_or_else(|| { + DeviceError::Scheduling("SingleStream policy used before init.".to_string()) + }) + } } impl SchedulingPolicy for SingleStream { @@ -230,7 +239,7 @@ impl SchedulingPolicy for SingleStream { &self, op: O, ) -> Result, DeviceError> { - let stream = self.stream.as_ref().unwrap(); + let stream = self.initialized_stream()?; // See `StreamPoolRoundRobin::schedule` for why the fields are // spelled out instead of using `..Default::default()`. Ok(DeviceFuture { @@ -245,7 +254,33 @@ impl SchedulingPolicy for SingleStream { /// Executes `op` synchronously on the single stream. fn sync>(&self, op: O) -> Result { - let stream = self.stream.as_ref().unwrap(); - op.sync_on(stream) + op.sync_on(self.initialized_stream()?) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::simt::device_operation::value; + + /// An uninitialized policy has no stream to hand out; that is a + /// scheduling error for the caller, not a panic inside the runtime. + #[test] + fn single_stream_reports_use_before_init_as_an_error() { + // SAFETY: the policy is never initialized on purpose; the contract + // says that yields an error rather than undefined behavior. + let policy = unsafe { SingleStream::new() }; + + let scheduled = policy.schedule(value(7_u32)).map(drop); + assert!( + matches!(scheduled, Err(DeviceError::Scheduling(_))), + "{scheduled:?}" + ); + + let synced = policy.sync(value(7_u32)); + assert!( + matches!(synced, Err(DeviceError::Scheduling(_))), + "{synced:?}" + ); } } diff --git a/cuda-async/tests/simt_execute_once.rs b/cuda-async/tests/simt_execute_once.rs new file mode 100644 index 0000000000..4d6e1e6a90 --- /dev/null +++ b/cuda-async/tests/simt_execute_once.rs @@ -0,0 +1,178 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! The two halves of a simt `unzip` execute their input exactly once, no +//! matter how the halves are driven: from two threads at the same time, as +//! two futures joined on one thread, or one from inside the other. +//! +//! Requires a GPU. + +use cuda_async::simt::device_context::init_device_contexts; +use cuda_async::simt::device_future::DeviceFuture; +use cuda_async::simt::device_operation::{value, DeviceOperation, ExecutionContext, Unzippable2}; +use cuda_async::simt::error::DeviceError; +use std::future::IntoFuture; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Barrier, Mutex}; +use std::time::Duration; + +/// Runs `f` on a fresh thread so the thread-local device map starts clean. +fn on_fresh_thread(f: F) { + std::thread::spawn(f).join().expect("test thread panicked"); +} + +/// An operation that bumps `counter` every time it executes and yields `val`. +fn counted_op( + counter: &Arc, + val: T, +) -> impl DeviceOperation { + let counter = Arc::clone(counter); + value(()).and_then(move |()| { + counter.fetch_add(1, Ordering::SeqCst); + value(val) + }) +} + +/// Like [`counted_op`], but slow enough that a second executor on another +/// thread is guaranteed to arrive while the first is still running. +fn slow_counted_op( + counter: &Arc, + val: T, +) -> impl DeviceOperation { + let counter = Arc::clone(counter); + value(()).and_then(move |()| { + counter.fetch_add(1, Ordering::SeqCst); + std::thread::sleep(Duration::from_millis(150)); + value(val) + }) +} + +/// An operation that fails at execute time. +struct FailingOp; + +impl DeviceOperation for FailingOp { + type Output = (u64, u64); + unsafe fn execute(self, _context: &ExecutionContext) -> Result<(u64, u64), DeviceError> { + Err(DeviceError::Internal("failing op".into())) + } +} + +impl IntoFuture for FailingOp { + type Output = Result<(u64, u64), DeviceError>; + type IntoFuture = DeviceFuture<(u64, u64), FailingOp>; + fn into_future(self) -> Self::IntoFuture { + DeviceFuture::failed(DeviceError::Internal("not used".into())) + } +} + +/// The two halves are executed at the same time on two threads. Each thread +/// initializes its own device contexts and therefore its own stream pool, so +/// the half that arrives second consumes the value on a different stream and +/// exercises the wait on the producer's completion event. The input runs +/// once and each half receives its side. +#[test] +fn unzip_concurrent_halves_run_input_once() { + on_fresh_thread(|| { + init_device_contexts(0, 1).expect("init failed (requires GPU)"); + + let counter = Arc::new(AtomicUsize::new(0)); + let (left, right) = slow_counted_op(&counter, (1_u64, 2_u64)).unzip(); + let barrier = Arc::new(Barrier::new(2)); + + let left_thread = { + let barrier = Arc::clone(&barrier); + std::thread::spawn(move || { + init_device_contexts(0, 1).expect("per-thread init failed"); + barrier.wait(); + left.sync().expect("left failed") + }) + }; + let right_thread = { + let barrier = Arc::clone(&barrier); + std::thread::spawn(move || { + init_device_contexts(0, 1).expect("per-thread init failed"); + barrier.wait(); + right.sync().expect("right failed") + }) + }; + + assert_eq!(left_thread.join().expect("left panicked"), 1); + assert_eq!(right_thread.join().expect("right panicked"), 2); + assert_eq!( + counter.load(Ordering::SeqCst), + 1, + "the input must run exactly once" + ); + }); +} + +/// Both halves driven concurrently as futures on one thread. +#[test] +fn unzip_join_on_one_thread() { + on_fresh_thread(|| { + init_device_contexts(0, 1).expect("init failed (requires GPU)"); + + let counter = Arc::new(AtomicUsize::new(0)); + let (left, right) = counted_op(&counter, (3_u64, 4_u64)).unzip(); + + let (a, b) = futures::executor::block_on(async { + futures::join!(left.into_future(), right.into_future()) + }); + assert_eq!(a.expect("left failed"), 3); + assert_eq!(b.expect("right failed"), 4); + assert_eq!(counter.load(Ordering::SeqCst), 1); + }); +} + +/// A failing input reports its own error to both halves; the second half +/// does not see an "already taken" error and the input is not retried. +#[test] +fn unzip_failure_is_reported_to_both_halves() { + on_fresh_thread(|| { + init_device_contexts(0, 1).expect("init failed (requires GPU)"); + + let (left, right) = FailingOp.unzip(); + for result in [left.sync(), right.sync()] { + assert!( + matches!(&result, Err(DeviceError::Internal(m)) if m == "failing op"), + "expected the input's own error, got {result:?}" + ); + } + }); +} + +/// An input that executes the other half from inside its own execution must +/// get an error, not deadlock on the shared state. +#[test] +fn unzip_reentrant_execution_is_an_error_not_a_deadlock() { + on_fresh_thread(|| { + init_device_contexts(0, 1).expect("init failed (requires GPU)"); + + // The right half is handed to the input through a type-erased slot: + // naming its type inside the closure that produces it would be + // circular. + type RunRight = Box Result + Send>; + let slot: Arc>> = Arc::new(Mutex::new(None)); + let slot_in_op = Arc::clone(&slot); + let (left, right) = value(()) + .and_then(move |()| { + let run_right = slot_in_op + .lock() + .unwrap() + .take() + .expect("slot is filled before execution"); + let nested = run_right(); + assert!( + matches!(&nested, Err(DeviceError::Internal(m)) if m.contains("re-entered")), + "re-entrant execution must be reported, got {nested:?}" + ); + value((1_u64, 2_u64)) + }) + .unzip(); + *slot.lock().unwrap() = Some(Box::new(move || right.sync())); + + assert_eq!(left.sync().expect("outer execution failed"), 1); + }); +} diff --git a/cuda-core/README.md b/cuda-core/README.md index 2814763b15..fcdc38c6c1 100644 --- a/cuda-core/README.md +++ b/cuda-core/README.md @@ -1,12 +1,56 @@ -# CUDA +# cuda-core -`cuda-core` wraps `cuda-bindings` with the CUDA types and functions used by -the rest of the workspace. +`cuda-core` is the thin safe wrapper layer over `cuda-bindings`. It exposes the +lower-level CUDA concepts used by the rest of the workspace without requiring +most crates to touch raw FFI directly. -# Testing +## Two host models -Run the crate tests with: +The crate currently carries two host-side models of the same driver: + +- **The crate root** (`Device`, `Stream`, `Module`, `MemPool`, `LaunchConfig`, + and the reviewed `vmm` module) is the model the rest of the workspace + (`cutile`, `cuda-async`) is written against. +- **`cuda_core::simt`** is the cuda-oxide host surface (`CudaContext`, + `CudaStream`, `CudaModule`, `DeviceBuffer`, `PinnedHostBuffer`, + `ConstantHandle`, the typed launch contracts, and its own `vmm` and `peer` + modules), carried over as-is for the shared host-crate migration. It is a + temporary bridge: the two models get reconciled after the repository + migration, and `simt` is deletable as a unit once that happens. Every + `simt` name that does not collide with the root surface is also re-exported + at the crate root; the two that collide, `simt::LaunchConfig` and + `simt::vmm`, are reachable only through `simt::`. + +New code in this workspace targets the root types. Code ported from +cuda-oxide keeps using `simt` until the reconciliation. + +The multicast half of `simt::vmm` (`MulticastObject`, `Mapping::new_multicast`) +is compiled only when the CUDA toolkit headers declare the `cuMulticast*` API +(CUDA 12.1+); `build.rs` probes for it and sets the `cuda_has_multicast` cfg. + +## Features + +- `f16` (default off, nightly only): implements `DeviceCopy` for the unstable + `f16` primitive. `half::f16` and `half::bf16` are supported unconditionally. + +## Companion crate + +`cuda-core-derive` provides `#[derive(DeviceCopy)]` for plain-data structs and +unions. It is re-exported next to the trait, so `use cuda_core::DeviceCopy;` +brings both the trait and the derive into scope. + +## Testing + +Unit and doc tests run without a GPU: ```bash -cargo test -p cuda-core +cargo test -p cuda-core --lib +cargo test -p cuda-core --doc ``` + +The integration tests need a GPU. `tests/vmm.rs` and `tests/vmm_multicast.rs` +exercise the root `vmm` module; `tests/simt_*.rs` exercise the `simt` surface. +Tests that need hardware this machine may lack self-skip: a second GPU for +`simt_vmm_p2p` and `vmm_multicast`, NVLink switch multicast for the multicast +tests. `scripts/run_gpu_tests.sh` lists every GPU test target and is the +reference for what must pass; a new `simt_*` test file must be added there. diff --git a/cuda-core/src/lib.rs b/cuda-core/src/lib.rs index ce2a802e5c..fb0de6c3f1 100644 --- a/cuda-core/src/lib.rs +++ b/cuda-core/src/lib.rs @@ -31,9 +31,9 @@ pub use simt::{ launch_kernel_cooperative, launch_kernel_cooperative_on_stream, launch_kernel_ex, launch_kernel_ex_cooperative, launch_kernel_ex_cooperative_on_stream, launch_kernel_ex_on_stream, launch_kernel_on_stream, BlockRequirement, ConstantHandle, - ContextLimit, CudaContext, CudaEvent, CudaFunction, CudaModule, CudaStream, DeviceBuffer, - DeviceCopy, DeviceLaunchLimits, DynamicSharedMemoryRequirement, EmbeddedModule, - EmbeddedModuleError, KernelLaunchConfig, KernelLaunchContract, LaunchAxis, LaunchConfig1D, - LaunchConfig2D, LaunchConfig3D, LaunchContractError, LaunchContractSpec, LaunchDimension, - PinnedHostBuffer, PreparedLaunch, StreamPriorityRange, SyncPolicy, + ContextLimit, CoordinateRequirement, CudaContext, CudaEvent, CudaFunction, CudaModule, + CudaStream, DeviceBuffer, DeviceCopy, DeviceLaunchLimits, DynamicSharedMemoryRequirement, + EmbeddedModule, EmbeddedModuleError, KernelLaunchConfig, KernelLaunchContract, LaunchAxis, + LaunchConfig1D, LaunchConfig2D, LaunchConfig3D, LaunchContractError, LaunchContractSpec, + LaunchDimension, PinnedHostBuffer, PreparedLaunch, StreamPriorityRange, SyncPolicy, }; diff --git a/cuda-core/src/simt/context.rs b/cuda-core/src/simt/context.rs index 2a76c5a389..72dcdefc47 100644 --- a/cuda-core/src/simt/context.rs +++ b/cuda-core/src/simt/context.rs @@ -21,7 +21,7 @@ use crate::simt::launch::DeviceLaunchLimits; use crate::simt::stream::CudaStream; use std::ffi::c_int; use std::mem::MaybeUninit; -use std::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicU32, AtomicUsize, Ordering}; use std::sync::Arc; /// Owns a reference to a CUDA device's primary context. @@ -40,11 +40,10 @@ pub struct CudaContext { pub(crate) cu_ctx: cuda_bindings::CUcontext, /// Zero-based device ordinal passed to [`CudaContext::new`]. pub(crate) ordinal: usize, - /// Number of live [`CudaStream`] instances sharing this context. + /// Number of live non-default [`CudaStream`] instances created through + /// this handle. Incremented once `cuStreamCreate` has succeeded, + /// decremented by `CudaStream::drop`. pub(crate) num_streams: AtomicUsize, - /// When `true`, the first [`new_stream`](CudaContext::new_stream) call - /// synchronizes the context to establish a clean ordering baseline. - pub(crate) event_tracking: AtomicBool, /// Sticky error state recorded by [`record_err`](CudaContext::record_err). /// Stores the raw `CUresult` value, or `0` if no error. pub(crate) error_state: AtomicU32, @@ -53,8 +52,8 @@ pub struct CudaContext { /// # Safety /// /// `CUdevice` and `CUcontext` are process-wide handles. All mutable state -/// (`num_streams`, `event_tracking`, `error_state`) uses atomics. The CUDA -/// driver itself is thread-safe for distinct contexts, and the +/// (`num_streams`, `error_state`) uses atomics. The CUDA driver itself is +/// thread-safe for distinct contexts, and the /// [`bind_to_thread`](CudaContext::bind_to_thread) mechanism ensures the /// correct context is current before each call. unsafe impl Send for CudaContext {} @@ -63,12 +62,14 @@ unsafe impl Sync for CudaContext {} /// Releases the primary context on drop. /// -/// Binds the context to the current thread first (required by -/// `cuDevicePrimaryCtxRelease`). Errors during teardown are recorded via -/// [`record_err`](CudaContext::record_err) rather than panicking. +/// `cuDevicePrimaryCtxRelease` is addressed by device and does not require +/// the context to be current, so nothing is bound here: making a context +/// current only to release it would leave the calling thread with a +/// dangling current context when this was the last reference. Errors during +/// teardown are recorded via [`record_err`](CudaContext::record_err) rather +/// than panicking. impl Drop for CudaContext { fn drop(&mut self) { - self.record_err(self.bind_to_thread()); let ctx = std::mem::replace(&mut self.cu_ctx, std::ptr::null_mut()); if !ctx.is_null() { self.record_err(unsafe { @@ -343,7 +344,6 @@ impl CudaContext { cu_ctx, ordinal, num_streams: AtomicUsize::new(0), - event_tracking: AtomicBool::new(true), error_state: AtomicU32::new(0), }); ctx.bind_to_thread()?; @@ -368,11 +368,15 @@ impl CudaContext { /// Binds this context to the calling thread if not already current. /// /// Checks [`check_err`](Self::check_err) first and propagates any sticky - /// error. Skips the `cuCtxSetCurrent` call when the context is already - /// bound, avoiding an unnecessary driver round-trip. + /// error, in which case the context is **not** made current. Skips the + /// `cuCtxSetCurrent` call when the context is already bound, avoiding an + /// unnecessary driver round-trip. /// /// Most methods on [`CudaStream`], [`CudaEvent`](crate::CudaEvent), and - /// [`CudaModule`](crate::CudaModule) call this internally. + /// [`CudaModule`](crate::CudaModule) call this internally before they + /// enqueue work, so a recorded error surfaces before anything new is + /// submitted. The synchronize wrappers deliberately do not: see + /// [`synchronize`](Self::synchronize). /// /// CUcontext is the backing runtime object, and CUmodule / CUfunction / CUstream /// are opaque handles to objects created under that context. bind_to_thread() @@ -381,6 +385,16 @@ impl CudaContext { /// can fail. pub fn bind_to_thread(&self) -> Result<(), DriverError> { self.check_err()?; + self.make_current() + } + + /// Makes this context current on the calling thread without consulting + /// the sticky error state. + /// + /// This is the half of [`bind_to_thread`](Self::bind_to_thread) that + /// synchronize wrappers and `Drop` impls use: a recorded error must not + /// stop a wait or a release from reaching the driver. + pub(crate) fn make_current(&self) -> Result<(), DriverError> { let mut current = MaybeUninit::uninit(); unsafe { cuda_bindings::cuCtxGetCurrent(current.as_mut_ptr()).result()?; @@ -395,9 +409,24 @@ impl CudaContext { /// Blocks the calling thread until all preceding work in this context /// completes. /// - /// Binds the context first, then calls `cuCtxSynchronize`. + /// The driver call is issued unconditionally, and only afterwards is the + /// sticky error state drained and returned. A `Result::Err` from this + /// method therefore never means the wait was skipped: either + /// `cuCtxSynchronize` itself failed, or it completed and a previously + /// recorded error is being reported. Safe wrappers that hand out host + /// memory after an enqueued copy rely on this ordering. pub fn synchronize(&self) -> Result<(), DriverError> { - self.bind_to_thread()?; + self.synchronize_driver()?; + self.check_err() + } + + /// `cuCtxSynchronize` without consulting the sticky error state. + /// + /// Returns `Err` only when the driver refused or failed the wait, which + /// is the signal callers use to decide that host memory may still be a + /// DMA target and must be leaked rather than freed. + pub(crate) fn synchronize_driver(&self) -> Result<(), DriverError> { + self.make_current()?; unsafe { cuda_bindings::cuCtxSynchronize() }.result() } @@ -418,9 +447,8 @@ impl CudaContext { /// The stream is created with `CU_STREAM_NON_BLOCKING`, so it does not /// implicitly synchronize with the default stream. /// - /// On the first call (when `num_streams` transitions from 0 to 1), the - /// context is synchronized to establish a clean ordering baseline if - /// `event_tracking` is enabled. + /// When this handle has no live streams yet, the context is synchronized + /// first to establish a clean ordering baseline for the new stream. pub fn new_stream(self: &Arc) -> Result, DriverError> { self.create_stream(None) } @@ -459,8 +487,9 @@ impl CudaContext { priority: Option, ) -> Result, DriverError> { self.bind_to_thread()?; - let prev = self.num_streams.fetch_add(1, Ordering::Relaxed); - if prev == 0 && self.event_tracking.load(Ordering::Relaxed) { + // Baseline for the first stream. Two threads racing here may both + // synchronize, which is harmless. + if self.num_streams.load(Ordering::Relaxed) == 0 { self.synchronize()?; } let flags = cuda_bindings::CUstream_flags_enum_CU_STREAM_NON_BLOCKING; @@ -477,6 +506,9 @@ impl CudaContext { .result()?; cu_stream.assume_init() }; + // Count the stream only once it exists: a failed create has no + // `CudaStream` whose drop would undo the increment. + self.num_streams.fetch_add(1, Ordering::Relaxed); Ok(Arc::new(CudaStream { cu_stream, ctx: self.clone(), @@ -764,12 +796,65 @@ impl CudaContext { /// Records a driver error into the sticky error state. /// /// Used during [`Drop`] paths where returning a `Result` is not possible. - /// If `result` is `Err`, the raw error code is stored; subsequent - /// [`check_err`](Self::check_err) or [`bind_to_thread`](Self::bind_to_thread) - /// calls will surface it. A later store overwrites an earlier one. + /// If `result` is `Err`, the raw error code is stored; a later store + /// overwrites an earlier one, and `Ok` values are ignored. + /// + /// The recorded error is reported by the next fallible call on this + /// context: [`check_err`](Self::check_err) and + /// [`bind_to_thread`](Self::bind_to_thread) return it up front, before + /// doing anything else, so methods that enqueue work refuse to submit + /// more while an error is pending. The synchronize wrappers + /// ([`synchronize`](Self::synchronize), [`CudaStream::synchronize`], + /// [`CudaEvent::synchronize`](crate::CudaEvent::synchronize)) instead + /// issue their driver wait first and report the recorded error + /// afterwards, and `Drop` impls never consult it, so a recorded error + /// cannot skip a wait or leak a release. This method is public so + /// callers with their own `Drop` types can participate in the same + /// scheme; it is not a way to abort in-flight work. pub fn record_err(&self, result: Result) { if let Err(err) = result { self.error_state.store(err.0, Ordering::Relaxed) } } } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::Ordering; + + /// `num_streams` counts exactly the live streams created through the + /// handle: it moves only when a create has succeeded and when a stream + /// drops, and the default stream is never counted. + #[test] + fn live_stream_count_tracks_creates_and_drops() { + let ctx = match CudaContext::new(0) { + Ok(ctx) => ctx, + Err(error) => { + eprintln!("SKIPPED: CUDA unavailable ({error:?})"); + return; + } + }; + let count = || ctx.num_streams.load(Ordering::Relaxed); + assert_eq!(count(), 0); + + let default = ctx.default_stream(); + assert_eq!(count(), 0, "the default stream is not a created stream"); + + let first = ctx.new_stream().expect("new_stream"); + assert_eq!(count(), 1); + let forked = first.fork().expect("fork"); + assert_eq!(count(), 2); + let prioritized = ctx + .new_stream_with_priority(0) + .expect("new_stream_with_priority"); + assert_eq!(count(), 3); + + drop(forked); + assert_eq!(count(), 2); + drop(prioritized); + drop(first); + drop(default); + assert_eq!(count(), 0); + } +} diff --git a/cuda-core/src/simt/device_buffer.rs b/cuda-core/src/simt/device_buffer.rs index 9057de9be9..560847aef0 100644 --- a/cuda-core/src/simt/device_buffer.rs +++ b/cuda-core/src/simt/device_buffer.rs @@ -19,6 +19,17 @@ //! freeing on a chosen stream; its caller takes over the obligation to //! order every other stream that uses the buffer before that stream. //! +//! # Context binding +//! +//! Every method that reaches the driver first binds `stream.context()` to the +//! calling thread ([`CudaContext::bind_to_thread`]). `cuMemAlloc` needs a +//! current context, and the default stream (a null `CUstream`) is resolved +//! against whichever context is current, so without the bind a thread that +//! has never touched CUDA fails with `CUDA_ERROR_INVALID_CONTEXT` and a +//! thread bound to another device works in the wrong context. The bind also +//! drains the context's sticky error state, so a pending recorded error is +//! returned before any transfer is enqueued. +//! //! # Quick start //! //! ```ignore @@ -59,6 +70,17 @@ use crate::simt::stream::CudaStream; /// `NonZeroU32` are `Copy`, but not every byte pattern is a valid value of /// those types. `DeviceCopy` is the stronger promise required when /// `DeviceBuffer` turns raw device bytes back into initialized Rust values. +/// +/// Raw pointers implement `DeviceCopy` only for `Sized` pointees. A fat +/// pointer such as `*const dyn Trait` carries a vtable address that only the +/// host's own code image can supply, so one read back from device memory is +/// not a valid pointer of its type; slice pointers are excluded on the same +/// grounds until something needs them. +/// +/// ```compile_fail,E0277 +/// fn assert_device_copy() {} +/// assert_device_copy::<*const dyn std::fmt::Debug>(); +/// ``` pub unsafe trait DeviceCopy: Copy {} macro_rules! impl_device_copy { @@ -88,8 +110,11 @@ impl_device_copy!( ); unsafe impl DeviceCopy for [T; N] {} -unsafe impl DeviceCopy for *const T {} -unsafe impl DeviceCopy for *mut T {} +// Thin pointers only: every address is a valid `*const T` / `*mut T` value, +// which is not true of the metadata half of a fat pointer (see the trait +// docs). The pointee itself need not be `DeviceCopy`; nothing dereferences it. +unsafe impl DeviceCopy for *const T {} +unsafe impl DeviceCopy for *mut T {} // Wrapper types that don't change the byte representation: a value of the // wrapper has the same layout and validity invariants as the inner `T`. @@ -153,28 +178,37 @@ pub struct DeviceBuffer { _marker: PhantomData, } -// SAFETY: CUdeviceptr is a u64 handle valid across threads when the owning -// context is bound. The PhantomData is Send if T is Send. +// SAFETY: the handle fields are plain values (`CUdeviceptr`, lengths) and +// `Arc`s that are themselves `Send + Sync`; a buffer of `T` elements moves +// between threads under the same rule as `Vec`. unsafe impl Send for DeviceBuffer {} -// SAFETY: &DeviceBuffer only exposes cu_deviceptr() and len(), both of which -// return Copy values. No interior mutability. +// SAFETY: no interior mutability. Every `&self` method either reads the handle +// fields or enqueues a device-side *read* of the buffer (`to_host_vec`, +// `copy_to_host`, `copy_to_pinned_host[_async]`, and the source side of +// `copy_from_device_async`); all writes go through `&mut self` or `self`. +// Shared access hands out `T` values by bit copy, so the `T: Sync` bound +// mirrors the one on `&[T]: Send`. unsafe impl Sync for DeviceBuffer {} impl Drop for DeviceBuffer { fn drop(&mut self) { if self.ptr != 0 { - self.ctx.record_err(self.ctx.bind_to_thread()); - // Safe buffer operations can enqueue work on any stream in this - // context. Synchronize all of them before implicitly freeing a - // stream-ordered allocation. + // Neither the bind nor the context-wide wait consults the sticky + // error state: a recorded error must not turn this drop into a + // leak, and it stays recorded for the next fallible call to + // report. Only a failed driver call skips the free, since memory + // the device may still be using cannot be released. let result = match &self.dealloc_stream { - Some(stream) => match self.ctx.synchronize() { - Ok(()) => unsafe { - crate::simt::memory::free_async(self.ptr, stream.cu_stream()) - }, - Err(error) => Err(error), - }, - None => unsafe { crate::simt::memory::free_sync(self.ptr) }, + // Safe buffer operations can enqueue work on any stream in + // this context. Synchronize all of them before implicitly + // freeing a stream-ordered allocation. + Some(stream) => self.ctx.synchronize_driver().and_then(|()| unsafe { + crate::simt::memory::free_async(self.ptr, stream.cu_stream()) + }), + None => self + .ctx + .make_current() + .and_then(|()| unsafe { crate::simt::memory::free_sync(self.ptr) }), }; self.ctx.record_err(result); } @@ -407,6 +441,15 @@ impl DeviceBuffer { /// An empty `data` slice yields an empty buffer without touching the /// driver allocator. /// + /// # Sticky errors + /// + /// A pending recorded error ([`CudaContext::record_err`]) is returned + /// before anything is allocated or enqueued. One recorded while the copy + /// is in flight is returned only after the stream has been synchronized, + /// because [`CudaStream::synchronize`] always issues the driver wait + /// first; an `Err` from this function never leaves `data` in use by the + /// device unless `cuStreamSynchronize` itself failed. + /// /// # Allocation safety on error /// /// The buffer takes ownership of the device allocation immediately after @@ -416,6 +459,7 @@ impl DeviceBuffer { /// device memory is leaked. pub fn from_host(stream: &CudaStream, data: &[T]) -> Result { let ctx = stream.context().clone(); + ctx.bind_to_thread()?; let len = data.len(); let num_bytes = allocation_size::(len)?; @@ -432,17 +476,17 @@ impl DeviceBuffer { // stream's context; ownership transfers to `buf` here so any early // return below frees it through the buffer's own `Drop`. let buf = unsafe { Self::from_raw_parts(ptr, len, ctx) }; - let enqueue_result = unsafe { + // A failed enqueue leaves nothing in flight, so there is nothing to + // wait for before `data`'s borrow ends. + unsafe { crate::simt::memory::memcpy_htod_async( buf.ptr, data.as_ptr(), num_bytes, stream.cu_stream(), - ) - }; - let sync_result = stream.synchronize(); - enqueue_result?; - sync_result?; + )?; + } + stream.synchronize()?; Ok(buf) } @@ -461,6 +505,7 @@ impl DeviceBuffer { data: &[T], ) -> Result { let ctx = stream.context().clone(); + ctx.bind_to_thread()?; let len = data.len(); let num_bytes = std::mem::size_of_val(data); @@ -544,6 +589,7 @@ impl DeviceBuffer { /// no device memory is leaked. pub fn zeroed(stream: &CudaStream, len: usize) -> Result { let ctx = stream.context().clone(); + ctx.bind_to_thread()?; let num_bytes = allocation_size::(len)?; // cuMemAlloc rejects zero-byte requests with CUDA_ERROR_INVALID_VALUE, @@ -568,9 +614,17 @@ impl DeviceBuffer { /// Copies the entire buffer back to the host, returning a `Vec`. /// /// Synchronizes on `stream` before returning so the host vector is safe - /// to read immediately. + /// to read immediately. See [`Self::from_host`] for how a sticky error + /// is ordered against the wait. If `cuStreamSynchronize` itself fails, + /// the vector is leaked instead of freed, since the DMA may still be + /// writing it. pub fn to_host_vec(&self, stream: &CudaStream) -> Result, DriverError> { + stream.context().bind_to_thread()?; let mut host = Vec::with_capacity(self.len); + if self.num_bytes() == 0 { + return Ok(host); + } + // A failed enqueue leaves nothing in flight; `host` drops normally. unsafe { crate::simt::memory::memcpy_dtoh_async( host.as_mut_ptr(), @@ -579,15 +633,22 @@ impl DeviceBuffer { stream.cu_stream(), )?; } - stream.synchronize()?; + if let Err(err) = stream.synchronize_driver() { + // The driver did not complete the wait, so the copy may still + // target `host`: leaking is the only sound way out. + std::mem::forget(host); + return Err(err); + } unsafe { host.set_len(self.len) }; + stream.context().check_err()?; Ok(host) } /// Copies the buffer contents into an existing host slice. /// /// Synchronizes on `stream` before returning. Panics if - /// `dst.len() < self.len()`. + /// `dst.len() < self.len()`. See [`Self::from_host`] for how a sticky + /// error is ordered against the wait. pub fn copy_to_host(&self, stream: &CudaStream, dst: &mut [T]) -> Result<(), DriverError> { assert!( dst.len() >= self.len, @@ -595,6 +656,7 @@ impl DeviceBuffer { dst.len(), self.len ); + stream.context().bind_to_thread()?; unsafe { crate::simt::memory::memcpy_dtoh_async( dst.as_mut_ptr(), @@ -611,7 +673,8 @@ impl DeviceBuffer { /// /// Panics if `dst.len() < self.len()`. Use pinned destinations when you /// need the transfer to avoid pageable-memory staging; this helper still - /// waits for completion before returning, matching [`Self::copy_to_host`]. + /// waits for completion before returning, matching [`Self::copy_to_host`], + /// and orders a sticky error after that wait the same way. /// /// For true DtoH overlap, use [`Self::copy_to_pinned_host_async`] and /// synchronize the stream later. @@ -662,6 +725,7 @@ impl DeviceBuffer { dst.len(), self.len ); + stream.context().bind_to_thread()?; unsafe { crate::simt::memory::memcpy_dtoh_async( dst.as_mut_ptr(), @@ -714,6 +778,7 @@ impl DeviceBuffer { src.len(), self.len ); + stream.context().bind_to_thread()?; let num_bytes = src.num_bytes(); unsafe { crate::simt::memory::memcpy_htod_async( @@ -745,6 +810,7 @@ impl DeviceBuffer { len: usize, ) -> Result { let ctx = stream.context().clone(); + ctx.bind_to_thread()?; let num_bytes = allocation_size::(len)?; if num_bytes == 0 { // SAFETY: a null pointer with zero bytes is never dereferenced @@ -776,6 +842,7 @@ impl DeviceBuffer { "device-to-device copy length mismatch: dst {} != src {}", self.len, other.len ); + stream.context().bind_to_thread()?; if self.num_bytes() == 0 { return Ok(()); } @@ -794,18 +861,17 @@ impl DeviceBuffer { /// /// The synchronization keeps this safe for borrowed host slices: `src` /// may be dropped, reused, or mutated immediately after this function - /// returns. Panics if `src.len() != self.len()`. + /// returns. Panics if `src.len() != self.len()`. See [`Self::from_host`] + /// for how a sticky error is ordered against the wait. pub fn copy_from_host(&mut self, stream: &CudaStream, src: &[T]) -> Result<(), DriverError> { // SAFETY: this safe wrapper synchronizes `stream` before returning, // so the borrowed host slice cannot be used by CUDA after this call. - let enqueue_result = unsafe { self.copy_from_host_async_unchecked(stream, src) }; - let sync_result = if self.num_bytes() == 0 { - Ok(()) - } else { - stream.synchronize() - }; - enqueue_result?; - sync_result + // A failed enqueue leaves nothing in flight to wait for. + unsafe { self.copy_from_host_async_unchecked(stream, src)? }; + if self.num_bytes() == 0 { + return Ok(()); + } + stream.synchronize() } /// Copies `src` into `self` host-to-device, enqueued on `stream`, and @@ -831,6 +897,7 @@ impl DeviceBuffer { self.len, src.len() ); + stream.context().bind_to_thread()?; if self.num_bytes() == 0 { return Ok(()); } @@ -891,6 +958,7 @@ impl DeviceBuffer { /// Zeroes every byte in the buffer asynchronously on `stream`. pub fn zero_async(&mut self, stream: &CudaStream) -> Result<(), DriverError> { + stream.context().bind_to_thread()?; if self.num_bytes() == 0 { return Ok(()); } diff --git a/cuda-core/src/simt/embedded.rs b/cuda-core/src/simt/embedded.rs index c146f58753..fed3e373cc 100644 --- a/cuda-core/src/simt/embedded.rs +++ b/cuda-core/src/simt/embedded.rs @@ -4,6 +4,14 @@ */ //! Loading CUDA modules from embedded device artifact bundles. +//! +//! cuda-oxide's build backend links each compiled device module into the +//! host executable as an artifact bundle (a named blob carrying a cubin +//! and/or PTX payload plus the compile options that produced it; see the +//! `oxide-artifacts` crate, re-exported here as needed). This module reads +//! those bundles back out of the running executable, or out of any object +//! file on disk, and loads the ones that carry a loadable payload into a +//! [`CudaContext`]. use crate::{CudaContext, CudaModule, DriverError}; use oxide_artifacts::ArtifactError; @@ -15,34 +23,49 @@ use std::fmt; use std::path::{Path, PathBuf}; use std::sync::Arc; +/// An artifact bundle that carries a payload the driver can load: a cubin, +/// or failing that PTX for the driver to JIT. +/// +/// Construct with [`EmbeddedModule::new`], or collect every loadable bundle +/// linked into the running executable with +/// [`embedded_modules_from_current_exe`]. #[derive(Clone, Debug, PartialEq, Eq)] pub struct EmbeddedModule { bundle: OwnedArtifactBundle, } impl EmbeddedModule { + /// Wraps `bundle` if it carries a cubin or PTX payload; `None` for a + /// bundle with neither (metadata-only, or an unsupported payload kind). pub fn new(bundle: OwnedArtifactBundle) -> Option { loadable_payload(&bundle) .is_some() .then_some(Self { bundle }) } + /// The bundle's name, as assigned by the build that produced it; this is + /// what [`load_embedded_module`] matches on. pub fn name(&self) -> &str { &self.bundle.name } + /// The GPU target the payload was compiled for (e.g. `sm_90`). pub fn target(&self) -> &str { &self.bundle.target } + /// The underlying bundle, with its compile options and every payload. pub fn bundle(&self) -> &OwnedArtifactBundle { &self.bundle } + /// The bytes of the payload of `kind`, if the bundle carries one. pub fn payload(&self, kind: ArtifactPayloadKind) -> Option<&[u8]> { self.bundle.payload(kind) } + /// Loads the bundle's cubin (preferred) or PTX payload into `ctx` as a + /// [`CudaModule`]. pub fn load(&self, ctx: &Arc) -> Result, EmbeddedModuleError> { let image = loadable_payload(&self.bundle).expect("EmbeddedModule always has a loadable payload"); @@ -51,6 +74,8 @@ impl EmbeddedModule { } } +/// Reads every artifact bundle linked into the running executable, loadable +/// or not. pub fn artifact_bundles_from_current_exe() -> Result, EmbeddedModuleError> { let path = @@ -58,6 +83,8 @@ pub fn artifact_bundles_from_current_exe() -> Result, E artifact_bundles_from_binary_path(path) } +/// Reads every artifact bundle linked into the executable or object file at +/// `path`, loadable or not. pub fn artifact_bundles_from_binary_path( path: impl AsRef, ) -> Result, EmbeddedModuleError> { @@ -70,6 +97,8 @@ pub fn artifact_bundles_from_binary_path( .map_err(EmbeddedModuleError::Artifacts) } +/// The loadable bundles linked into the running executable, in link order. +/// Bundles without a cubin or PTX payload are skipped. pub fn embedded_modules_from_current_exe() -> Result, EmbeddedModuleError> { Ok(artifact_bundles_from_current_exe()? .into_iter() @@ -77,6 +106,10 @@ pub fn embedded_modules_from_current_exe() -> Result, Embedd .collect()) } +/// Loads the embedded module called `name` into `ctx`. +/// +/// Rereads the executable on every call; callers that load several modules +/// should collect [`embedded_modules_from_current_exe`] once instead. pub fn load_embedded_module( ctx: &Arc, name: &str, @@ -90,6 +123,8 @@ pub fn load_embedded_module( module.load(ctx) } +/// Loads the first loadable embedded module into `ctx`: the common case of an +/// executable built with a single device module. pub fn load_first_embedded_module( ctx: &Arc, ) -> Result, EmbeddedModuleError> { @@ -106,20 +141,31 @@ fn loadable_payload(bundle: &OwnedArtifactBundle) -> Option<&[u8]> { .or_else(|| bundle.payload(ArtifactPayloadKind::Ptx)) } +/// Why an embedded module could not be found or loaded. #[derive(Debug)] pub enum EmbeddedModuleError { + /// The path of the running executable could not be determined. CurrentExe { + /// The underlying I/O error. source: std::io::Error, }, + /// The executable or object file at `path` could not be read. Io { + /// The file that could not be read. path: PathBuf, + /// The underlying I/O error. source: std::io::Error, }, + /// The file was read but its artifact section could not be parsed. Artifacts(ArtifactError), + /// No loadable bundle in the executable has the requested name. ModuleNotFound { + /// The name that was requested. name: String, }, + /// The executable contains no loadable bundle at all. NoModules, + /// The driver rejected the payload (`cuModuleLoadData`). Driver(DriverError), } diff --git a/cuda-core/src/simt/event.rs b/cuda-core/src/simt/event.rs index 0ef390afa3..ddbd150e60 100644 --- a/cuda-core/src/simt/event.rs +++ b/cuda-core/src/simt/event.rs @@ -45,12 +45,12 @@ unsafe impl Sync for CudaEvent {} /// Destroys the underlying `CUevent` on drop. /// -/// Binds the context to the current thread first (required by -/// `cuEventDestroy`). Errors are recorded on the context rather than -/// panicking. +/// Makes the context current first without consulting the sticky error +/// state, so a recorded error cannot stop the destroy from being issued. +/// Errors are recorded on the context rather than panicking. impl Drop for CudaEvent { fn drop(&mut self) { - self.ctx.record_err(self.ctx.bind_to_thread()); + self.ctx.record_err(self.ctx.make_current()); self.ctx .record_err(unsafe { cuda_bindings::cuEventDestroy_v2(self.cu_event).result() }); } @@ -103,9 +103,15 @@ impl CudaEvent { /// Blocks the calling thread until this event has been recorded and all /// preceding stream work has completed. + /// + /// Like [`CudaStream::synchronize`], the driver wait is issued + /// unconditionally and the context's sticky error state is drained and + /// returned only afterwards, so an `Err` never means the wait was + /// skipped. pub fn synchronize(&self) -> Result<(), DriverError> { - self.ctx.bind_to_thread()?; - unsafe { cuda_bindings::cuEventSynchronize(self.cu_event).result() } + self.ctx.make_current()?; + unsafe { cuda_bindings::cuEventSynchronize(self.cu_event).result()? }; + self.ctx.check_err() } /// Returns `true` when all work captured by this event has completed, diff --git a/cuda-core/src/simt/mod.rs b/cuda-core/src/simt/mod.rs index 7f07b4c361..5511015528 100644 --- a/cuda-core/src/simt/mod.rs +++ b/cuda-core/src/simt/mod.rs @@ -15,8 +15,10 @@ //! Naming: everything cuda-oxide exposed at its crate root is re-exported //! here, so `cuda_core::simt::X` is oxide's `cuda_core::X`. The crate root //! additionally re-exports every name that does not collide with the -//! existing surface; the one collision is [`LaunchConfig`], which is -//! reachable only through this module. +//! existing surface. Two do collide and are reachable only through this +//! module: [`LaunchConfig`] (the crate has its own `LaunchConfig`) and +//! [`vmm`] (the crate's own reviewed `vmm` module forks the same cuda-oxide +//! ancestor). //! //! Differences from the cuda-oxide original, recorded for the //! reconciliation: @@ -53,9 +55,10 @@ pub use device_buffer::{DeviceBuffer, DeviceCopy}; pub use embedded::{EmbeddedModule, EmbeddedModuleError}; pub use event::CudaEvent; pub use launch::{ - BlockRequirement, DeviceLaunchLimits, DynamicSharedMemoryRequirement, KernelLaunchConfig, - KernelLaunchContract, LaunchAxis, LaunchConfig, LaunchConfig1D, LaunchConfig2D, LaunchConfig3D, - LaunchContractError, LaunchContractSpec, LaunchDimension, PreparedLaunch, + BlockRequirement, CoordinateRequirement, DeviceLaunchLimits, DynamicSharedMemoryRequirement, + KernelLaunchConfig, KernelLaunchContract, LaunchAxis, LaunchConfig, LaunchConfig1D, + LaunchConfig2D, LaunchConfig3D, LaunchContractError, LaunchContractSpec, LaunchDimension, + PreparedLaunch, }; pub use module::{ConstantHandle, CudaFunction, CudaModule}; pub use pinned_host_buffer::PinnedHostBuffer; @@ -152,9 +155,12 @@ pub unsafe fn launch_kernel_ex( stream: cuda_bindings::CUstream, kernel_params: &mut [*mut std::ffi::c_void], ) -> Result<(), DriverError> { - // CUlaunchAttribute_st is opaque (see cuda-bindings/build.rs) for CUDA 13.2+ - // compatibility. C layout: { id: u32 @ 0, pad: [u8;4] @ 4, value: union @ 8 }. - // clusterDim is three u32 fields (x, y, z) at offset 0 within the value union. + // `CUlaunchAttribute_st` is a plain struct in the generated bindings: + // `id` at offset 0, a 4-byte pad, then the 8-aligned 64-byte `value` + // union at offset 8. The attribute is written through those byte offsets + // rather than the union's members so this does not depend on bindgen's + // anonymous-member naming; `clusterDim.{x,y,z}` are the first three u32s + // of the union, at 8/12/16. let mut cluster_attr: cuda_bindings::CUlaunchAttribute_st = unsafe { std::mem::zeroed() }; unsafe { let base = &mut cluster_attr as *mut _ as *mut u8; @@ -276,10 +282,9 @@ pub unsafe fn launch_kernel_cooperative( stream: cuda_bindings::CUstream, kernel_params: &mut [*mut std::ffi::c_void], ) -> Result<(), DriverError> { - // CUlaunchAttribute_st is opaque (see cuda-bindings/build.rs) for CUDA 13.2+ - // compatibility. C layout: { id: u32 @ 0, pad: [u8;4] @ 4, value: union @ 8 }. - // For the COOPERATIVE attribute the value union holds a single `int cooperative` - // at offset 0 — set to 1 to enable, 0 to disable. + // Same layout as in `launch_kernel_ex` (`id` at 0, `value` union at 8); + // the COOPERATIVE attribute's payload is a single `int cooperative` at the + // start of the union: 1 enables, 0 disables. let mut coop_attr: cuda_bindings::CUlaunchAttribute_st = unsafe { std::mem::zeroed() }; unsafe { let base = &mut coop_attr as *mut _ as *mut u8; @@ -385,10 +390,9 @@ pub unsafe fn launch_kernel_ex_cooperative( stream: cuda_bindings::CUstream, kernel_params: &mut [*mut std::ffi::c_void], ) -> Result<(), DriverError> { - // CUlaunchAttribute_st is opaque (see cuda-bindings/build.rs) for CUDA 13.2+ - // compatibility. C layout: { id: u32 @ 0, pad: [u8;4] @ 4, value: union @ 8 }. - // attrs[0]: clusterDim — three u32 fields (x, y, z) at offset 0 of the union. - // attrs[1]: cooperative — a single `int` at offset 0 of the union; 1 = enabled. + // Same layout as in `launch_kernel_ex` (`id` at 0, `value` union at 8). + // attrs[0]: clusterDim, three u32s at the start of the union. + // attrs[1]: cooperative, a single `int` at the start of the union; 1 = enabled. let mut attrs: [cuda_bindings::CUlaunchAttribute_st; 2] = unsafe { std::mem::zeroed() }; unsafe { let base = &mut attrs[0] as *mut _ as *mut u8; diff --git a/cuda-core/src/simt/module.rs b/cuda-core/src/simt/module.rs index a2b4cb646f..e06aa4e985 100644 --- a/cuda-core/src/simt/module.rs +++ b/cuda-core/src/simt/module.rs @@ -56,12 +56,13 @@ unsafe impl Sync for CudaModule {} /// Unloads the module on drop. /// -/// Binds the context to the current thread first (required by -/// `cuModuleUnload`). Errors are recorded on the context rather than +/// Makes the context current first (required by `cuModuleUnload`) without +/// consulting the sticky error state, so a recorded error cannot stop the +/// unload from being issued. Errors are recorded on the context rather than /// panicking. impl Drop for CudaModule { fn drop(&mut self) { - self.ctx.record_err(self.ctx.bind_to_thread()); + self.ctx.record_err(self.ctx.make_current()); self.ctx .record_err(unsafe { cuda_bindings::cuModuleUnload(self.cu_module).result() }); } @@ -166,7 +167,6 @@ pub struct CudaFunction { /// Raw CUDA function handle. pub(crate) cu_function: cuda_bindings::CUfunction, /// Owning module. Prevents unloading while this function handle exists. - #[allow(unused)] pub(crate) module: Arc, } @@ -258,38 +258,123 @@ impl CudaModule { } } -/// A resolved handle to a `#[constant]` device global. Macro-generated -/// `set_` methods resolve these lazily on first use and cache the -/// handle on the `LoadedModule` struct. Callers pass `size_of::()` on -/// every write; correctness depends on the resolver asserting that the -/// driver-reported size matches the host-side type. -#[derive(Clone, Copy, Debug)] +/// A resolved handle to a module-scope device global (a `#[constant]` or +/// `__device__` symbol), as `cuModuleGetGlobal` reports it. +/// +/// The handle owns an `Arc` to its [`CudaModule`], so the address it holds +/// stays valid for as long as the handle exists, and it carries the +/// driver-reported byte size, so every write is bounds-checked against the +/// symbol instead of trusting the caller's length. Macro-generated +/// `set_` methods resolve these lazily on first use and cache them on +/// the `LoadedModule` struct; a resolver that also wants the symbol to match +/// a host type asserts `size() == size_of::()` once, at resolution. +/// +/// Cloning is an `Arc` bump. The handle is deliberately not `Copy`: a `Copy` +/// handle is a bare address that can outlive the module it points into. +/// +/// ```compile_fail,E0382 +/// # fn demo(handle: cuda_core::ConstantHandle) { +/// let first = handle; +/// let second = handle; +/// # let _ = (first, second); +/// # } +/// ``` +#[derive(Clone, Debug)] pub struct ConstantHandle { pub(crate) dptr: cuda_bindings::CUdeviceptr, + /// Size of the symbol in bytes, as reported by `cuModuleGetGlobal`. + pub(crate) size: usize, + /// The module that owns the symbol; keeps `dptr` valid. + pub(crate) module: Arc, } impl ConstantHandle { - /// Construct from a raw device pointer. Used by macro-generated - /// `LoadedModule` initializers after [`CudaModule::get_global`] has - /// resolved the symbol and the size has been asserted against - /// `size_of::()`. + /// Resolves the global `name` in `module`. + /// + /// Binds the module's context and calls [`CudaModule::get_global`]; the + /// returned handle records the symbol's address and size and keeps + /// `module` loaded. + /// + /// # Errors + /// + /// Whatever `get_global` returns: a bind failure, or `CUDA_ERROR_NOT_FOUND` + /// for a name the module does not export. + /// + /// # Panics + /// + /// Panics if `name` contains interior null bytes. + pub fn new(module: &Arc, name: &str) -> Result { + let (dptr, size) = module.get_global(name)?; + Ok(Self { + dptr, + size, + module: Arc::clone(module), + }) + } + + /// Builds a handle from an already-resolved address and size. + /// + /// [`new`](Self::new) is the ordinary constructor; this exists for callers + /// that resolved the symbol themselves through [`CudaModule::get_global`] + /// and want to avoid a second lookup. /// /// # Safety /// - /// `dptr` must point to at least `size_of::()` bytes of constant - /// memory in a still-loaded module. - pub unsafe fn from_raw(dptr: cuda_bindings::CUdeviceptr) -> Self { - Self { dptr } + /// - `dptr` and `size` must be exactly the pair `cuModuleGetGlobal` + /// returned for one symbol of `module`. Every bounds check this handle + /// performs trusts `size`, so an overstated size lets the *safe* + /// [`write_async_staged`](Self::write_async_staged) overrun the symbol, + /// and an address from another module, or from an earlier load of the + /// same image, is a wild device pointer. + /// - `module` must be the module that owns the symbol; holding it is what + /// keeps the address valid. + pub unsafe fn from_raw( + module: Arc, + dptr: cuda_bindings::CUdeviceptr, + size: usize, + ) -> Self { + Self { dptr, size, module } + } + + /// Size of the symbol in bytes, as reported by the driver. + pub fn size(&self) -> usize { + self.size + } + + /// The module that owns the symbol. + pub fn module(&self) -> &Arc { + &self.module + } + + /// Rejects a write of `num_bytes` that would run past the symbol, and a + /// stream from a context other than the module's, before any driver call. + fn check_write(&self, stream: &crate::CudaStream, num_bytes: usize) -> Result<(), DriverError> { + if num_bytes > self.size { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, + )); + } + if stream.context().as_ref() != self.module.context().as_ref() { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_CONTEXT, + )); + } + Ok(()) } -} -impl ConstantHandle { /// Stream-ordered `cuMemcpyHtoDAsync` from `src` (`num_bytes` of host /// memory) into the device global. /// + /// # Errors + /// + /// `CUDA_ERROR_INVALID_VALUE` if `num_bytes` exceeds [`size`](Self::size), + /// and `CUDA_ERROR_INVALID_CONTEXT` if `stream` belongs to a different + /// context than the module; both are checked before anything is enqueued. + /// /// # Safety /// - /// - `src` must point to at least `num_bytes` of readable host memory. + /// - `src` must point to at least `num_bytes` of readable host memory that + /// stays valid and unmodified until the copy completes. /// - The bytes must have a layout compatible with the device-side type. pub unsafe fn write_async( &self, @@ -297,6 +382,7 @@ impl ConstantHandle { src: *const u8, num_bytes: usize, ) -> Result<(), DriverError> { + self.check_write(stream, num_bytes)?; stream.context().bind_to_thread()?; unsafe { crate::simt::memory::memcpy_htod_async(self.dptr, src, num_bytes, stream.cu_stream()) @@ -309,13 +395,22 @@ impl ConstantHandle { /// The bytes are kept alive until the stream reaches a host callback /// enqueued after the copy. This makes safe setters sound even when the /// caller passes a temporary such as `&3.0`. + /// + /// # Errors + /// + /// `CUDA_ERROR_INVALID_VALUE` if `bytes` is longer than the symbol + /// ([`size`](Self::size)), and `CUDA_ERROR_INVALID_CONTEXT` if `stream` + /// belongs to a different context than the module; both are checked + /// before anything is enqueued, so a refused write leaves the symbol + /// untouched. pub fn write_async_staged( &self, stream: &crate::CudaStream, bytes: Box<[MaybeUninit]>, ) -> Result<(), DriverError> { - stream.context().bind_to_thread()?; let num_bytes = bytes.len(); + self.check_write(stream, num_bytes)?; + stream.context().bind_to_thread()?; if num_bytes == 0 { return Ok(()); } @@ -351,7 +446,12 @@ impl ConstantHandle { callback_data as *mut Box<[MaybeUninit]>, ) }; - if let Err(sync_err) = stream.synchronize() { + // The copy is enqueued but nothing will free the bytes after it: + // wait for it here. Only a failed driver wait means the DMA may + // still be reading `staged`, in which case leaking is the only + // sound option; a sticky error recorded meanwhile does not make + // the completed wait any less complete. + if let Err(sync_err) = stream.synchronize_driver() { Box::leak(staged); return Err(sync_err); } @@ -365,16 +465,32 @@ impl ConstantHandle { /// Synchronous `cuMemcpyHtoD` from `src` into the device global. Blocks /// the calling thread. /// + /// The module is the one this handle was resolved from; it no longer + /// needs to be passed in. + /// + /// # Errors + /// + /// `CUDA_ERROR_INVALID_VALUE` if `num_bytes` exceeds [`size`](Self::size), + /// checked before any driver call. + /// /// # Safety /// - /// Same contract as [`write_async`](Self::write_async). + /// Same contract as [`write_async`](Self::write_async), except that `src` + /// only needs to stay valid for the duration of the call. pub unsafe fn write_blocking( &self, - module: &Arc, src: *const u8, num_bytes: usize, ) -> Result<(), DriverError> { - unsafe { module.copy_bytes_to_device_global_sync(self.dptr, src, num_bytes) } + if num_bytes > self.size { + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, + )); + } + unsafe { + self.module + .copy_bytes_to_device_global_sync(self.dptr, src, num_bytes) + } } } @@ -640,10 +756,11 @@ impl CudaFunction { ) -> Result { self.context().bind_to_thread()?; - // CUlaunchAttribute_st is opaque in the generated CUDA 13.2+ - // bindings. Its C layout stores the id at offset 0 and the value union - // at offset 8; clusterDim.x/y/z occupy the first three u32 values in - // that union. This matches the launch helpers in cuda-core's root. + // `CUlaunchAttribute_st` is a plain struct in the generated bindings + // (`id` at offset 0, a 4-byte pad, the 8-aligned `value` union at 8). + // Written through byte offsets, as the launch helpers in `simt` do, + // so this does not depend on bindgen's anonymous-member naming; + // `clusterDim.{x,y,z}` are the first three u32s of the union. let mut cluster_attribute: cuda_bindings::CUlaunchAttribute_st = unsafe { std::mem::zeroed() }; unsafe { diff --git a/cuda-core/src/simt/pinned_host_buffer.rs b/cuda-core/src/simt/pinned_host_buffer.rs index f39d823d82..cf29b9fc22 100644 --- a/cuda-core/src/simt/pinned_host_buffer.rs +++ b/cuda-core/src/simt/pinned_host_buffer.rs @@ -129,9 +129,21 @@ impl PinnedHostBuffer { } else { ctx.bind_to_thread()?; let ptr = unsafe { crate::simt::memory::malloc_host(num_bytes)? }; - NonNull::new(ptr.cast::()).ok_or(DriverError( - cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, - ))? + // `cuMemAllocHost` returns page-aligned memory in practice, but + // the slices handed out below assume `align_of::()`, so check + // rather than trust: a misaligned `&[T]` is undefined behavior. + if ptr.is_null() || !(ptr as usize).is_multiple_of(std::mem::align_of::()) { + if !ptr.is_null() { + // SAFETY: `ptr` came from `malloc_host` above and has not + // been handed out. + ctx.record_err(unsafe { crate::simt::memory::free_host(ptr) }); + } + return Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, + )); + } + // SAFETY: checked non-null just above. + unsafe { NonNull::new_unchecked(ptr.cast::()) } }; Ok(Self { @@ -147,7 +159,8 @@ impl PinnedHostBuffer { impl Drop for PinnedHostBuffer { fn drop(&mut self) { if self.num_bytes != 0 { - self.ctx.record_err(self.ctx.bind_to_thread()); + // Not `bind_to_thread`: a recorded error must not skip the free. + self.ctx.record_err(self.ctx.make_current()); self.ctx .record_err(unsafe { crate::simt::memory::free_host(self.ptr.as_ptr().cast()) }); } diff --git a/cuda-core/src/simt/stream.rs b/cuda-core/src/simt/stream.rs index 8a97417cc4..5e60a933a7 100644 --- a/cuda-core/src/simt/stream.rs +++ b/cuda-core/src/simt/stream.rs @@ -57,10 +57,11 @@ unsafe impl Sync for CudaStream {} /// live stream count. /// /// The default stream (null handle) is never destroyed. Errors during -/// teardown are recorded on the context rather than panicking. +/// teardown are recorded on the context rather than panicking, and a +/// previously recorded error does not stop the destroy from being issued. impl Drop for CudaStream { fn drop(&mut self) { - self.ctx.record_err(self.ctx.bind_to_thread()); + self.ctx.record_err(self.ctx.make_current()); if !self.cu_stream.is_null() { self.ctx.num_streams.fetch_sub(1, Ordering::Relaxed); self.ctx @@ -99,8 +100,26 @@ impl CudaStream { /// Blocks the calling thread until all work enqueued on this stream /// completes. + /// + /// `cuStreamSynchronize` is issued unconditionally; the context's sticky + /// error state is drained and returned only afterwards. An `Err` from + /// this method therefore never means the wait was skipped: either the + /// driver failed the wait, or the wait completed and a previously + /// recorded error is being reported. The `DeviceBuffer` copy wrappers + /// depend on that ordering to hand host memory back only after the DMA + /// has finished. pub fn synchronize(&self) -> Result<(), DriverError> { - self.ctx.bind_to_thread()?; + self.synchronize_driver()?; + self.ctx.check_err() + } + + /// `cuStreamSynchronize` without consulting the sticky error state. + /// + /// `Err` means the driver refused or failed the wait, so an enqueued + /// transfer may still be in flight; callers that own the host buffer + /// leak it in that case rather than free memory the device may write. + pub(crate) fn synchronize_driver(&self) -> Result<(), DriverError> { + self.ctx.make_current()?; unsafe { cuda_bindings::cuStreamSynchronize(self.cu_stream) }.result() } @@ -142,7 +161,6 @@ impl CudaStream { /// in the stream DAG. pub fn fork(&self) -> Result, DriverError> { self.ctx.bind_to_thread()?; - self.ctx.num_streams.fetch_add(1, Ordering::Relaxed); let mut cu_stream = MaybeUninit::uninit(); let cu_stream = unsafe { cuda_bindings::cuStreamCreate( @@ -152,6 +170,9 @@ impl CudaStream { .result()?; cu_stream.assume_init() }; + // Count the stream only once it exists: a failed create has no + // `CudaStream` whose drop would undo the increment. + self.ctx.num_streams.fetch_add(1, Ordering::Relaxed); let stream = Arc::new(CudaStream { cu_stream, ctx: self.ctx.clone(), @@ -208,26 +229,65 @@ impl CudaStream { /// `oneshot::Sender::send` or `Waker::wake` in `host_func` to unblock a /// future when GPU work finishes. /// - /// `host_func` is boxed, leaked into a raw pointer, and passed as user - /// data to `cuLaunchHostFunc`. The driver calls - /// `callback_wrapper` on a driver-internal - /// thread, which reclaims the box and invokes the closure. + /// `host_func` is boxed and handed to `cuLaunchHostFunc` as user data. + /// The driver calls a private `extern "C"` trampoline on a driver-internal + /// thread, which reclaims the box and invokes the closure. If the enqueue + /// itself fails the box is reclaimed here and the closure is dropped + /// without running. /// - /// Panics inside the closure are caught and discarded to prevent unwinding - /// across the FFI boundary. - pub fn launch_host_function( + /// # What the callback may do + /// + /// - It runs at an unknown later time on a thread the driver owns, so + /// `F: 'static`: a closure borrowing a local of the enclosing scope is + /// rejected at compile time. + /// + /// ```compile_fail,E0373 + /// # fn demo(stream: &cuda_core::CudaStream) { + /// let hits = std::sync::atomic::AtomicUsize::new(0); + /// let _ = stream.launch_host_function(|| { + /// hits.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + /// }); + /// # } + /// ``` + /// + /// - It must not call CUDA driver APIs. The driver documents this for + /// `cuLaunchHostFunc` (the call may fail with + /// `CUDA_ERROR_NOT_PERMITTED`, or deadlock). Publishing a result to the + /// host (an atomic store, a channel send, a `Waker::wake`) is what the + /// hook is for; [`CudaContext::record_err`] is a plain atomic store and + /// is fine. + /// - It must not block waiting on a thread that may itself be inside the + /// driver: host callbacks run serially on the driver's thread and hold + /// up every later callback and every context-wide wait in the process + /// while they run. + /// - Panics inside the closure are caught and **discarded** so the unwind + /// never crosses the C ABI boundary; a panicking callback reports + /// nothing. Record failures somewhere the host can observe them. + /// - Unlike the legacy `cuStreamAddCallback`, the driver does not invoke + /// the function once the context is in an error state, in which case + /// the boxed closure is leaked. + pub fn launch_host_function( &self, host_func: F, ) -> Result<(), DriverError> { - let boxed = Box::new(host_func); - unsafe { + self.ctx.bind_to_thread()?; + let boxed = Box::into_raw(Box::new(host_func)); + let result = unsafe { cuda_bindings::cuLaunchHostFunc( self.cu_stream, Some(Self::callback_wrapper::), - Box::into_raw(boxed) as *mut c_void, + boxed as *mut c_void, ) .result() + }; + if result.is_err() { + // The driver did not take the callback and will never invoke + // `callback_wrapper` for this pointer. + // SAFETY: `boxed` came from `Box::into_raw` above and nothing + // else has been given the chance to free it. + drop(unsafe { Box::from_raw(boxed) }); } + result } /// `extern "C"` trampoline invoked by the CUDA driver on a driver-internal @@ -243,7 +303,7 @@ impl CudaStream { /// where `f: F`. /// - Must be called exactly once per enqueued callback (double-free /// otherwise). - unsafe extern "C" fn callback_wrapper(callback: *mut c_void) { + unsafe extern "C" fn callback_wrapper(callback: *mut c_void) { let _ = std::panic::catch_unwind(|| { let callback: Box = unsafe { Box::from_raw(callback as *mut F) }; callback(); diff --git a/cuda-core/src/simt/vmm.rs b/cuda-core/src/simt/vmm.rs index 4784fc6099..1393103192 100644 --- a/cuda-core/src/simt/vmm.rs +++ b/cuda-core/src/simt/vmm.rs @@ -13,15 +13,39 @@ //! //! All handle types are RAII: `PhysicalAllocation` releases via `cuMemRelease`, //! `VirtualReservation` frees via `cuMemAddressFree`, and `Mapping` unmaps -//! via `cuMemUnmap`. Drop order matters -- mappings must be dropped before the -//! physical allocation or virtual reservation they reference. +//! via `cuMemUnmap`. A `Mapping<'a>` borrows both the reservation it lives in +//! and the memory it maps, so the teardown order (mapping first, then +//! reservation and allocation) is checked by the compiler rather than left +//! to convention. +//! +//! Unmapping is **not stream-ordered**: `cuMemUnmap` acts on the host +//! timeline without waiting for enqueued work, so a mapping must not be +//! dropped while a kernel or copy that touches its range is in flight (see +//! [`Mapping`]). //! //! The module also wraps CUDA's multicast objects (`cuMulticast*`, CUDA //! 12.1+): `MulticastObject` (compiled out on pre-12.1 toolkits, see the //! build script probe) builds an NVLink SHARP (NVLS) team whose //! mapped VA ranges respond to device-side `multimem.*` instructions with -//! switch-side reduction and broadcast across every bound GPU. See -//! `tests/vmm_multicast.rs` and the `nvls_all_reduce` example. +//! switch-side reduction and broadcast across every bound GPU. The host-side +//! plumbing is exercised by `tests/simt_vmm_multicast.rs`; the switch-side +//! semantics need device code and are not tested in this crate. +//! +//! # Contexts +//! +//! Nothing in this module binds a [`CudaContext`](crate::CudaContext): every +//! entry point is addressed by `CUdevice` and operates on the process-wide +//! VA space, and none of them holds a context to bind. On CUDA 13.3 (driver +//! 610.43) all of them, `set_access` included, succeed on a thread with no +//! current context and even in a process that never created one, so a +//! `CUDA_ERROR_INVALID_VALUE` from `set_access` is about its arguments (a +//! `va`/`size` that is not an exact mapped range, or a device outside the +//! system), not about the calling thread's binding. The work that *uses* a +//! mapping does run in a context: the kernels and copies that touch the VA +//! run in the context of the device they were issued on, which must be in +//! the `set_access` device list. `MulticastObject::bind_mem` is the one call +//! CUDA documents as needing a context for the bound device current; bind +//! one with `CudaContext::bind_to_thread` first. use crate::error::{DriverError, IntoResult}; use cuda_bindings::CUdeviceptr; @@ -152,34 +176,110 @@ impl Drop for VirtualReservation { /// A mapping of physical memory into a virtual address range. /// -/// Created by [`Mapping::new`], which calls `cuMemMap` to bind a -/// `PhysicalAllocation` (or a portion of it) to a region within a -/// `VirtualReservation`. Dropped via `cuMemUnmap`. -pub struct Mapping { +/// Created by [`Mapping::new`] (`cuMemMap` of a [`PhysicalAllocation`]) or +/// `Mapping::new_multicast` (of a `MulticastObject`, CUDA 12.1+); dropped +/// via `cuMemUnmap`. The mapping borrows the [`VirtualReservation`] it lives +/// in and the memory it maps for its whole lifetime, so it cannot outlive +/// either. Freeing a reservation that is still mapped is exactly the teardown +/// mistake `Drop` cannot report, and the borrow turns it into a compile +/// error: +/// +/// ```compile_fail,E0505 +/// # use cuda_core::simt::vmm::{Mapping, PhysicalAllocation, VirtualReservation}; +/// # fn demo(va: VirtualReservation, phys: PhysicalAllocation) -> Result<(), cuda_core::DriverError> { +/// let mapping = Mapping::new(&va, 0, &phys, 0, va.size())?; +/// drop(va); // the range is still mapped +/// drop(mapping); +/// # Ok(()) +/// # } +/// ``` +/// +/// # Unmapping is not stream-ordered +/// +/// `cuMemUnmap` takes effect on the host timeline; it does not wait for work +/// already enqueued on any stream. Dropping a `Mapping` while a kernel or a +/// copy that touches its range is still in flight is a device-side +/// use-after-unmap: the access faults (`CUDA_ERROR_ILLEGAL_ADDRESS`) or, if +/// the range has been remapped meanwhile, lands in someone else's memory. +/// Synchronize every stream that uses the range before the mapping drops; +/// nothing here does it for you. +pub struct Mapping<'a> { va: CUdeviceptr, size: usize, + /// The range `va` lies in; borrowed so the mapping is unmapped first. + reservation: &'a VirtualReservation, + /// What `va` is backed by; borrowed so it outlives the mapping. + backing: MappingBacking<'a>, } -impl Mapping { - /// Maps `size` bytes of `phys` at `offset` into virtual address `va`. +/// The memory a [`Mapping`] maps. +enum MappingBacking<'a> { + Physical(&'a PhysicalAllocation), + #[cfg(cuda_has_multicast)] + Multicast(&'a MulticastObject), +} + +/// Bounds a mapping of `size` bytes at `va_offset` into `reservation` and at +/// `backing_offset` into a backing of `backing_size` bytes, returning the +/// mapped address. +/// +/// The driver cannot check the reservation bound: VA ranges are process-wide, +/// so a `cuMemMap` past the end of this reservation would succeed if it +/// happened to land in another one. +fn mapping_range( + reservation: &VirtualReservation, + va_offset: usize, + backing_size: usize, + backing_offset: usize, + size: usize, +) -> Result { + let invalid = || DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE); + let va_end = va_offset.checked_add(size).ok_or_else(invalid)?; + let backing_end = backing_offset.checked_add(size).ok_or_else(invalid)?; + if va_end > reservation.size() || backing_end > backing_size { + return Err(invalid()); + } + let va_offset = CUdeviceptr::try_from(va_offset).map_err(|_| invalid())?; + reservation + .base() + .checked_add(va_offset) + .ok_or_else(invalid) +} + +impl<'a> Mapping<'a> { + /// Maps `size` bytes of `phys`, starting `phys_offset` bytes into it, at + /// `va_offset` bytes into `reservation`. + /// + /// Both offsets and `size` must be multiples of the allocation + /// granularity; the mapping is not accessible until [`set_access`] grants + /// a device permission on it. + /// + /// # Errors /// - /// `va` must lie within a `VirtualReservation`. `offset` is the byte - /// offset into the physical allocation (typically 0 for full mappings). - /// `size` must be a multiple of the allocation granularity. + /// `CUDA_ERROR_INVALID_VALUE` when the range does not fit in + /// `reservation` or in `phys` (checked before the driver call), and + /// whatever `cuMemMap` reports otherwise. pub fn new( - va: CUdeviceptr, + reservation: &'a VirtualReservation, + va_offset: usize, + phys: &'a PhysicalAllocation, + phys_offset: usize, size: usize, - phys: &PhysicalAllocation, - offset: usize, ) -> Result { + let va = mapping_range(reservation, va_offset, phys.size(), phys_offset, size)?; unsafe { - cuda_bindings::cuMemMap(va, size, offset, phys.handle(), 0).result()?; + cuda_bindings::cuMemMap(va, size, phys_offset, phys.handle(), 0).result()?; } - Ok(Self { va, size }) + Ok(Self { + va, + size, + reservation, + backing: MappingBacking::Physical(phys), + }) } - /// Maps `size` bytes of a multicast object at `offset` into virtual - /// address `va`. + /// Maps `size` bytes of `multicast`, starting `mc_offset` bytes into it, + /// at `va_offset` bytes into `reservation`. /// /// The resulting VA is a *multicast* view: `multimem.*` PTX instructions /// issued against it operate on every copy bound to the object (see @@ -188,17 +288,28 @@ impl Mapping { /// /// All devices must have been added to the multicast object (via /// [`MulticastObject::add_device`]) before mapping it. + /// + /// # Errors + /// + /// As [`Mapping::new`], with the object's per-device size as the bound. #[cfg(cuda_has_multicast)] pub fn new_multicast( - va: CUdeviceptr, + reservation: &'a VirtualReservation, + va_offset: usize, + multicast: &'a MulticastObject, + mc_offset: usize, size: usize, - multicast: &MulticastObject, - offset: usize, ) -> Result { + let va = mapping_range(reservation, va_offset, multicast.size(), mc_offset, size)?; unsafe { - cuda_bindings::cuMemMap(va, size, offset, multicast.handle(), 0).result()?; + cuda_bindings::cuMemMap(va, size, mc_offset, multicast.handle(), 0).result()?; } - Ok(Self { va, size }) + Ok(Self { + va, + size, + reservation, + backing: MappingBacking::Multicast(multicast), + }) } /// Returns the virtual address this mapping occupies. @@ -210,9 +321,36 @@ impl Mapping { pub fn size(&self) -> usize { self.size } + + /// Returns the reservation this mapping lies in. + pub fn reservation(&self) -> &'a VirtualReservation { + self.reservation + } + + /// Returns the physical allocation this mapping maps, or `None` for a + /// multicast view. + pub fn physical(&self) -> Option<&'a PhysicalAllocation> { + match self.backing { + MappingBacking::Physical(phys) => Some(phys), + #[cfg(cuda_has_multicast)] + MappingBacking::Multicast(_) => None, + } + } + + /// Returns the multicast object this mapping is a view of, or `None` for + /// a mapping of physical memory. + #[cfg(cuda_has_multicast)] + pub fn multicast(&self) -> Option<&'a MulticastObject> { + match self.backing { + MappingBacking::Multicast(multicast) => Some(multicast), + MappingBacking::Physical(_) => None, + } + } } -impl Drop for Mapping { +impl Drop for Mapping<'_> { + /// `cuMemUnmap`, immediately; see the type docs for why in-flight work on + /// the range must be complete first. fn drop(&mut self) { unsafe { let _ = cuda_bindings::cuMemUnmap(self.va, self.size).result(); @@ -227,6 +365,13 @@ impl Drop for Mapping { /// `[va, va + size)`. /// /// Typically called once after all mappings within a reservation are established. +/// +/// # Errors +/// +/// `CUDA_ERROR_INVALID_VALUE` means the arguments: `[va, va + size)` must be +/// exactly one or more whole mappings (granularity-aligned, fully mapped), +/// and every device must exist. It is not a context error; see the module +/// docs on contexts. pub fn set_access( va: CUdeviceptr, size: usize, @@ -267,8 +412,56 @@ pub fn allocation_granularity(device: cuda_bindings::CUdevice) -> Result usize { - (size + granularity - 1) & !(granularity - 1) + assert!(granularity != 0, "granularity must be nonzero"); + let remainder = size % granularity; + if remainder == 0 { + size + } else { + size.checked_add(granularity - remainder) + .expect("aligned size overflows usize") + } +} + +#[cfg(test)] +mod align_size_tests { + use super::align_size; + + #[test] + fn handles_power_of_two_granularities() { + assert_eq!(align_size(0, 2 << 20), 0); + assert_eq!(align_size(1, 2 << 20), 2 << 20); + assert_eq!(align_size(2 << 20, 2 << 20), 2 << 20); + assert_eq!(align_size((2 << 20) + 1, 2 << 20), 4 << 20); + } + + /// The previous mask-based version, `(size + g - 1) & !(g - 1)`, assumed + /// a power of two: for `g = 6` it maps 7 to `12 & !5 = 8`, which is not + /// a multiple of 6 at all. + #[test] + fn handles_general_granularities() { + assert_eq!(align_size(0, 6), 0); + assert_eq!(align_size(12, 6), 12); + assert_eq!(align_size(13, 6), 18); + assert_eq!(align_size(7, 6), 12); + } + + #[test] + #[should_panic(expected = "granularity must be nonzero")] + fn rejects_zero_granularity() { + align_size(1, 0); + } + + #[test] + #[should_panic(expected = "aligned size overflows usize")] + fn rejects_overflow() { + align_size(usize::MAX, 2); + } } /// Granularity flavor for [`multicast_granularity`]. @@ -381,7 +574,10 @@ pub fn multicast_granularity( /// 4. [`Mapping::new_multicast`] + [`set_access`] per device. /// /// Teardown reverses it: mappings first, then bindings ([`MulticastBinding`]), -/// then this object, then the physical allocations. +/// then this object, then the physical allocations. Mappings and bindings +/// borrow this object, so the compiler enforces that they go first; the +/// physical allocations are not borrowed by the bindings (CUDA keeps bound +/// memory alive), so releasing them last stays a convention. /// /// Dropping releases the handle via `cuMemRelease` (the documented release /// path for multicast objects). @@ -429,15 +625,15 @@ impl MulticastObject { /// and every team device must already have been added. A CUDA context /// for the bound device must be current. /// - /// The returned [`MulticastBinding`] unbinds on drop; it must be dropped - /// before this object and before `phys`. + /// The returned [`MulticastBinding`] unbinds on drop. It borrows this + /// object, so it cannot outlive it; drop it before `phys` as well. pub fn bind_mem( &self, mc_offset: usize, phys: &PhysicalAllocation, mem_offset: usize, size: usize, - ) -> Result { + ) -> Result, DriverError> { unsafe { cuda_bindings::cuMulticastBindMem( self.handle, @@ -450,7 +646,7 @@ impl MulticastObject { .result()?; } Ok(MulticastBinding { - mc_handle: self.handle, + multicast: self, device: phys.device(), mc_offset, size, @@ -485,30 +681,38 @@ impl Drop for MulticastObject { /// One device's physical memory bound into a [`MulticastObject`]. /// /// Created by [`MulticastObject::bind_mem`]; unbinds via `cuMulticastUnbind` -/// on drop. Must be dropped before the multicast object and before the -/// physical allocation it binds. +/// on drop. It borrows the multicast object, so it cannot outlive it (an +/// unbind against a released object handle would be a wild handle); drop it +/// before the physical allocation it binds as well. #[cfg(cuda_has_multicast)] -pub struct MulticastBinding { - mc_handle: cuda_bindings::CUmemGenericAllocationHandle, +pub struct MulticastBinding<'a> { + /// The object the memory is bound into; borrowed so the unbind runs + /// while the handle is still valid. + multicast: &'a MulticastObject, device: cuda_bindings::CUdevice, mc_offset: usize, size: usize, } #[cfg(cuda_has_multicast)] -impl MulticastBinding { +impl<'a> MulticastBinding<'a> { /// Returns the device whose memory is bound. pub fn device(&self) -> cuda_bindings::CUdevice { self.device } + + /// Returns the multicast object the memory is bound into. + pub fn multicast(&self) -> &'a MulticastObject { + self.multicast + } } #[cfg(cuda_has_multicast)] -impl Drop for MulticastBinding { +impl Drop for MulticastBinding<'_> { fn drop(&mut self) { unsafe { let _ = cuda_bindings::cuMulticastUnbind( - self.mc_handle, + self.multicast.handle(), self.device, self.mc_offset, self.size, diff --git a/cuda-core/src/vmm.rs b/cuda-core/src/vmm.rs index 3c8ffb648f..400ee7fe31 100644 --- a/cuda-core/src/vmm.rs +++ b/cuda-core/src/vmm.rs @@ -13,17 +13,25 @@ //! //! All handle types are RAII: `PhysicalAllocation` releases its handle via //! `cuMemRelease`, `VirtualReservation` frees via `cuMemAddressFree`, and -//! `Mapping` unmaps via `cuMemUnmap`. CUDA keeps a physical allocation alive -//! while mappings still reference it, so releasing its handle first does not -//! invalidate those mappings. Mappings must still be dropped before their -//! virtual reservation for deterministic, leak-free cleanup because `Drop` -//! cannot report a failed `cuMemAddressFree` call. +//! `Mapping` unmaps via `cuMemUnmap`. A `Mapping<'a>` borrows both the +//! reservation it lives in and the memory it maps, so the teardown order +//! (mapping first, then reservation and allocation) is checked by the +//! compiler: a reservation cannot be freed while a mapping still points into +//! it, which `Drop` could not report. CUDA does keep a physical allocation +//! alive while mappings reference it, so the borrow of the allocation is +//! about determinism rather than a driver requirement. +//! +//! Unmapping is **not stream-ordered**: `cuMemUnmap` acts on the host +//! timeline without waiting for enqueued work, so a mapping must not be +//! dropped while a kernel or copy that touches its range is in flight (see +//! [`Mapping`]). //! //! The module also wraps CUDA's multicast objects (`cuMulticast*`, CUDA -//! 12.1+): `MulticastObject` (builds an NVLink SHARP (NVLS) team whose +//! 12.1+): `MulticastObject` builds an NVLink SHARP (NVLS) team whose //! mapped VA ranges respond to device-side `multimem.*` instructions with -//! switch-side reduction and broadcast across every bound GPU. See -//! `tests/vmm_multicast.rs`. +//! switch-side reduction and broadcast across every bound GPU. The host-side +//! plumbing is exercised by `tests/vmm_multicast.rs`; the switch-side +//! semantics need device code and are not tested in this crate. use crate::error::{DriverError, IntoResult}; use cuda_bindings::CUdeviceptr; @@ -152,37 +160,111 @@ impl Drop for VirtualReservation { /// A mapping of physical memory into a virtual address range. /// -/// Created by [`Mapping::new`], which calls `cuMemMap` to bind a -/// `PhysicalAllocation` (or a portion of it) to a region within a -/// `VirtualReservation`. Dropped via `cuMemUnmap`. -pub struct Mapping { +/// Created by [`Mapping::new`] (`cuMemMap` of a [`PhysicalAllocation`]) or +/// [`Mapping::new_multicast`] (of a [`MulticastObject`]); dropped via +/// `cuMemUnmap`. The mapping borrows the [`VirtualReservation`] it lives in +/// and the memory it maps for its whole lifetime, so it cannot outlive either. +/// Freeing a reservation that is still mapped is exactly the teardown mistake +/// `Drop` cannot report, and the borrow turns it into a compile error: +/// +/// ```compile_fail,E0505 +/// # use cuda_core::vmm::{Mapping, PhysicalAllocation, VirtualReservation}; +/// # fn demo(va: VirtualReservation, phys: PhysicalAllocation) -> Result<(), cuda_core::DriverError> { +/// let mapping = Mapping::new(&va, 0, &phys, 0, va.size())?; +/// drop(va); // the range is still mapped +/// drop(mapping); +/// # Ok(()) +/// # } +/// ``` +/// +/// # Unmapping is not stream-ordered +/// +/// `cuMemUnmap` takes effect on the host timeline; it does not wait for work +/// already enqueued on any stream. Dropping a `Mapping` while a kernel or a +/// copy that touches its range is still in flight is a device-side +/// use-after-unmap: the access faults (`CUDA_ERROR_ILLEGAL_ADDRESS`) or, if +/// the range has been remapped meanwhile, lands in someone else's memory. +/// Synchronize every stream that uses the range before the mapping drops; +/// nothing here does it for you. +pub struct Mapping<'a> { va: CUdeviceptr, size: usize, + /// The range `va` lies in; borrowed so the mapping is unmapped first. + reservation: &'a VirtualReservation, + /// What `va` is backed by; borrowed so it outlives the mapping. + backing: MappingBacking<'a>, +} + +/// The memory a [`Mapping`] maps. +enum MappingBacking<'a> { + Physical(&'a PhysicalAllocation), + Multicast(&'a MulticastObject), +} + +/// Bounds a mapping of `size` bytes at `va_offset` into `reservation` and at +/// `backing_offset` into a backing of `backing_size` bytes, returning the +/// mapped address. +/// +/// The driver cannot check the reservation bound: VA ranges are process-wide, +/// so a `cuMemMap` past the end of this reservation would succeed if it +/// happened to land in another one. +fn mapping_range( + reservation: &VirtualReservation, + va_offset: usize, + backing_size: usize, + backing_offset: usize, + size: usize, +) -> Result { + let invalid = || DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE); + let va_end = va_offset.checked_add(size).ok_or_else(invalid)?; + let backing_end = backing_offset.checked_add(size).ok_or_else(invalid)?; + if va_end > reservation.size() || backing_end > backing_size { + return Err(invalid()); + } + let va_offset = CUdeviceptr::try_from(va_offset).map_err(|_| invalid())?; + reservation + .base() + .checked_add(va_offset) + .ok_or_else(invalid) } -impl Mapping { - /// Maps `size` bytes of `phys` at `offset` into virtual address `va`. +impl<'a> Mapping<'a> { + /// Maps `size` bytes of `phys`, starting `phys_offset` bytes into it, at + /// `va_offset` bytes into `reservation`. + /// + /// Both offsets and `size` must be multiples of the allocation + /// granularity; the mapping is not accessible until [`set_access`] grants + /// a device permission on it. /// - /// `va` must lie within a `VirtualReservation`. `offset` is the byte - /// offset into the physical allocation (typically 0 for full mappings). - /// `size` must be a multiple of the allocation granularity. + /// CUDA retains the physical memory while the mapping exists, so the + /// borrow of `phys` is about deterministic teardown, not validity. /// - /// CUDA retains the physical allocation while the mapping exists, so the - /// mapping remains valid if `phys` releases its handle first. + /// # Errors + /// + /// `CUDA_ERROR_INVALID_VALUE` when the range does not fit in + /// `reservation` or in `phys` (checked before the driver call), and + /// whatever `cuMemMap` reports otherwise. pub fn new( - va: CUdeviceptr, + reservation: &'a VirtualReservation, + va_offset: usize, + phys: &'a PhysicalAllocation, + phys_offset: usize, size: usize, - phys: &PhysicalAllocation, - offset: usize, ) -> Result { + let va = mapping_range(reservation, va_offset, phys.size(), phys_offset, size)?; unsafe { - cuda_bindings::cuMemMap(va, size, offset, phys.handle(), 0).result()?; + cuda_bindings::cuMemMap(va, size, phys_offset, phys.handle(), 0).result()?; } - Ok(Self { va, size }) + Ok(Self { + va, + size, + reservation, + backing: MappingBacking::Physical(phys), + }) } - /// Maps `size` bytes of a multicast object at `offset` into virtual - /// address `va`. + /// Maps `size` bytes of `multicast`, starting `mc_offset` bytes into it, + /// at `va_offset` bytes into `reservation`. /// /// The resulting VA is a *multicast* view: `multimem.*` PTX instructions /// issued against it operate on every copy bound to the object (see @@ -191,17 +273,27 @@ impl Mapping { /// /// All devices must have been added to the multicast object (via /// [`MulticastObject::add_device`]) before mapping it. - /// CUDA retains the multicast allocation while the mapping exists. + /// + /// # Errors + /// + /// As [`Mapping::new`], with the object's per-device size as the bound. pub fn new_multicast( - va: CUdeviceptr, + reservation: &'a VirtualReservation, + va_offset: usize, + multicast: &'a MulticastObject, + mc_offset: usize, size: usize, - multicast: &MulticastObject, - offset: usize, ) -> Result { + let va = mapping_range(reservation, va_offset, multicast.size(), mc_offset, size)?; unsafe { - cuda_bindings::cuMemMap(va, size, offset, multicast.handle(), 0).result()?; + cuda_bindings::cuMemMap(va, size, mc_offset, multicast.handle(), 0).result()?; } - Ok(Self { va, size }) + Ok(Self { + va, + size, + reservation, + backing: MappingBacking::Multicast(multicast), + }) } /// Returns the virtual address this mapping occupies. @@ -213,9 +305,34 @@ impl Mapping { pub fn size(&self) -> usize { self.size } + + /// Returns the reservation this mapping lies in. + pub fn reservation(&self) -> &'a VirtualReservation { + self.reservation + } + + /// Returns the physical allocation this mapping maps, or `None` for a + /// multicast view. + pub fn physical(&self) -> Option<&'a PhysicalAllocation> { + match self.backing { + MappingBacking::Physical(phys) => Some(phys), + MappingBacking::Multicast(_) => None, + } + } + + /// Returns the multicast object this mapping is a view of, or `None` for + /// a mapping of physical memory. + pub fn multicast(&self) -> Option<&'a MulticastObject> { + match self.backing { + MappingBacking::Multicast(multicast) => Some(multicast), + MappingBacking::Physical(_) => None, + } + } } -impl Drop for Mapping { +impl Drop for Mapping<'_> { + /// `cuMemUnmap`, immediately; see the type docs for why in-flight work on + /// the range must be complete first. fn drop(&mut self) { unsafe { let _ = cuda_bindings::cuMemUnmap(self.va, self.size).result(); @@ -379,7 +496,10 @@ pub fn multicast_granularity( /// 4. [`Mapping::new_multicast`] + [`set_access`] per device. /// /// Teardown reverses it: mappings first, then bindings ([`MulticastBinding`]), -/// then this object, then the physical allocations. +/// then this object, then the physical allocations. Mappings and bindings +/// borrow this object, so the compiler enforces that they go first; the +/// physical allocations are not borrowed by the bindings (CUDA keeps bound +/// memory alive), so releasing them last stays a convention. /// /// Dropping releases the handle via `cuMemRelease` (the documented release /// path for multicast objects). @@ -425,15 +545,15 @@ impl MulticastObject { /// and every team device must already have been added. A CUDA context /// for the bound device must be current. /// - /// The returned [`MulticastBinding`] unbinds on drop; it must be dropped - /// before this object and before `phys`. + /// The returned [`MulticastBinding`] unbinds on drop. It borrows this + /// object, so it cannot outlive it; drop it before `phys` as well. pub fn bind_mem( &self, mc_offset: usize, phys: &PhysicalAllocation, mem_offset: usize, size: usize, - ) -> Result { + ) -> Result, DriverError> { unsafe { cuda_bindings::cuMulticastBindMem( self.handle, @@ -446,7 +566,7 @@ impl MulticastObject { .result()?; } Ok(MulticastBinding { - mc_handle: self.handle, + multicast: self, device: phys.device(), mc_offset, size, @@ -480,27 +600,35 @@ impl Drop for MulticastObject { /// One device's physical memory bound into a [`MulticastObject`]. /// /// Created by [`MulticastObject::bind_mem`]; unbinds via `cuMulticastUnbind` -/// on drop. Must be dropped before the multicast object and before the -/// physical allocation it binds. -pub struct MulticastBinding { - mc_handle: cuda_bindings::CUmemGenericAllocationHandle, +/// on drop. It borrows the multicast object, so it cannot outlive it (an +/// unbind against a released object handle would be a wild handle); drop it +/// before the physical allocation it binds as well. +pub struct MulticastBinding<'a> { + /// The object the memory is bound into; borrowed so the unbind runs + /// while the handle is still valid. + multicast: &'a MulticastObject, device: cuda_bindings::CUdevice, mc_offset: usize, size: usize, } -impl MulticastBinding { +impl<'a> MulticastBinding<'a> { /// Returns the device whose memory is bound. pub fn device(&self) -> cuda_bindings::CUdevice { self.device } + + /// Returns the multicast object the memory is bound into. + pub fn multicast(&self) -> &'a MulticastObject { + self.multicast + } } -impl Drop for MulticastBinding { +impl Drop for MulticastBinding<'_> { fn drop(&mut self) { unsafe { let _ = cuda_bindings::cuMulticastUnbind( - self.mc_handle, + self.multicast.handle(), self.device, self.mc_offset, self.size, diff --git a/cuda-core/tests/simt_device_buffer.rs b/cuda-core/tests/simt_device_buffer.rs index 9d67b0471d..c3dbd02d70 100644 --- a/cuda-core/tests/simt_device_buffer.rs +++ b/cuda-core/tests/simt_device_buffer.rs @@ -4,16 +4,28 @@ */ use cuda_core::{CudaContext, CudaStream, DeviceBuffer, DriverError, PinnedHostBuffer}; -use std::sync::{mpsc, Mutex, MutexGuard}; +use std::sync::{mpsc, Arc, Mutex, MutexGuard}; use std::thread::JoinHandle; use std::time::Duration; const BLOCKED_TIMEOUT: Duration = Duration::from_millis(100); const COMPLETION_TIMEOUT: Duration = Duration::from_secs(1); -static GATED_TEST_LOCK: Mutex<()> = Mutex::new(()); - -fn lock_gated_test() -> MutexGuard<'static, ()> { - GATED_TEST_LOCK +/// The error the sticky-state tests record; any code that is not `Ok` works. +const INJECTED: DriverError = DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE); +static TEST_LOCK: Mutex<()> = Mutex::new(()); + +/// Serializes every test in this file. +/// +/// The gated tests park a `cuLaunchHostFunc` callback on a channel until the +/// test thread releases it. While that callback blocks, a synchronizing driver +/// call made by *any other* test thread in this process (`cuMemFree` from a +/// `DeviceBuffer` drop, `cuStreamDestroy`, `cuCtxSynchronize`) waits for the +/// gated stream while holding the driver's context lock, and the gate owner +/// then blocks on that lock inside an unrelated call such as `cuEventCreate` +/// before it can release the gate. Observed as a hard deadlock under the +/// default parallel test runner; running the file serially removes it. +fn serialize_test() -> MutexGuard<'static, ()> { + TEST_LOCK .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) } @@ -80,6 +92,7 @@ fn gate_stream(stream: &CudaStream) -> mpsc::Sender<()> { #[test] fn device_buffer_from_host_roundtrip() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); @@ -97,8 +110,62 @@ fn device_buffer_from_host_roundtrip() { assert_eq!(host_vec, data); } +/// Every `DeviceBuffer` entry point binds `stream.context()` itself. A +/// freshly spawned thread has no current CUDA context, so a constructor that +/// reached `cuMemAlloc` without binding would fail with +/// `CUDA_ERROR_INVALID_CONTEXT` there (and, on a multi-GPU host, a thread +/// bound to another device would allocate in the wrong context). +#[test] +fn device_buffer_binds_the_stream_context_on_an_unbound_thread() { + let _guard = serialize_test(); + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let stream = ctx.new_stream().expect("failed to create CUDA stream"); + let pinned = + PinnedHostBuffer::from_slice(&ctx, &[5_u32, 8, 13, 21]).expect("failed to allocate pinned"); + + std::thread::spawn(move || { + let mut current: cuda_bindings::CUcontext = std::ptr::null_mut(); + assert_eq!( + unsafe { cuda_bindings::cuCtxGetCurrent(&mut current) }, + cuda_bindings::cudaError_enum_CUDA_SUCCESS + ); + assert!( + current.is_null(), + "a new thread must start without a context" + ); + + let zeroed = + DeviceBuffer::::zeroed(&stream, 4).expect("zeroed must bind the stream context"); + let from_host = DeviceBuffer::from_host(&stream, &[1_u32, 2, 3, 4]) + .expect("from_host must bind the stream context"); + // SAFETY: `pinned` outlives the `to_host_vec` synchronization below. + let from_pinned = unsafe { DeviceBuffer::from_pinned_host(&stream, &pinned) } + .expect("from_pinned_host must bind the stream context"); + // SAFETY: the buffer is fully written by `copy_from_host` before it is read. + let mut uninit = unsafe { DeviceBuffer::::uninitialized_async(&stream, 4) } + .expect("uninitialized_async must bind the stream context"); + uninit + .copy_from_host(&stream, &[9_u32, 9, 9, 9]) + .expect("copy_from_host must bind the stream context"); + + assert_eq!(zeroed.to_host_vec(&stream).expect("readback"), [0, 0, 0, 0]); + assert_eq!( + from_host.to_host_vec(&stream).expect("readback"), + [1, 2, 3, 4] + ); + assert_eq!( + from_pinned.to_host_vec(&stream).expect("readback"), + pinned.as_slice() + ); + assert_eq!(uninit.to_host_vec(&stream).expect("readback"), [9, 9, 9, 9]); + }) + .join() + .expect("worker thread panicked"); +} + #[test] fn device_buffer_zeroed_initializes_with_zeros() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); @@ -116,6 +183,7 @@ fn device_buffer_zeroed_initializes_with_zeros() { #[test] fn device_buffer_supports_empty_allocations() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); @@ -134,6 +202,7 @@ fn device_buffer_supports_empty_allocations() { #[test] fn device_buffer_rejects_allocation_size_overflow() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); let overflowing_len = usize::MAX / std::mem::size_of::() + 1; @@ -146,6 +215,7 @@ fn device_buffer_rejects_allocation_size_overflow() { #[test] fn device_buffer_async_compat_methods_roundtrip() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); @@ -192,7 +262,7 @@ fn device_buffer_async_compat_methods_roundtrip() { #[test] fn async_allocation_ordinary_drop_waits_for_cross_stream_work() { - let _gated_test_guard = lock_gated_test(); + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let allocation_stream = ctx .new_stream() @@ -231,7 +301,7 @@ fn async_allocation_ordinary_drop_waits_for_cross_stream_work() { #[test] fn async_allocation_drop_async_orders_free_after_allocation_stream() { - let _gated_test_guard = lock_gated_test(); + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let allocation_stream = ctx .new_stream() @@ -272,7 +342,7 @@ fn async_allocation_drop_async_orders_free_after_allocation_stream() { #[test] fn sync_allocation_allows_async_drop_after_queued_work() { - let _gated_test_guard = lock_gated_test(); + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); @@ -306,8 +376,198 @@ fn sync_allocation_allows_async_drop_after_queued_work() { .expect("synchronous allocation async free failed"); } +/// Gates `stream`, enqueues a device-to-host copy into pinned memory behind +/// the gate, records a sticky error, and runs the wait built by `make_wait` +/// on a worker thread. The wait must block until the gate is released, then +/// report the recorded error, and the copy must be complete by the time it +/// returns: a wait that returned the recorded error *instead of* waiting +/// would hand the caller host memory the DMA is still writing. +fn assert_wait_completes_under_sticky_error( + label: &str, + make_wait: impl FnOnce(&Arc, &Arc) -> W, +) where + W: FnOnce() -> Result<(), DriverError> + Send + 'static, +{ + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let stream = ctx.new_stream().expect("failed to create CUDA stream"); + let expected = [1_u32, 2, 3, 4]; + let src = DeviceBuffer::from_host(&stream, &expected).expect("failed to allocate source"); + let mut dst = PinnedHostBuffer::::zeroed(&ctx, expected.len()) + .expect("failed to allocate pinned destination"); + + let release_gate = gate_stream(&stream); + // SAFETY: `dst` stays alive, unread, and unaliased until the worker's + // wait has returned and been checked below. + unsafe { src.copy_to_pinned_host_async(&stream, &mut dst) } + .expect("failed to enqueue the gated copy"); + let wait = make_wait(&ctx, &stream); + ctx.record_err::<()>(Err(INJECTED)); + + let (started_tx, started_rx) = mpsc::channel(); + let (completion_tx, completion_rx) = mpsc::channel(); + let worker = std::thread::spawn(move || { + started_tx + .send(()) + .expect("failed to send wait worker start"); + completion_tx + .send(wait()) + .expect("failed to send wait result"); + }); + let result = finish_gated_worker(label, release_gate, started_rx, completion_rx, worker); + + assert_eq!( + result, + Err(INJECTED), + "{label}: the recorded error must be reported after the wait, not swallowed" + ); + assert_eq!( + ctx.check_err(), + Ok(()), + "{label}: the wait must have drained the recorded error" + ); + assert_eq!( + stream.query(), + Ok(true), + "{label}: the wait must have covered the in-flight copy" + ); + assert_eq!( + dst.as_slice(), + &expected, + "{label}: the copy must have landed before the wait returned" + ); +} + +#[test] +fn stream_synchronize_waits_for_in_flight_work_under_a_sticky_error() { + let _guard = serialize_test(); + assert_wait_completes_under_sticky_error("CudaStream::synchronize", |_, stream| { + let stream = Arc::clone(stream); + move || stream.synchronize() + }); +} + +#[test] +fn context_synchronize_waits_for_in_flight_work_under_a_sticky_error() { + let _guard = serialize_test(); + assert_wait_completes_under_sticky_error("CudaContext::synchronize", |ctx, _| { + let ctx = Arc::clone(ctx); + move || ctx.synchronize() + }); +} + +#[test] +fn event_synchronize_waits_for_in_flight_work_under_a_sticky_error() { + let _guard = serialize_test(); + assert_wait_completes_under_sticky_error("CudaEvent::synchronize", |_, stream| { + let event = stream + .record_event(None) + .expect("failed to record the completion event"); + move || event.synchronize() + }); +} + +/// Runs a safe enqueue-then-wait `DeviceBuffer` wrapper while a sticky error +/// lands *between* its enqueue and its wait. The gate callback records the +/// error when released, which the stream orders after the wrapper's own +/// pre-enqueue drain and before the copy it gates, so the wrapper cannot +/// return early: it must finish the copy and then report the error. +/// +/// `op` performs the copy from the source buffer into host memory it owns, +/// returning the wrapper's result and the host bytes it can still observe +/// (`None` when the wrapper's error path consumed them). +/// +/// Only a pinned destination discriminates: the driver stages a +/// device-to-host copy into pageable memory synchronously, so `copy_to_host` +/// and `to_host_vec` complete inside the enqueue call regardless of what the +/// wait does afterwards. They are exercised for coverage of the reporting +/// path; `copy_to_pinned_host` is the case that fails without the fix. +fn assert_copy_wrapper_finishes_before_reporting_sticky_error(label: &str, op: F) +where + F: FnOnce(&DeviceBuffer, &Arc) -> (Result<(), DriverError>, Option>) + + Send + + 'static, +{ + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let stream = ctx.new_stream().expect("failed to create CUDA stream"); + let expected = vec![1_u32, 2, 3, 4]; + let src = DeviceBuffer::from_host(&stream, &expected).expect("failed to allocate source"); + + let (release_gate, release_rx) = mpsc::channel(); + let ctx_for_gate = Arc::clone(&ctx); + stream + .launch_host_function(move || { + let _ = release_rx.recv(); + ctx_for_gate.record_err::<()>(Err(INJECTED)); + }) + .expect("failed to enqueue the recording gate"); + + let (started_tx, started_rx) = mpsc::channel(); + let (completion_tx, completion_rx) = mpsc::channel(); + let stream_for_thread = Arc::clone(&stream); + let worker = std::thread::spawn(move || { + started_tx + .send(()) + .expect("failed to send copy worker start"); + completion_tx + .send(op(&src, &stream_for_thread)) + .expect("failed to send copy result"); + }); + let (result, observed) = + finish_gated_worker(label, release_gate, started_rx, completion_rx, worker); + + assert_eq!( + result, + Err(INJECTED), + "{label}: the error recorded during the copy must be reported" + ); + assert_eq!( + stream.query(), + Ok(true), + "{label}: the wrapper must not return while its copy is in flight" + ); + if let Some(observed) = observed { + assert_eq!( + observed, expected, + "{label}: the copy must have completed before the wrapper returned" + ); + } +} + +#[test] +fn copy_to_pinned_host_finishes_the_copy_before_reporting_a_sticky_error() { + let _guard = serialize_test(); + assert_copy_wrapper_finishes_before_reporting_sticky_error( + "copy_to_pinned_host", + |src, stream| { + let mut dst = PinnedHostBuffer::::zeroed(stream.context(), src.len()) + .expect("failed to allocate pinned destination"); + let result = src.copy_to_pinned_host(stream, &mut dst); + (result, Some(dst.to_vec())) + }, + ); +} + +#[test] +fn copy_to_host_finishes_the_copy_before_reporting_a_sticky_error() { + let _guard = serialize_test(); + assert_copy_wrapper_finishes_before_reporting_sticky_error("copy_to_host", |src, stream| { + let mut dst = vec![0_u32; src.len()]; + let result = src.copy_to_host(stream, &mut dst); + (result, Some(dst)) + }); +} + +#[test] +fn to_host_vec_finishes_the_copy_before_reporting_a_sticky_error() { + let _guard = serialize_test(); + assert_copy_wrapper_finishes_before_reporting_sticky_error("to_host_vec", |src, stream| { + (src.to_host_vec(stream).map(drop), None) + }); +} + #[test] fn drop_async_bind_error_preserves_ordinary_cleanup() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); let dev = DeviceBuffer::::zeroed(&stream, 4096).expect("failed to allocate device buffer"); @@ -341,6 +601,7 @@ fn drop_async_bind_error_preserves_ordinary_cleanup() { #[test] fn from_host_with_pinned_source_allows_source_drop_after_return() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); @@ -360,6 +621,7 @@ fn from_host_with_pinned_source_allows_source_drop_after_return() { #[test] fn copy_from_host_with_pinned_source_allows_source_reuse_after_return() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); @@ -387,6 +649,7 @@ fn copy_from_host_with_pinned_source_allows_source_reuse_after_return() { // Run under `compute-sanitizer --tool memcheck` to catch a regression. #[test] fn uninitialized_async_implicit_drop_waits_for_pending_work() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); @@ -406,6 +669,7 @@ fn uninitialized_async_implicit_drop_waits_for_pending_work() { #[test] fn uninitialized_async_cast_elem_implicit_drop_is_stream_ordered() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); diff --git a/cuda-core/tests/simt_device_buffer_leaks.rs b/cuda-core/tests/simt_device_buffer_leaks.rs index 2e51c654ea..b6d6e0688c 100644 --- a/cuda-core/tests/simt_device_buffer_leaks.rs +++ b/cuda-core/tests/simt_device_buffer_leaks.rs @@ -18,9 +18,115 @@ //! pin the zero-length construction path, which must not call the driver //! allocator at all (`cuMemAlloc` rejects zero-byte requests). -use std::sync::Arc; +use std::ffi::c_void; +use std::mem::MaybeUninit; +use std::sync::{Arc, Mutex, MutexGuard}; + +use cuda_core::{CudaContext, DeviceBuffer, DriverError, IntoResult}; + +/// The error the sticky-state tests record; any code that is not `Ok` works. +const INJECTED: DriverError = DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE); +static TEST_LOCK: Mutex<()> = Mutex::new(()); + +/// Serializes the tests in this file: they account for device memory with +/// process-wide counters (`cuMemGetInfo`, the default pool's in-use bytes), +/// which a concurrently running test's allocations would perturb. +fn serialize_test() -> MutexGuard<'static, ()> { + TEST_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +/// Bytes of `device`'s default memory pool currently handed out by +/// `cuMemAllocAsync` (`CU_MEMPOOL_ATTR_USED_MEM_CURRENT`). Unlike +/// `cuMemGetInfo`, this moves by exactly the freed size once a +/// `cuMemFreeAsync` has executed, whether or not the pool releases the +/// backing memory to the OS. +fn pool_used_bytes(device: cuda_bindings::CUdevice) -> u64 { + let mut pool = MaybeUninit::uninit(); + let mut used = 0_u64; + unsafe { + cuda_bindings::cuDeviceGetDefaultMemPool(pool.as_mut_ptr(), device) + .result() + .expect("cuDeviceGetDefaultMemPool failed"); + cuda_bindings::cuMemPoolGetAttribute( + pool.assume_init(), + cuda_bindings::CUmemPool_attribute_enum_CU_MEMPOOL_ATTR_USED_MEM_CURRENT, + (&mut used as *mut u64).cast::(), + ) + .result() + .expect("cuMemPoolGetAttribute failed"); + } + used +} + +/// A recorded error must not turn a drop into a leak, and must still be +/// reported afterwards. `Drop` used to reach the driver through +/// `bind_to_thread`, which returns the recorded error *instead of* binding, +/// so `cuMemFree` ran against whatever context the thread happened to have +/// current (none, on a fresh thread) and the error was re-recorded. +#[test] +fn sync_allocation_is_freed_when_dropped_under_a_sticky_error() { + let _guard = serialize_test(); + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let stream = ctx.new_stream().expect("failed to create CUDA stream"); + let buf = DeviceBuffer::::zeroed(&stream, 4096).expect("failed to allocate device buffer"); + let ptr = buf.cu_deviceptr(); + stream.synchronize().expect("sync failed"); + + ctx.record_err::<()>(Err(INJECTED)); + drop(buf); -use cuda_core::{CudaContext, DeviceBuffer}; + let mut base = 0; + let mut size = 0; + assert_eq!( + unsafe { cuda_bindings::cuMemGetAddressRange_v2(&mut base, &mut size, ptr) }, + cuda_bindings::cudaError_enum_CUDA_ERROR_NOT_FOUND, + "drop must free a cuMemAlloc allocation with an error recorded" + ); + assert_eq!( + ctx.check_err(), + Err(INJECTED), + "the recorded error must survive the drop for the next fallible call" + ); +} + +/// The stream-ordered variant: `Drop` synchronizes the context and then +/// enqueues `cuMemFreeAsync`. It used to synchronize through the checked +/// wrapper, which reported the recorded error and skipped the free. +#[test] +fn async_allocation_is_freed_when_dropped_under_a_sticky_error() { + let _guard = serialize_test(); + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let stream = ctx.new_stream().expect("failed to create CUDA stream"); + let size = 1_usize << 20; + // SAFETY: the buffer is never read. + let buf = unsafe { DeviceBuffer::::uninitialized_async(&stream, size) } + .expect("failed to allocate stream-ordered device buffer"); + stream.synchronize().expect("sync failed"); + let used_with_buffer = pool_used_bytes(ctx.cu_device()); + assert!( + used_with_buffer >= size as u64, + "pool accounting must include the live allocation" + ); + + ctx.record_err::<()>(Err(INJECTED)); + drop(buf); + // The free is stream-ordered; the wait that retires it is also the next + // fallible call, which reports the recorded error after waiting. + assert_eq!( + stream.synchronize(), + Err(INJECTED), + "the recorded error must survive the drop for the next fallible call" + ); + + let used_after = pool_used_bytes(ctx.cu_device()); + assert!( + used_after + size as u64 <= used_with_buffer, + "drop must free a cuMemAllocAsync allocation with an error recorded: \ + pool bytes in use went {used_with_buffer} -> {used_after}" + ); +} /// A live buffer must hold exactly one extra context strong count, and the /// count must return to baseline once the buffer drops. A construction @@ -28,6 +134,7 @@ use cuda_core::{CudaContext, DeviceBuffer}; /// guard would) fails the live-count assertion by one per construction. #[test] fn ctx_strong_count_returns_to_baseline_after_buffer_lifecycle() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); @@ -70,6 +177,7 @@ fn ctx_strong_count_returns_to_baseline_after_buffer_lifecycle() { /// `cuMemGetInfo` before and after the cycles. #[test] fn vram_returns_to_baseline_after_buffer_cycles() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); @@ -111,6 +219,7 @@ fn vram_returns_to_baseline_after_buffer_cycles() { /// driver allocation. #[test] fn zero_length_construction_succeeds_for_both_constructors() { + let _guard = serialize_test(); let ctx = CudaContext::new(0).expect("failed to create CUDA context"); let stream = ctx.new_stream().expect("failed to create CUDA stream"); diff --git a/cuda-core/tests/simt_device_copy_impls.rs b/cuda-core/tests/simt_device_copy_impls.rs index 2429f7c7ec..1676500bad 100644 --- a/cuda-core/tests/simt_device_copy_impls.rs +++ b/cuda-core/tests/simt_device_copy_impls.rs @@ -33,3 +33,15 @@ fn device_copy_covers_core_parity_types() { assert_device_copy::>(); assert_device_copy::>(); } + +#[test] +fn device_copy_covers_thin_raw_pointers_only() { + // Thin pointers to any sized pointee, `DeviceCopy` or not: the pointee is + // never dereferenced, and every address is a valid pointer value. + assert_device_copy::<*const u8>(); + assert_device_copy::<*mut f32>(); + assert_device_copy::<*const String>(); + assert_device_copy::<*mut [u32; 4]>(); + // Fat pointers (`*const dyn Trait`, `*const [T]`) are intentionally NOT + // `DeviceCopy`; the compile_fail doctest on the trait pins the `dyn` case. +} diff --git a/cuda-core/tests/simt_module_globals.rs b/cuda-core/tests/simt_module_globals.rs new file mode 100644 index 0000000000..c3ee5372f9 --- /dev/null +++ b/cuda-core/tests/simt_module_globals.rs @@ -0,0 +1,171 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! `ConstantHandle` resolves a module-scope global with its driver-reported +//! size, keeps the module loaded, and refuses writes past the symbol. +//! +//! The module is hand-written PTX with two data symbols and no kernel, so +//! this needs the driver's JIT but no device-code toolchain. + +use cuda_core::{ConstantHandle, CudaContext, CudaModule, CudaStream, DriverError, IntoResult}; +use std::mem::MaybeUninit; +use std::sync::Arc; + +const PTX: &str = r#" +.version 7.8 +.target sm_75 +.address_size 64 + +.global .align 4 .b8 payload[16]; +.const .align 4 .b8 tuning[8]; +"#; + +const INVALID_VALUE: DriverError = + DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE); + +fn load() -> (Arc, Arc, Arc) { + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let stream = ctx.new_stream().expect("failed to create CUDA stream"); + let module = ctx + .load_module_from_ptx_src(PTX) + .expect("failed to JIT the test PTX"); + (ctx, stream, module) +} + +/// Reads `len` bytes back from the global `name` resolves to, synchronously. +fn read_back(module: &Arc, name: &str, len: usize) -> Vec { + let (dptr, size) = module.get_global(name).expect("get_global failed"); + assert!(len <= size, "test reads within the symbol"); + let mut host = vec![0_u8; len]; + module.context().bind_to_thread().expect("bind failed"); + unsafe { cuda_bindings::cuMemcpyDtoH_v2(host.as_mut_ptr().cast(), dptr, len) } + .result() + .expect("cuMemcpyDtoH failed"); + host +} + +fn staged(bytes: &[u8]) -> Box<[MaybeUninit]> { + bytes.iter().map(|&b| MaybeUninit::new(b)).collect() +} + +#[test] +fn constant_handle_carries_the_driver_reported_size() { + let (_ctx, _stream, module) = load(); + + let payload = ConstantHandle::new(&module, "payload").expect("payload must resolve"); + assert_eq!(payload.size(), 16); + assert!(Arc::ptr_eq(payload.module(), &module)); + + let tuning = ConstantHandle::new(&module, "tuning").expect("tuning must resolve"); + assert_eq!(tuning.size(), 8); + + assert_eq!( + ConstantHandle::new(&module, "missing").map(|handle| handle.size()), + Err(DriverError( + cuda_bindings::cudaError_enum_CUDA_ERROR_NOT_FOUND + )), + "an unknown symbol is the driver's error, not a handle" + ); +} + +#[test] +fn write_async_staged_rejects_oversize_writes_before_touching_the_device() { + let (_ctx, stream, module) = load(); + let handle = ConstantHandle::new(&module, "payload").expect("payload must resolve"); + + assert_eq!( + handle.write_async_staged(&stream, staged(&[0xAB; 17])), + Err(INVALID_VALUE), + "one byte past the symbol must be refused" + ); + stream.synchronize().expect("sync failed"); + assert_eq!( + read_back(&module, "payload", 16), + vec![0_u8; 16], + "a refused write must leave the symbol untouched" + ); + + let bytes: Vec = (1..=16).collect(); + handle + .write_async_staged(&stream, staged(&bytes)) + .expect("an exact-size write is accepted"); + stream.synchronize().expect("sync failed"); + assert_eq!(read_back(&module, "payload", 16), bytes); + + handle + .write_async_staged(&stream, staged(&[7, 7])) + .expect("a shorter write is accepted"); + stream.synchronize().expect("sync failed"); + assert_eq!(&read_back(&module, "payload", 16)[..2], &[7, 7]); +} + +#[test] +fn raw_writers_reject_oversize_lengths() { + let (_ctx, stream, module) = load(); + let handle = ConstantHandle::new(&module, "tuning").expect("tuning must resolve"); + let src = [0x5A_u8; 9]; + + // SAFETY: `src` outlives the calls and is long enough for every length + // passed; the oversize lengths are rejected before the pointer is used. + unsafe { + assert_eq!( + handle.write_async(&stream, src.as_ptr(), 9), + Err(INVALID_VALUE) + ); + assert_eq!(handle.write_blocking(src.as_ptr(), 9), Err(INVALID_VALUE)); + handle + .write_async(&stream, src.as_ptr(), 8) + .expect("exact-size async write"); + } + stream.synchronize().expect("sync failed"); + assert_eq!(read_back(&module, "tuning", 8), vec![0x5A; 8]); + + let blocking = [0x3C_u8; 8]; + // SAFETY: as above. + unsafe { handle.write_blocking(blocking.as_ptr(), 8) }.expect("exact-size blocking write"); + assert_eq!(read_back(&module, "tuning", 8), vec![0x3C; 8]); +} + +#[test] +fn handle_keeps_its_module_loaded() { + let (_ctx, stream, module) = load(); + let handle = ConstantHandle::new(&module, "payload").expect("payload must resolve"); + let clone = handle.clone(); + assert_eq!( + Arc::strong_count(&module), + 3, + "each handle holds the module" + ); + drop(module); + + // The module's only owners are now the handles; the write must still + // land in a loaded module rather than in an unloaded one's address. + handle + .write_async_staged(&stream, staged(&[9; 16])) + .expect("write through the sole owner of the module"); + stream.synchronize().expect("sync failed"); + assert_eq!(read_back(clone.module(), "payload", 16), vec![9; 16]); + drop(handle); + assert_eq!(Arc::strong_count(clone.module()), 1); +} + +/// The context check compares driver handles, not `Arc` identity: a second +/// `CudaContext` for the same device retains the same primary context and +/// its streams must be accepted. (A genuinely different context needs a +/// second GPU, which this suite does not assume.) +#[test] +fn writes_accept_a_stream_from_another_handle_to_the_same_context() { + let (ctx, _stream, module) = load(); + let handle = ConstantHandle::new(&module, "payload").expect("payload must resolve"); + + let same_device = CudaContext::new(ctx.ordinal()).expect("second context object"); + let stream = same_device + .new_stream() + .expect("stream on second context object"); + handle + .write_async_staged(&stream, staged(&[1; 16])) + .expect("the same primary context is the same context"); + stream.synchronize().expect("sync failed"); +} diff --git a/cuda-core/tests/simt_vmm_multicast.rs b/cuda-core/tests/simt_vmm_multicast.rs index a5716b9860..223b13e4a6 100644 --- a/cuda-core/tests/simt_vmm_multicast.rs +++ b/cuda-core/tests/simt_vmm_multicast.rs @@ -11,9 +11,9 @@ //! and tear down in the required order. //! //! The switch-side reduction semantics can only be exercised from device -//! code (`multimem.ld_reduce` / `multimem.st`); that lives in the -//! `nvls_all_reduce` example. Host code must not dereference the multicast -//! mapping, so this test only establishes and tears it down. +//! code (`multimem.ld_reduce` / `multimem.st`), which this crate does not +//! build. Host code must not dereference the multicast mapping, so this test +//! only establishes and tears it down. //! //! Requires an NVLink-switch system (e.g. HGX H100/B200); skips elsewhere. //! Compiled out entirely when the CUDA toolkit predates the multicast API @@ -73,10 +73,11 @@ fn multicast_two_gpu_team_roundtrip() { } // The physical size must satisfy the allocation granularity as well as - // the multicast granularity, so align to both. + // the multicast granularity, so align to both. Bindings borrow the team + // and mappings borrow their reservation and backing, so each kind of + // handle lives in its own collection, created after what it borrows. let contexts = [&ctx0, &ctx1]; let mut physes = Vec::new(); - let mut bindings = Vec::new(); for (i, ctx) in contexts.iter().enumerate() { ctx.bind_to_thread().expect("bind context"); let alloc_gran = @@ -84,30 +85,41 @@ fn multicast_two_gpu_team_roundtrip() { let phys_size = vmm::align_size(size, alloc_gran); let phys = vmm::PhysicalAllocation::new(devices[i], phys_size).expect("cuMemCreate failed"); assert_eq!(phys.device(), devices[i]); + physes.push(phys); + } + let mut bindings = Vec::new(); + for (i, (ctx, phys)) in contexts.iter().zip(&physes).enumerate() { + ctx.bind_to_thread().expect("bind context"); let binding = team - .bind_mem(0, &phys, 0, size) + .bind_mem(0, phys, 0, size) .expect("cuMulticastBindMem failed"); assert_eq!(binding.device(), devices[i]); - physes.push(phys); + assert!(std::ptr::eq(binding.multicast(), &team)); bindings.push(binding); } - // Per GPU: (uc_va, uc_map, mc_va, mc_map) - let mut mappings = Vec::new(); - for (i, ctx) in contexts.iter().enumerate() { + // Per GPU: a unicast reservation over the device's own memory and a + // multicast reservation over the team. + let mut reservations = Vec::new(); + for ctx in &contexts { ctx.bind_to_thread().expect("bind context"); - let uc_va = vmm::VirtualReservation::new(size, 0).expect("unicast VA reserve"); - let uc_map = - vmm::Mapping::new(uc_va.base(), size, &physes[i], 0).expect("unicast cuMemMap"); + let mc_va = vmm::VirtualReservation::new(size, granularity).expect("multicast VA reserve"); + reservations.push((uc_va, mc_va)); + } + let mut mappings = Vec::new(); + for (i, ((uc_va, mc_va), phys)) in reservations.iter().zip(&physes).enumerate() { + contexts[i].bind_to_thread().expect("bind context"); + + let uc_map = vmm::Mapping::new(uc_va, 0, phys, 0, size).expect("unicast cuMemMap"); vmm::set_access(uc_va.base(), size, &[devices[i]]).expect("unicast set_access"); - let mc_va = vmm::VirtualReservation::new(size, granularity).expect("multicast VA reserve"); let mc_map = - vmm::Mapping::new_multicast(mc_va.base(), size, &team, 0).expect("multicast cuMemMap"); + vmm::Mapping::new_multicast(mc_va, 0, &team, 0, size).expect("multicast cuMemMap"); + assert!(mc_map.multicast().is_some_and(|t| std::ptr::eq(t, &team))); vmm::set_access(mc_va.base(), size, &[devices[i]]).expect("multicast set_access"); - mappings.push((uc_va, uc_map, mc_va, mc_map)); + mappings.push((uc_map, mc_map)); } ctx0.bind_to_thread().expect("bind ctx0"); @@ -116,7 +128,7 @@ fn multicast_two_gpu_team_roundtrip() { let stream0 = ctx0.default_stream(); unsafe { cuda_core::simt::memory::memcpy_htod_async( - mappings[0].0.base(), + mappings[0].0.va(), pattern.as_ptr(), byte_len, stream0.cu_stream(), @@ -129,7 +141,7 @@ fn multicast_two_gpu_team_roundtrip() { unsafe { cuda_core::simt::memory::memcpy_dtoh_async( readback.as_mut_ptr(), - mappings[0].0.base(), + mappings[0].0.va(), byte_len, stream0.cu_stream(), ) @@ -142,9 +154,11 @@ fn multicast_two_gpu_team_roundtrip() { ); // Teardown order: mappings, then bindings, then team, then physical - // allocations. + // allocations and reservations. The borrows enforce the first three + // steps; the last is convention. drop(mappings); drop(bindings); drop(team); drop(physes); + drop(reservations); } diff --git a/cuda-core/tests/simt_vmm_p2p.rs b/cuda-core/tests/simt_vmm_p2p.rs index 8adca11b01..ba69c27153 100644 --- a/cuda-core/tests/simt_vmm_p2p.rs +++ b/cuda-core/tests/simt_vmm_p2p.rs @@ -26,9 +26,12 @@ use cuda_core::simt::context::CudaContext; use cuda_core::simt::peer; use cuda_core::simt::vmm; -use cuda_core::IntoResult; +use cuda_core::{DriverError, IntoResult}; use std::mem::MaybeUninit; +const INVALID_VALUE: DriverError = + DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE); + fn gpu_count() -> Result { unsafe { cuda_core::init(0)? }; let mut count = MaybeUninit::uninit(); @@ -54,8 +57,11 @@ fn vmm_alloc_map_set_access_roundtrip() { let va = vmm::VirtualReservation::new(alloc_size, 0).expect("cuMemAddressReserve failed"); assert_ne!(va.base(), 0); - let mapping = vmm::Mapping::new(va.base(), alloc_size, &phys, 0).expect("cuMemMap failed"); + let mapping = vmm::Mapping::new(&va, 0, &phys, 0, alloc_size).expect("cuMemMap failed"); assert_eq!(mapping.va(), va.base()); + assert_eq!(mapping.size(), alloc_size); + assert!(std::ptr::eq(mapping.reservation(), &va)); + assert!(mapping.physical().is_some_and(|p| std::ptr::eq(p, &phys))); vmm::set_access(va.base(), alloc_size, &[ctx.cu_device()]).expect("cuMemSetAccess failed"); @@ -86,11 +92,47 @@ fn vmm_alloc_map_set_access_roundtrip() { assert_eq!(readback, pattern, "single-GPU VMM roundtrip failed"); + // The mapping borrows both, so this is the only order that compiles. drop(mapping); drop(va); drop(phys); } +/// The driver cannot tell where our reservation ends (VA ranges are +/// process-wide), so `Mapping::new` bounds the range itself before calling +/// `cuMemMap`, and reports a wrap or an overrun as `CUDA_ERROR_INVALID_VALUE`. +#[test] +fn mapping_rejects_ranges_outside_the_reservation_or_allocation() { + let ctx = CudaContext::new(0).expect("failed to create context for GPU 0"); + let granularity = + vmm::allocation_granularity(ctx.cu_device()).expect("failed to query granularity"); + let size = vmm::align_size(4096, granularity); + + let phys = vmm::PhysicalAllocation::new(ctx.cu_device(), size).expect("cuMemCreate failed"); + // Twice the allocation, so there is a valid second half to map into. + let va = vmm::VirtualReservation::new(2 * size, 0).expect("cuMemAddressReserve failed"); + + assert_eq!( + vmm::Mapping::new(&va, 2 * size, &phys, 0, size).err(), + Some(INVALID_VALUE), + "a range past the end of the reservation is refused" + ); + assert_eq!( + vmm::Mapping::new(&va, 0, &phys, granularity, size).err(), + Some(INVALID_VALUE), + "a range past the end of the allocation is refused" + ); + assert_eq!( + vmm::Mapping::new(&va, usize::MAX, &phys, 0, size).err(), + Some(INVALID_VALUE), + "offset arithmetic that would wrap is refused" + ); + + let map = vmm::Mapping::new(&va, size, &phys, 0, size).expect("the second half is in range"); + assert_eq!(map.va(), va.base() + size as u64); + assert!(std::ptr::eq(map.reservation(), &va)); +} + #[test] fn p2p_vmm_cross_gpu_access() { let count = gpu_count().expect("failed to get device count"); @@ -118,7 +160,7 @@ fn p2p_vmm_cross_gpu_access() { // --- Reserve VA and map on GPU 0 --- let va0 = vmm::VirtualReservation::new(alloc_size, 0).expect("VA reserve for GPU 0"); - let map0 = vmm::Mapping::new(va0.base(), alloc_size, &phys, 0).expect("map on GPU 0"); + let map0 = vmm::Mapping::new(&va0, 0, &phys, 0, alloc_size).expect("map on GPU 0"); vmm::set_access( va0.base(), alloc_size, @@ -128,7 +170,7 @@ fn p2p_vmm_cross_gpu_access() { // --- Reserve VA and map the SAME physical memory on GPU 1 --- let va1 = vmm::VirtualReservation::new(alloc_size, 0).expect("VA reserve for GPU 1"); - let map1 = vmm::Mapping::new(va1.base(), alloc_size, &phys, 0).expect("map on GPU 1"); + let map1 = vmm::Mapping::new(&va1, 0, &phys, 0, alloc_size).expect("map on GPU 1"); vmm::set_access( va1.base(), alloc_size, diff --git a/cuda-core/tests/vmm.rs b/cuda-core/tests/vmm.rs index ac36a4d42c..76918a8ee3 100644 --- a/cuda-core/tests/vmm.rs +++ b/cuda-core/tests/vmm.rs @@ -9,13 +9,16 @@ //! granularity, create physical memory, reserve a virtual range, map the //! memory, grant access, roundtrip data through the mapping, and tear //! down in a leak-free order (mapping before reservation and physical -//! allocation handle). +//! allocation handle, which the mapping's borrows enforce). //! //! Skips when no GPU is present. -use cuda_core::{vmm, Device, IntoResult}; +use cuda_core::{vmm, Device, DriverError, IntoResult}; use std::mem::MaybeUninit; +const INVALID_VALUE: DriverError = + DriverError(cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE); + #[test] fn align_size_handles_general_granularities() { assert_eq!(vmm::align_size(0, 6), 0); @@ -73,7 +76,12 @@ fn vmm_single_gpu_roundtrip() { assert_eq!(va.size(), size); { - let _map = vmm::Mapping::new(va.base(), size, &phys, 0).expect("cuMemMap"); + let map = vmm::Mapping::new(&va, 0, &phys, 0, size).expect("cuMemMap"); + assert_eq!(map.va(), va.base()); + assert_eq!(map.size(), size); + assert!(std::ptr::eq(map.reservation(), &va)); + assert!(map.physical().is_some_and(|p| std::ptr::eq(p, &phys))); + assert!(map.multicast().is_none()); vmm::set_access(va.base(), size, &[cu_device]).expect("cuMemSetAccess"); // Roundtrip a pattern through the mapped range. @@ -100,6 +108,54 @@ fn vmm_single_gpu_roundtrip() { } assert_eq!(readback, pattern, "data must roundtrip through the mapping"); // The mapping drops here, before the reservation and the physical - // allocation handle, matching a leak-free teardown order. + // allocation handle it borrows; the borrows make any other order a + // compile error. + } +} + +/// The driver cannot tell where our reservation ends (VA ranges are +/// process-wide), so `Mapping::new` bounds the range itself before calling +/// `cuMemMap`, and reports a wrap or an overrun as `CUDA_ERROR_INVALID_VALUE`. +#[test] +fn mapping_rejects_ranges_outside_the_reservation_or_allocation() { + match gpu_count() { + Ok(count) if count > 0 => {} + Ok(_) => { + eprintln!("SKIPPED: mapping bounds test requires a GPU"); + return; + } + Err(error) => { + eprintln!("SKIPPED: CUDA unavailable ({error:?})"); + return; + } } + + let device = Device::new(0).expect("GPU 0 device"); + let cu_device = device.cu_device(); + let granularity = vmm::allocation_granularity(cu_device).expect("allocation granularity query"); + let size = vmm::align_size(1 << 20, granularity); + + let phys = vmm::PhysicalAllocation::new(cu_device, size).expect("cuMemCreate"); + // Twice the allocation, so there is a valid second half to map into. + let va = vmm::VirtualReservation::new(2 * size, 0).expect("cuMemAddressReserve"); + + assert_eq!( + vmm::Mapping::new(&va, 2 * size, &phys, 0, size).err(), + Some(INVALID_VALUE), + "a range past the end of the reservation is refused" + ); + assert_eq!( + vmm::Mapping::new(&va, 0, &phys, granularity, size).err(), + Some(INVALID_VALUE), + "a range past the end of the allocation is refused" + ); + assert_eq!( + vmm::Mapping::new(&va, usize::MAX, &phys, 0, size).err(), + Some(INVALID_VALUE), + "offset arithmetic that would wrap is refused" + ); + + let map = vmm::Mapping::new(&va, size, &phys, 0, size).expect("the second half is in range"); + assert_eq!(map.va(), va.base() + size as u64); + assert!(std::ptr::eq(map.reservation(), &va)); } diff --git a/cuda-core/tests/vmm_multicast.rs b/cuda-core/tests/vmm_multicast.rs index 0c8cb398d5..19ac36b0a9 100644 --- a/cuda-core/tests/vmm_multicast.rs +++ b/cuda-core/tests/vmm_multicast.rs @@ -11,9 +11,9 @@ //! and tear down in a clean, leak-free order. //! //! The switch-side reduction semantics can only be exercised from device -//! code (`multimem.ld_reduce` / `multimem.st`); that lives in the -//! `nvls_all_reduce` example. Host code must not dereference the multicast -//! mapping, so this test only establishes and tears it down. +//! code (`multimem.ld_reduce` / `multimem.st`), which this crate does not +//! build. Host code must not dereference the multicast mapping, so this test +//! only establishes and tears it down. //! //! Requires an NVLink-switch system (e.g. HGX H100/B200); skips elsewhere. @@ -68,10 +68,11 @@ fn multicast_two_gpu_team_roundtrip() { } // The physical size must satisfy the allocation granularity as well as - // the multicast granularity, so align to both. + // the multicast granularity, so align to both. Bindings borrow the team + // and mappings borrow their reservation and backing, so each kind of + // handle lives in its own collection, created after what it borrows. let contexts = [&ctx0, &ctx1]; let mut physes = Vec::new(); - let mut bindings = Vec::new(); for (i, ctx) in contexts.iter().enumerate() { ctx.bind_to_thread().expect("bind context"); let alloc_gran = @@ -79,30 +80,41 @@ fn multicast_two_gpu_team_roundtrip() { let phys_size = vmm::align_size(size, alloc_gran); let phys = vmm::PhysicalAllocation::new(devices[i], phys_size).expect("cuMemCreate failed"); assert_eq!(phys.device(), devices[i]); + physes.push(phys); + } + let mut bindings = Vec::new(); + for (i, (ctx, phys)) in contexts.iter().zip(&physes).enumerate() { + ctx.bind_to_thread().expect("bind context"); let binding = team - .bind_mem(0, &phys, 0, size) + .bind_mem(0, phys, 0, size) .expect("cuMulticastBindMem failed"); assert_eq!(binding.device(), devices[i]); - physes.push(phys); + assert!(std::ptr::eq(binding.multicast(), &team)); bindings.push(binding); } - // Per GPU: (uc_va, uc_map, mc_va, mc_map) - let mut mappings = Vec::new(); - for (i, ctx) in contexts.iter().enumerate() { + // Per GPU: a unicast reservation over the device's own memory and a + // multicast reservation over the team. + let mut reservations = Vec::new(); + for ctx in &contexts { ctx.bind_to_thread().expect("bind context"); - let uc_va = vmm::VirtualReservation::new(size, 0).expect("unicast VA reserve"); - let uc_map = - vmm::Mapping::new(uc_va.base(), size, &physes[i], 0).expect("unicast cuMemMap"); + let mc_va = vmm::VirtualReservation::new(size, granularity).expect("multicast VA reserve"); + reservations.push((uc_va, mc_va)); + } + let mut mappings = Vec::new(); + for (i, ((uc_va, mc_va), phys)) in reservations.iter().zip(&physes).enumerate() { + contexts[i].bind_to_thread().expect("bind context"); + + let uc_map = vmm::Mapping::new(uc_va, 0, phys, 0, size).expect("unicast cuMemMap"); vmm::set_access(uc_va.base(), size, &[devices[i]]).expect("unicast set_access"); - let mc_va = vmm::VirtualReservation::new(size, granularity).expect("multicast VA reserve"); let mc_map = - vmm::Mapping::new_multicast(mc_va.base(), size, &team, 0).expect("multicast cuMemMap"); + vmm::Mapping::new_multicast(mc_va, 0, &team, 0, size).expect("multicast cuMemMap"); + assert!(mc_map.multicast().is_some_and(|t| std::ptr::eq(t, &team))); vmm::set_access(mc_va.base(), size, &[devices[i]]).expect("multicast set_access"); - mappings.push((uc_va, uc_map, mc_va, mc_map)); + mappings.push((uc_map, mc_map)); } ctx0.bind_to_thread().expect("bind ctx0"); @@ -110,7 +122,7 @@ fn multicast_two_gpu_team_roundtrip() { let byte_len = std::mem::size_of_val(pattern.as_slice()); unsafe { cuda_core::sys::cuMemcpyHtoD_v2( - mappings[0].0.base(), + mappings[0].0.va(), pattern.as_ptr() as *const std::ffi::c_void, byte_len, ) @@ -122,7 +134,7 @@ fn multicast_two_gpu_team_roundtrip() { unsafe { cuda_core::sys::cuMemcpyDtoH_v2( readback.as_mut_ptr() as *mut std::ffi::c_void, - mappings[0].0.base(), + mappings[0].0.va(), byte_len, ) .result() @@ -134,9 +146,11 @@ fn multicast_two_gpu_team_roundtrip() { ); // Clean teardown order: mappings, then bindings, then team, then physical - // allocation handles. + // allocation handles and reservations. The borrows enforce the first + // three steps; the last is convention. drop(mappings); drop(bindings); drop(team); drop(physes); + drop(reservations); } diff --git a/cutile-rs/scripts/run_gpu_tests.sh b/cutile-rs/scripts/run_gpu_tests.sh index b40b2254a4..b10369a791 100755 --- a/cutile-rs/scripts/run_gpu_tests.sh +++ b/cutile-rs/scripts/run_gpu_tests.sh @@ -93,6 +93,7 @@ for test_target in \ simt_context_sync_policy \ simt_device_buffer \ simt_device_buffer_leaks \ + simt_module_globals \ simt_pinned_host_buffer \ simt_stream_priority \ simt_stream_query \ @@ -122,7 +123,8 @@ for test_target in \ execute_once \ execution_lock \ pool_allocation \ - reactor_correctness + reactor_correctness \ + simt_execute_once do run_step \ "cuda-async GPU integration test ${test_target}" \