refactor(cuda-core): share host handles and allocations across Tile and SIMT

.github/workflows/cutile-rs.yml

@@ -289,7 +289,7 @@ jobs:
 
       - name: Native f16 DeviceCopy
         working-directory: ${{ github.workspace }}
-        run: cargo test -p cuda-core --features f16 --test simt_device_copy_impls
+        run: cargo test -p cuda-core --features f16 --test device_copy_impls
 
       - name: loom model check (interleavings + weak memory)
         working-directory: ${{ github.workspace }}

cuda-async/README.md

@@ -17,7 +17,7 @@ until `.sync()`, `.sync_on()`, or `.await` is called.
 use cutile::prelude::*;
 
 fn main() -> Result<(), DeviceError> {
-    let device = cuda_core::Device::new(0)?;
+    let device = cuda_core::Context::new(0)?;
     let stream = device.new_stream()?;
 
     let mut z = api::zeros::<f32>(&[16, 16]).sync_on(&stream)?;

cuda-async/src/cuda_graph.rs

@@ -7,7 +7,7 @@ use crate::device_context::with_default_device_policy;
 use crate::device_future::DeviceFuture;
 use crate::device_operation::{DeviceOp, ExecutionContext, GraphNode, ReplayResource};
 use crate::error::DeviceError;
-use cuda_core::{sys, Device, IntoResult, Stream};
+use cuda_core::{sys, Context, IntoResult, Stream};
 use std::collections::HashMap;
 use std::future::IntoFuture;
 use std::mem::MaybeUninit;
@@ -67,7 +67,7 @@ struct GraphExecHandle {
     /// Contexts of `update` ops not yet consumed by a replay.
     pending: Mutex<Vec<ExecutionContext>>,
     /// Bound before destruction: the driver needs the owning context current.
-    device: Arc<Device>,
+    device: Arc<Context>,
     cu_graph: sys::CUgraph,
     cu_graph_exec: sys::CUgraphExec,
 }
@@ -85,7 +85,7 @@ impl GraphExecHandle {
     /// instantiation fails, and owned by the returned handle otherwise
     /// (whose `Drop` also covers an upload failure).
     fn instantiate(
-        device: Arc<Device>,
+        device: Arc<Context>,
         cu_graph: sys::CUgraph,
         stream: &Stream,
     ) -> Result<Self, DeviceError> {
@@ -165,7 +165,7 @@ fn capture_on<R>(
     stream: &Arc<Stream>,
     record: impl FnOnce() -> Result<R, DeviceError>,
 ) -> Result<(R, Arc<GraphExecHandle>), DeviceError> {
-    let device = stream.device().clone();
+    let device = stream.context().clone();
     device.bind_to_thread()?;
 
     unsafe {
@@ -204,7 +204,7 @@ fn capture_on<R>(
     let exec = GraphExecHandle::instantiate(device, cu_graph, stream)?;
 
     // Drain the upload. The recorded work itself has *not* run.
-    unsafe { stream.synchronize() }?;
+    stream.synchronize()?;
 
     Ok((output, Arc::new(exec)))
 }

cuda-async/src/device_buffer.rs

@@ -3,61 +3,11 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-//! Wrappers around CUDA device pointers: owned (freed on drop), foreign (kept
-//! alive by an external owner via a liveness token), and raw-borrowed (an
-//! `unsafe` escape that does nothing on drop).
+//! Shared allocation ownership and typed device addresses.
 
-use crate::device_context::with_deallocator_stream;
-use cuda_core::free_async;
 use cuda_core::sys::CUdeviceptr;
+pub use cuda_core::{DeviceAllocation, RawDeviceBuffer as DeviceBuffer};
 use std::marker::PhantomData;
-use std::sync::Arc;
-
-/// A CUDA device allocation owned by an external framework (cudarc, torch, a VMM
-/// range, ...).
-///
-/// A borrowed [`DeviceBuffer`] can hold an `Arc<dyn DeviceAllocation>` as a
-/// *liveness token*: while the buffer — and anything derived from it — is alive,
-/// the owner's refcount is nonzero, so the memory cannot be freed out from under
-/// cutile. This turns liveness from an ongoing caller promise into a fact
-/// enforced at construction.
-///
-/// # Safety
-/// Implementors must guarantee that, for as long as the implementor is alive:
-/// - `device_ptr()` returns a valid CUDA device pointer on `device_id()`,
-/// - the allocation there is at least `len_bytes()` bytes, and
-/// - all three getters are **stable**: they return the same values on every
-///   call for the implementor's entire lifetime. Consumers read them once at
-///   construction and cache the results (`DeviceBuffer::foreign`), so an
-///   implementor whose pointer moves would silently diverge from the cached
-///   copy. This is the same stability `std`'s allocator traits demand of a
-///   currently-allocated block.
-///
-/// This positional obligation is discharged **once, in the `impl`** — not at
-/// each borrow. The temporal obligation (nobody else touches the bytes while
-/// cutile work is in flight) cannot live here: it is per-borrow, and is why
-/// `Tensor::from_foreign` remains an `unsafe` call.
-pub unsafe trait DeviceAllocation: Send + Sync + 'static {
-    /// The device pointer of the allocation.
-    fn device_ptr(&self) -> CUdeviceptr;
-    /// The size of the allocation in bytes.
-    fn len_bytes(&self) -> usize;
-    /// The device ordinal the allocation lives on.
-    fn device_id(&self) -> usize;
-}
-
-/// How a [`DeviceBuffer`]'s memory is owned.
-enum Owner {
-    /// cutile allocated the memory; free it asynchronously on drop.
-    Owned,
-    /// A raw, unowned borrow (via `borrowed_from_raw_parts`); do nothing on drop.
-    /// Liveness is the caller's unverified obligation.
-    Borrowed,
-    /// A foreign allocation kept alive by its owner; do nothing on drop, and
-    /// hold the owner so the memory provably outlives this buffer. The `Arc` is
-    /// held purely for that RAII liveness effect — never read directly.
-    Foreign(#[allow(dead_code)] Arc<dyn DeviceAllocation>),
-}
 
 /// A non-owning, copyable handle to a typed device memory address.
 #[derive(Debug, Copy, Clone)]
@@ -85,141 +35,3 @@ impl<T> DevicePointer<T> {
         }
     }
 }
-
-/// A type-erased handle to a CUDA device memory allocation.
-///
-/// `DeviceBuffer` manages a raw byte buffer on the GPU. It stores only the device pointer,
-/// byte length, and device id — all type information lives in higher-level wrappers
-/// such as `Tensor<T>`.
-///
-/// An *owned* buffer (the default) is freed asynchronously on drop. A *foreign*
-/// buffer (see [`foreign`](Self::foreign)) wraps memory owned by an external
-/// allocator (cudarc, torch, a VMM range, ...), holds a liveness token, and is a
-/// no-op on drop. A raw *borrowed* buffer
-/// (see [`borrowed_from_raw_parts`](Self::borrowed_from_raw_parts)) is the same
-/// no-op-on-drop but without a token — an `unsafe` escape whose liveness is the
-/// caller's responsibility.
-pub struct DeviceBuffer {
-    device_id: usize,
-    cudptr: CUdeviceptr,
-    len: usize,
-    owner: Owner,
-}
-
-impl std::fmt::Debug for DeviceBuffer {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let kind = match self.owner {
-            Owner::Owned => "owned",
-            Owner::Borrowed => "borrowed",
-            Owner::Foreign(_) => "foreign",
-        };
-        f.debug_struct("DeviceBuffer")
-            .field("device_id", &self.device_id)
-            .field("cudptr", &self.cudptr)
-            .field("len", &self.len)
-            .field("owner", &kind)
-            .finish()
-    }
-}
-
-unsafe impl Send for DeviceBuffer {}
-unsafe impl Sync for DeviceBuffer {}
-
-impl Drop for DeviceBuffer {
-    fn drop(&mut self) {
-        // Only owned buffers are freed here. Borrowed/foreign buffers wrap memory
-        // owned elsewhere; freeing it would be a double free (or invalid, e.g.
-        // for a VMM-mapped range). A `Foreign` owner's `Arc` drops with `self`.
-        if !matches!(self.owner, Owner::Owned) {
-            return;
-        }
-        // A destructor has no way to report failure, and a panic here would
-        // abort the process if this drop runs during an unwind. Both a missing
-        // device context and a refused `cuMemFreeAsync` are therefore ignored:
-        // the allocation leaks until the primary context is released, which
-        // is the least bad outcome for memory the driver would not free.
-        // SAFETY: the pointer was allocated stream-ordered by this crate, and
-        // the driver completes queued work on the deallocator stream before
-        // the memory is handed out again.
-        let _ = unsafe {
-            with_deallocator_stream(self.device_id, |stream| free_async(self.cudptr, stream))
-        };
-    }
-}
-
-impl DeviceBuffer {
-    /// Constructs an owned `DeviceBuffer` from a raw device pointer, byte length, and device id.
-    ///
-    /// The buffer takes ownership of the allocation and frees it asynchronously on drop.
-    ///
-    /// # Safety
-    /// The caller must ensure `dptr` points to a valid device allocation of at least
-    /// `len_bytes` bytes that is safe to free via the deallocator stream.
-    pub unsafe fn from_raw_parts(dptr: CUdeviceptr, len_bytes: usize, device_id: usize) -> Self {
-        Self {
-            cudptr: dptr,
-            len: len_bytes,
-            device_id,
-            owner: Owner::Owned,
-        }
-    }
-
-    /// Wraps `len_bytes` of a foreign allocation, holding `owner` alive so the
-    /// memory outlives this buffer. Dropping it does **not** free the memory.
-    ///
-    /// `len_bytes` is the borrowing view's **logical** size (what the tensor
-    /// invariant checks), not the owner's allocation size — the owner may be
-    /// larger, and the addressable-extent-vs-owner check happens in
-    /// `Tensor::from_foreign` before this is called. Liveness is enforced by
-    /// `owner`'s refcount; pointer validity was asserted once when `owner`'s
-    /// type implemented [`DeviceAllocation`].
-    pub fn foreign(owner: Arc<dyn DeviceAllocation>, len_bytes: usize) -> Self {
-        Self {
-            cudptr: owner.device_ptr(),
-            len: len_bytes,
-            device_id: owner.device_id(),
-            owner: Owner::Foreign(owner),
-        }
-    }
-
-    /// Constructs a borrowed `DeviceBuffer` that wraps memory owned by an external
-    /// allocator. Dropping it does **not** free the allocation.
-    ///
-    /// # Safety
-    /// This borrows, rather than takes ownership of, `dptr`; dropping the buffer
-    /// never frees it. The caller must ensure `dptr` points to a valid device
-    /// allocation of at least `len_bytes` bytes, and must keep that allocation
-    /// mapped — never freed, reallocated, or resized — for the entire lifetime of
-    /// this buffer and every value derived from it. When wrapped in a tensor used
-    /// as a mutable output, the owner must additionally not access the same memory
-    /// concurrently; see `Tensor::borrow_raw_parts` in the `cutile` crate for the
-    /// full obligations.
-    pub unsafe fn borrowed_from_raw_parts(
-        dptr: CUdeviceptr,
-        len_bytes: usize,
-        device_id: usize,
-    ) -> Self {
-        Self {
-            cudptr: dptr,
-            len: len_bytes,
-            device_id,
-            owner: Owner::Borrowed,
-        }
-    }
-    /// Returns whether the allocation is empty (zero bytes).
-    pub fn is_empty(&self) -> bool {
-        self.len_bytes() == 0
-    }
-    /// Returns the byte length of the allocation.
-    pub fn len_bytes(&self) -> usize {
-        self.len
-    }
-    /// Returns the raw CUDA device pointer.
-    pub fn cu_deviceptr(&self) -> CUdeviceptr {
-        self.cudptr
-    }
-    /// Returns the device id this allocation belongs to.
-    pub fn device_id(&self) -> usize {
-        self.device_id
-    }
-}

cuda-async/src/device_context.rs

@@ -7,7 +7,7 @@
 //!
 //! ## Architecture
 //!
-//! - **Global (process-wide)**: one [`Device`] per device id, shared across all threads
+//! - **Global (process-wide)**: one [`Context`] per device id, shared across all threads
 //!   via [`OnceLock`].
 //!
 //! - **Per-thread**: Scheduling policy and deallocator stream remain thread-local, since
@@ -18,7 +18,7 @@
 
 use crate::error::{device_assert, device_error, DeviceError};
 use crate::scheduling_policies::{SchedulingPolicy, StreamPoolRoundRobin};
-use cuda_core::{Device, Function, MemPool, Module, Stream};
+use cuda_core::{Context, Function, MemPool, Module, Stream};
 use std::cell::Cell;
 use std::collections::HashMap;
 use std::hash::{DefaultHasher, Hash, Hasher};
@@ -89,28 +89,28 @@ pub struct Validator {
     pub launch_checks: Vec<crate::predicate::LaunchCheck>,
 }
 
-// ── Global Device (process-wide, per-device singleton) ─────────────────────
+// ── Global Context (process-wide, per-device singleton) ─────────────────────
 
 /// Global per-device handles. Shared across all threads so that
 /// `Module`/`Function` loaded against a device can be used from any thread.
-static DEVICES: OnceLock<Mutex<HashMap<usize, Arc<Device>>>> = OnceLock::new();
+static DEVICES: OnceLock<Mutex<HashMap<usize, Arc<Context>>>> = OnceLock::new();
 
-fn devices() -> &'static Mutex<HashMap<usize, Arc<Device>>> {
+fn devices() -> &'static Mutex<HashMap<usize, Arc<Context>>> {
     DEVICES.get_or_init(|| Mutex::new(HashMap::new()))
 }
 
-/// Get or create the global [`Device`] for a device ordinal.
+/// Get or create the global [`Context`] for a device ordinal.
 ///
 /// The first call for a given `device_id` creates the device handle; subsequent
-/// calls return the same `Arc<Device>`.
-fn get_or_init_device(device_id: usize) -> Result<Arc<Device>, DeviceError> {
+/// calls return the same `Arc<Context>`.
+fn get_or_init_device(device_id: usize) -> Result<Arc<Context>, DeviceError> {
     let mut devices = devices()
         .lock()
         .map_err(|_| device_error(device_id, "device map lock poisoned"))?;
     if let Some(device) = devices.get(&device_id) {
         return Ok(Arc::clone(device));
     }
-    let device = Device::new(device_id)?;
+    let device = Context::new(device_id)?;
     devices.insert(device_id, Arc::clone(&device));
     Ok(device)
 }
@@ -243,7 +243,7 @@ pub fn init_device(
 ) -> Result<(), DeviceError> {
     let device_context = new_device_context(device_id, policy)?;
     let pred = hashmap.insert(device_id, device_context).is_none();
-    device_assert(device_id, pred, "Device is already initialized.")
+    device_assert(device_id, pred, "Context is already initialized.")
 }
 
 pub fn init_with_default_policy(
@@ -260,7 +260,7 @@ pub fn init_with_default_policy(
         pool: None,
     };
     let pred = hashmap.insert(device_id, device_context).is_none();
-    device_assert(device_id, pred, "Device is already initialized.")
+    device_assert(device_id, pred, "Context is already initialized.")
 }
 
 /// Holds the thread's context map while a closure runs against it, and puts it
@@ -371,10 +371,10 @@ where
     })
 }
 
-/// Run a closure with a reference to the [`Device`] for `device_id`.
+/// Run a closure with a reference to the [`Context`] for `device_id`.
 pub fn with_device<F, R>(device_id: usize, f: F) -> Result<R, DeviceError>
 where
-    F: FnOnce(&Arc<Device>) -> R,
+    F: FnOnce(&Arc<Context>) -> R,
 {
     let device = get_or_init_device(device_id)?;
     Ok(f(&device))
@@ -429,7 +429,7 @@ pub fn set_default_device(default_device_id: usize) {
 /// Returns [`DeviceError::Context`](crate::error::DeviceError::Context) if
 /// `pool` was created on a different device than `device_id`.
 pub fn set_device_pool(device_id: usize, pool: Arc<MemPool>) -> Result<(), DeviceError> {
-    let pool_device = pool.device().ordinal();
+    let pool_device = pool.context().ordinal();
     device_assert(
         device_id,
         pool_device == device_id,
@@ -467,7 +467,10 @@ pub fn get_device_pool(device_id: usize) -> Result<Option<Arc<MemPool>>, DeviceE
 /// Errors from the device-context lookup are downgraded to `None`; this is the
 /// single choke-point for that decision so callers don't each re-derive it.
 pub fn pool_for_stream(stream: &Arc<Stream>) -> Option<Arc<MemPool>> {
-    get_device_pool(stream.device().ordinal()).ok().flatten()
+    get_device_pool(stream.context().ordinal())
+        .ok()
+        .flatten()
+        .filter(|pool| pool.context().as_ref() == stream.context().as_ref())
 }
 
 /// Run a closure with the scheduling policy of the current thread's default device.
@@ -495,12 +498,12 @@ pub fn load_module_from_file(filename: &str, device_id: usize) -> Result<Arc<Mod
 }
 
 /// Load a compiled CUDA module from an in-memory **cubin** image. Not for PTX
-/// (see [`cuda_core::Device::load_module_from_bytes`]); use
+/// (see [`cuda_core::Context::load_module_from_bytes`]); use
 /// [`load_module_from_ptx`] for that.
 ///
 /// # Safety
 ///
-/// Same contract as [`cuda_core::Device::load_module_from_bytes`]: `image`
+/// Same contract as [`cuda_core::Context::load_module_from_bytes`]: `image`
 /// must be a complete, well-formed cubin. The driver dereferences the offsets
 /// declared in the image header, so a truncated or malformed image is read
 /// past the end of the slice.