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}" \