Skip to main content

oxicuda_driver/
stream_ordered_alloc.rs

1//! Stream-ordered memory allocation (CUDA 11.2+ / 12.x+).
2//!
3//! Stream-ordered memory allocation allows memory operations (`alloc` / `free`)
4//! to participate in the stream execution order, eliminating the need for
5//! explicit synchronisation between allocation and kernel launch.
6//!
7//! This module provides:
8//!
9//! * [`StreamMemoryPool`] — a memory pool bound to a specific device.
10//! * [`StreamAllocation`] — a handle to a stream-ordered allocation.
11//! * [`StreamOrderedAllocConfig`] — pool configuration (sizes, thresholds).
12//! * [`PoolAttribute`] / [`PoolUsageStats`] — attribute queries and statistics.
13//! * [`PoolExportDescriptor`] / [`ShareableHandleType`] — IPC sharing metadata.
14//! * [`stream_alloc`] / [`stream_free`] — convenience free functions.
15//!
16//! # The stream-ordered model
17//!
18//! Independently of the GPU driver, every pool carries a faithful CPU
19//! simulation of the stream-ordered allocator (see
20//! [`stream_ordered_model`](crate::stream_ordered_model)).  Allocations
21//! (`alloc_async` / [`alloc_on`](StreamMemoryPool::alloc_on)) and frees
22//! (`free_async` / [`free_on`](StreamMemoryPool::free_on)) are sequenced per
23//! stream; freed blocks return to the pool once their stream reaches the free
24//! point and are then **reused** by a later same-or-larger request.  This makes
25//! the allocator's semantics — visibility ordering, block reuse, and
26//! reserved-vs-used accounting — observable on a plain CPU, with no GPU
27//! required.
28//!
29//! # Platform behaviour
30//!
31//! On platforms with a real CUDA driver (Linux, Windows),
32//! [`StreamMemoryPool::new`] additionally creates a driver-side pool via
33//! `cuMemPoolCreate`.  The lower-level `cuMem*Async` bindings remain available
34//! for real-GPU use.  [`StreamMemoryPool::cpu_pool`] builds a pool backed only
35//! by the CPU model, so the full stream-ordered API can be exercised without a
36//! driver on any platform.
37//!
38//! # Example
39//!
40//! ```rust
41//! use oxicuda_driver::stream_ordered_alloc::*;
42//! use oxicuda_driver::StreamOrderId;
43//!
44//! // A pool backed by the faithful CPU model — no GPU driver required.
45//! let config = StreamOrderedAllocConfig::default_for_device(0);
46//! let mut pool = StreamMemoryPool::cpu_pool(config)?;
47//!
48//! // A genuine stream-ordering identity.  In a GPU program this is derived
49//! // from a real `Stream` via `StreamMemoryPool::stream_id`; the model
50//! // sequences this token exactly like a real stream would.
51//! let stream = StreamOrderId::from(1);
52//!
53//! let mut alloc = pool.alloc_on(1024, stream)?;
54//! assert_eq!(alloc.size(), 1024);
55//! assert!(!alloc.is_freed());
56//!
57//! pool.free_on(&mut alloc, stream)?;
58//! assert!(alloc.is_freed());
59//! # Ok::<(), oxicuda_driver::CudaError>(())
60//! ```
61
62use std::fmt;
63
64use crate::error::{CudaError, CudaResult};
65use crate::ffi::CUdeviceptr;
66use crate::stream::Stream;
67use crate::stream_ordered_model::{ModelLimits, StreamOrderId, StreamOrderModel};
68
69// ---------------------------------------------------------------------------
70// Constants — CUmemPoolAttribute (mirrors CUDA header values)
71// ---------------------------------------------------------------------------
72
73/// Pool reuse policy: follow event dependencies.
74pub const CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES: u32 = 1;
75/// Pool reuse policy: allow opportunistic reuse.
76pub const CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC: u32 = 2;
77/// Pool reuse policy: allow internal dependency insertion.
78pub const CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES: u32 = 3;
79/// Release threshold in bytes (memory returned to OS when usage drops below).
80pub const CU_MEMPOOL_ATTR_RELEASE_THRESHOLD: u32 = 4;
81/// Current reserved memory (bytes) — read-only.
82pub const CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT: u32 = 5;
83/// High-water mark of reserved memory (bytes) — resettable.
84pub const CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH: u32 = 6;
85/// Current used memory (bytes) — read-only.
86pub const CU_MEMPOOL_ATTR_USED_MEM_CURRENT: u32 = 7;
87/// High-water mark of used memory (bytes) — resettable.
88pub const CU_MEMPOOL_ATTR_USED_MEM_HIGH: u32 = 8;
89
90// ---------------------------------------------------------------------------
91// StreamOrderedAllocConfig
92// ---------------------------------------------------------------------------
93
94/// Configuration for a stream-ordered memory pool.
95///
96/// All sizes are in bytes.
97#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct StreamOrderedAllocConfig {
99    /// Initial pool size in bytes.  The pool pre-reserves this amount of
100    /// device memory when created.
101    pub initial_pool_size: usize,
102
103    /// Maximum pool size in bytes.  `0` means unlimited — the pool will grow
104    /// as needed (subject to device memory limits).
105    pub max_pool_size: usize,
106
107    /// Release threshold in bytes.  When the pool is trimmed, at least this
108    /// much memory is kept reserved for future allocations.
109    pub release_threshold: usize,
110
111    /// The device ordinal to create the pool on.
112    pub device: i32,
113}
114
115impl StreamOrderedAllocConfig {
116    /// Validate that the configuration is internally consistent.
117    ///
118    /// # Rules
119    ///
120    /// * `initial_pool_size` must not exceed `max_pool_size` (when
121    ///   `max_pool_size > 0`).
122    /// * `release_threshold` must not exceed `max_pool_size` (when
123    ///   `max_pool_size > 0`).
124    /// * `device` must be non-negative.
125    ///
126    /// # Errors
127    ///
128    /// Returns [`CudaError::InvalidValue`] if any rule is violated.
129    pub fn validate(&self) -> CudaResult<()> {
130        if self.device < 0 {
131            return Err(CudaError::InvalidValue);
132        }
133
134        if self.max_pool_size > 0 {
135            if self.initial_pool_size > self.max_pool_size {
136                return Err(CudaError::InvalidValue);
137            }
138            if self.release_threshold > self.max_pool_size {
139                return Err(CudaError::InvalidValue);
140            }
141        }
142
143        Ok(())
144    }
145
146    /// Returns a sensible default configuration for the given device.
147    ///
148    /// * `initial_pool_size` = 0 (grow on demand)
149    /// * `max_pool_size` = 0 (unlimited)
150    /// * `release_threshold` = 0 (release everything on trim)
151    pub fn default_for_device(device: i32) -> Self {
152        Self {
153            initial_pool_size: 0,
154            max_pool_size: 0,
155            release_threshold: 0,
156            device,
157        }
158    }
159}
160
161// ---------------------------------------------------------------------------
162// PoolAttribute
163// ---------------------------------------------------------------------------
164
165/// Attributes that can be queried or set on a [`StreamMemoryPool`].
166#[derive(Debug, Clone, Copy, PartialEq, Eq)]
167pub enum PoolAttribute {
168    /// Whether freed blocks can be reused by following event dependencies.
169    ReuseFollowEventDependencies,
170    /// Whether freed blocks can be opportunistically reused (without ordering).
171    ReuseAllowOpportunistic,
172    /// Whether the pool may insert internal dependencies for reuse.
173    ReuseAllowInternalDependencies,
174    /// The release threshold in bytes.
175    ReleaseThreshold(u64),
176    /// Current reserved memory (read-only query).
177    ReservedMemCurrent,
178    /// High-water mark of reserved memory.
179    ReservedMemHigh,
180    /// Current used memory (read-only query).
181    UsedMemCurrent,
182    /// High-water mark of used memory.
183    UsedMemHigh,
184}
185
186impl PoolAttribute {
187    /// Convert to the raw CUDA attribute constant.
188    pub fn to_raw(self) -> u32 {
189        match self {
190            Self::ReuseFollowEventDependencies => CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES,
191            Self::ReuseAllowOpportunistic => CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC,
192            Self::ReuseAllowInternalDependencies => {
193                CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES
194            }
195            Self::ReleaseThreshold(_) => CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
196            Self::ReservedMemCurrent => CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
197            Self::ReservedMemHigh => CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
198            Self::UsedMemCurrent => CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
199            Self::UsedMemHigh => CU_MEMPOOL_ATTR_USED_MEM_HIGH,
200        }
201    }
202}
203
204// ---------------------------------------------------------------------------
205// PoolUsageStats
206// ---------------------------------------------------------------------------
207
208/// Snapshot of pool memory usage.
209#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
210pub struct PoolUsageStats {
211    /// Bytes currently reserved from the device allocator.
212    pub reserved_current: u64,
213    /// Peak bytes reserved (since creation or last reset).
214    pub reserved_high: u64,
215    /// Bytes currently in use by outstanding allocations.
216    pub used_current: u64,
217    /// Peak bytes in use (since creation or last reset).
218    pub used_high: u64,
219    /// Number of active (not-yet-freed) allocations.
220    pub active_allocations: usize,
221    /// Peak number of concurrent allocations.
222    pub peak_allocations: usize,
223}
224
225// ---------------------------------------------------------------------------
226// ShareableHandleType / PoolExportDescriptor
227// ---------------------------------------------------------------------------
228
229/// Handle type used for IPC sharing of memory pools.
230#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
231pub enum ShareableHandleType {
232    /// No sharing.
233    #[default]
234    None,
235    /// POSIX file descriptor (Linux).
236    PosixFileDescriptor,
237    /// Win32 handle (Windows).
238    Win32Handle,
239    /// Win32 KMT handle (Windows, legacy).
240    Win32KmtHandle,
241}
242
243/// Descriptor for exporting a pool for IPC sharing.
244#[derive(Debug, Clone, Copy, PartialEq, Eq)]
245pub struct PoolExportDescriptor {
246    /// The handle type to use for sharing.
247    pub shareable_handle_type: ShareableHandleType,
248    /// The device ordinal that owns the pool.
249    pub pool_device: i32,
250}
251
252// ---------------------------------------------------------------------------
253// StreamAllocation
254// ---------------------------------------------------------------------------
255
256/// Handle to a stream-ordered memory allocation.
257///
258/// The allocation is produced by the pool's faithful CPU model — the source of
259/// truth for byte accounting, block reuse, and per-stream ordering. It is
260/// associated with a specific stream and memory pool: it becomes available when
261/// all preceding work on the stream has completed, and is returned to the pool
262/// when freed (also stream-ordered).
263///
264/// # Pointer semantics
265///
266/// [`as_ptr`](StreamAllocation::as_ptr) returns the *model* address, which is
267/// the allocator's identity token — **not** a device pointer you can pass to a
268/// kernel. Obtaining a real on-device `CUdeviceptr` requires the driver-backed
269/// `cuMemAllocAsync` / `cuMemAllocFromPoolAsync` path (see the module's
270/// `gpu_*` bindings); the model address is stable and unique for reuse and
271/// ordering bookkeeping only.
272pub struct StreamAllocation {
273    /// Model address (allocator identity token), typed as `CUdeviceptr`. This
274    /// is **not** a dereferenceable on-device pointer — see the type docs.
275    ptr: CUdeviceptr,
276    /// Size of the allocation in bytes.
277    size: usize,
278    /// The stream this allocation is ordered on (raw ordering token).
279    stream: u64,
280    /// The pool handle that owns this allocation.
281    pool: u64,
282    /// Sequence number at which the allocation becomes valid on its stream,
283    /// in the owning pool's [`StreamOrderModel`].
284    ready_seq: u64,
285    /// Whether this allocation has already been freed.
286    freed: bool,
287}
288
289impl StreamAllocation {
290    /// Returns the allocator's *model* address as a raw `u64`.
291    ///
292    /// This is a stable, unique identity token used for reuse and ordering
293    /// bookkeeping — **not** a dereferenceable on-device `CUdeviceptr`. A real
294    /// device pointer must be obtained via the driver-backed `cuMemAllocAsync`
295    /// path (see the type-level documentation).
296    #[inline]
297    pub fn as_ptr(&self) -> u64 {
298        self.ptr
299    }
300
301    /// Returns the allocation size in bytes.
302    #[inline]
303    pub fn size(&self) -> usize {
304        self.size
305    }
306
307    /// Returns `true` if this allocation has been freed.
308    #[inline]
309    pub fn is_freed(&self) -> bool {
310        self.freed
311    }
312
313    /// Returns the stream handle this allocation is ordered on.
314    #[inline]
315    pub fn stream(&self) -> u64 {
316        self.stream
317    }
318
319    /// Returns the ordering identifier of the stream this allocation is bound
320    /// to in the owning pool's stream-ordered model.
321    #[inline]
322    pub fn stream_id(&self) -> StreamOrderId {
323        StreamOrderId(self.stream)
324    }
325
326    /// Returns the sequence number at which this allocation becomes valid on
327    /// its stream within the owning pool's [`StreamOrderModel`].
328    ///
329    /// The allocation is safe to read on its own stream only once that stream
330    /// has executed past this point (queryable via
331    /// [`StreamMemoryPool::is_ready`]).
332    #[inline]
333    pub fn ready_seq(&self) -> u64 {
334        self.ready_seq
335    }
336
337    /// Returns the pool handle that owns this allocation.
338    #[inline]
339    pub fn pool(&self) -> u64 {
340        self.pool
341    }
342}
343
344impl fmt::Debug for StreamAllocation {
345    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
346        f.debug_struct("StreamAllocation")
347            .field("ptr", &format_args!("0x{:016x}", self.ptr))
348            .field("size", &self.size)
349            .field("stream", &format_args!("0x{:016x}", self.stream))
350            .field("freed", &self.freed)
351            .finish()
352    }
353}
354
355// ---------------------------------------------------------------------------
356// StreamMemoryPool
357// ---------------------------------------------------------------------------
358
359/// A memory pool for stream-ordered allocations.
360///
361/// Every pool drives a faithful CPU model of the stream-ordered allocator (the
362/// source of truth for byte accounting, block reuse, and per-stream ordering).
363/// On platforms with a real CUDA driver (Linux, Windows),
364/// [`StreamMemoryPool::new`] *additionally* creates a driver-side pool via
365/// `cuMemPoolCreate`.  [`StreamMemoryPool::cpu_pool`] builds a pool backed only
366/// by the CPU model and never touches the driver, so the API can be exercised
367/// on any platform.
368///
369/// # Allocation tracking
370///
371/// The pool maintains running allocation counts and byte totals (mirrored from
372/// the CPU model) for diagnostics; these are available everywhere via
373/// [`StreamMemoryPool::stats`].
374pub struct StreamMemoryPool {
375    /// Raw `CUmemoryPool` handle (0 if not backed by a real driver pool).
376    handle: u64,
377    /// Whether this wrapper *owns* the driver-side pool and must destroy it on
378    /// drop. `true` only for pools created by [`StreamMemoryPool::new`] via
379    /// `cuMemPoolCreate`; `false` for the CPU model, and crucially for the
380    /// driver-owned default pool from [`StreamMemoryPool::default_pool`] (which
381    /// must never be passed to `cuMemPoolDestroy`).
382    owned: bool,
383    /// Device ordinal.
384    device: i32,
385    /// Configuration used to create this pool.
386    config: StreamOrderedAllocConfig,
387    /// Number of currently active (not freed) allocations (mirror of the
388    /// model's live count, kept for cheap field access).
389    active_allocations: usize,
390    /// Total bytes currently in use (mirror of the model's `used`).
391    total_allocated: usize,
392    /// Peak bytes ever in use concurrently (mirror of the model's `used_high`).
393    peak_allocated: usize,
394    /// Peak number of concurrent allocations (mirror of the model's peak).
395    peak_allocation_count: usize,
396    /// Faithful CPU model of the stream-ordered allocator.  This is the
397    /// authority for pointers, reuse, and stream ordering on every platform.
398    model: StreamOrderModel,
399}
400
401impl fmt::Debug for StreamMemoryPool {
402    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
403        f.debug_struct("StreamMemoryPool")
404            .field("handle", &format_args!("0x{:016x}", self.handle))
405            .field("device", &self.device)
406            .field("active_allocations", &self.active_allocations)
407            .field("total_allocated", &self.total_allocated)
408            .field("reserved", &self.model.reserved())
409            .finish()
410    }
411}
412
413impl StreamMemoryPool {
414    /// Create a new memory pool for the given device.
415    ///
416    /// The configuration is validated and the CPU model is initialised.  On
417    /// platforms with a real CUDA driver, `cuMemPoolCreate` is also invoked and
418    /// its handle stored; without a driver this fails cleanly.
419    ///
420    /// To obtain a pool that never touches the driver (e.g. for CPU-only use of
421    /// the stream-ordered API), use [`StreamMemoryPool::cpu_pool`].
422    ///
423    /// # Errors
424    ///
425    /// * [`CudaError::InvalidValue`] if the config fails validation.
426    /// * On non-macOS, any [`CudaError`] from `cuMemPoolCreate` (e.g.
427    ///   [`CudaError::NotInitialized`] when no driver is loadable).
428    pub fn new(config: StreamOrderedAllocConfig) -> CudaResult<Self> {
429        config.validate()?;
430
431        #[cfg_attr(target_os = "macos", allow(unused_mut))]
432        let mut pool = Self::with_model(config);
433
434        // On real GPU platforms, create the driver-side pool via
435        // `cuMemPoolCreate` and store the returned handle.  When the driver
436        // is absent the call returns `Err` and pool creation fails cleanly.
437        #[cfg(not(target_os = "macos"))]
438        {
439            pool.handle = Self::gpu_create_pool(&pool.config)?;
440            // We created the driver-side pool, so we own it and must destroy
441            // it on drop.
442            pool.owned = true;
443        }
444
445        Ok(pool)
446    }
447
448    /// Create a pool backed solely by the faithful CPU model, without touching
449    /// the CUDA driver.
450    ///
451    /// This always succeeds (given a valid configuration) on every platform and
452    /// is the recommended entry point for using the stream-ordered allocation
453    /// semantics on a CPU.
454    ///
455    /// # Errors
456    ///
457    /// * [`CudaError::InvalidValue`] if the config fails validation.
458    pub fn cpu_pool(config: StreamOrderedAllocConfig) -> CudaResult<Self> {
459        config.validate()?;
460        Ok(Self::with_model(config))
461    }
462
463    /// Build a pool value with a fresh CPU model (no driver interaction).
464    fn with_model(config: StreamOrderedAllocConfig) -> Self {
465        let model = StreamOrderModel::new(Self::model_limits(&config));
466        Self {
467            handle: 0,
468            owned: false,
469            device: config.device,
470            config,
471            active_allocations: 0,
472            total_allocated: 0,
473            peak_allocated: 0,
474            peak_allocation_count: 0,
475            model,
476        }
477    }
478
479    /// Build the model limits from a pool configuration.
480    fn model_limits(config: &StreamOrderedAllocConfig) -> ModelLimits {
481        ModelLimits {
482            max_pool_size: config.max_pool_size,
483            release_threshold: config.release_threshold,
484        }
485    }
486
487    /// Derive a stream-ordering identifier from a genuine [`Stream`].
488    ///
489    /// The stream's raw handle is a stable, unique token for the lifetime of
490    /// the stream, which the CPU model uses as the stream's ordering identity.
491    #[inline]
492    pub fn stream_id(stream: &Stream) -> StreamOrderId {
493        StreamOrderId(stream.raw().0 as usize as u64)
494    }
495
496    /// Allocate memory on a stream (stream-ordered), identified by a raw
497    /// ordering token.
498    ///
499    /// The allocation becomes valid on the stream once the stream reaches the
500    /// allocation point.  A previously-freed block of the same-or-larger size
501    /// is reused when available; otherwise a fresh block is carved from the
502    /// pool.
503    ///
504    /// # Errors
505    ///
506    /// * [`CudaError::InvalidValue`] if `size` is zero.
507    /// * [`CudaError::OutOfMemory`] if `max_pool_size` would be exceeded.
508    pub fn alloc_async(&mut self, size: usize, stream: u64) -> CudaResult<StreamAllocation> {
509        self.alloc_on(size, StreamOrderId(stream))
510    }
511
512    /// Allocate memory ordered on a genuine [`Stream`].
513    ///
514    /// This is the recommended entry point when a real CUDA [`Stream`] is
515    /// available: the allocation is sequenced against that exact stream in the
516    /// pool's [`StreamOrderModel`] (see [`StreamMemoryPool::stream_id`]).
517    ///
518    /// # Errors
519    ///
520    /// * [`CudaError::InvalidValue`] if `size` is zero.
521    /// * [`CudaError::OutOfMemory`] if `max_pool_size` would be exceeded.
522    pub fn alloc_async_on_stream(
523        &mut self,
524        size: usize,
525        stream: &Stream,
526    ) -> CudaResult<StreamAllocation> {
527        self.alloc_on(size, Self::stream_id(stream))
528    }
529
530    /// Allocate memory ordered on the stream identified by `stream`.
531    ///
532    /// The block is carved from the pool — reusing a previously-freed block of
533    /// the same-or-larger size when one is available — and sequenced on the
534    /// stream so that it only becomes valid once the stream reaches the
535    /// allocation point (queryable via [`StreamMemoryPool::is_ready`]).
536    ///
537    /// # Errors
538    ///
539    /// * [`CudaError::InvalidValue`] if `size` is zero.
540    /// * [`CudaError::OutOfMemory`] if `max_pool_size` would be exceeded.
541    pub fn alloc_on(&mut self, size: usize, stream: StreamOrderId) -> CudaResult<StreamAllocation> {
542        let model_alloc = self.model.alloc(size, stream)?;
543        self.sync_mirror_stats();
544
545        Ok(StreamAllocation {
546            ptr: model_alloc.ptr,
547            size: model_alloc.size,
548            stream: stream.raw(),
549            pool: self.handle,
550            ready_seq: model_alloc.ready_seq,
551            freed: false,
552        })
553    }
554
555    /// Free memory on a stream (stream-ordered).
556    ///
557    /// The memory is returned to the pool once all prior work on the
558    /// allocation's stream has completed.  The allocation is marked freed and
559    /// cannot be freed again.
560    ///
561    /// # Errors
562    ///
563    /// * [`CudaError::InvalidValue`] if the allocation is already freed, or its
564    ///   pointer is not live in this pool (foreign-pointer free).
565    pub fn free_async(&mut self, alloc: &mut StreamAllocation) -> CudaResult<()> {
566        let stream = alloc.stream_id();
567        self.free_on(alloc, stream)
568    }
569
570    /// Free `alloc` ordered on a genuine [`Stream`].
571    ///
572    /// CUDA permits freeing on a stream different from the one the allocation
573    /// was made on; the free still completes only once *that* stream reaches
574    /// the free point.
575    ///
576    /// # Errors
577    ///
578    /// * [`CudaError::InvalidValue`] if the allocation is already freed or its
579    ///   pointer is not live in this pool.
580    pub fn free_async_on_stream(
581        &mut self,
582        alloc: &mut StreamAllocation,
583        stream: &Stream,
584    ) -> CudaResult<()> {
585        self.free_on(alloc, Self::stream_id(stream))
586    }
587
588    /// Free `alloc` ordered on the stream identified by `stream`.
589    ///
590    /// # Errors
591    ///
592    /// * [`CudaError::InvalidValue`] if the allocation is already freed or its
593    ///   pointer is not live in this pool.
594    pub fn free_on(
595        &mut self,
596        alloc: &mut StreamAllocation,
597        stream: StreamOrderId,
598    ) -> CudaResult<()> {
599        if alloc.freed {
600            return Err(CudaError::InvalidValue);
601        }
602
603        self.model.free(alloc.ptr, stream)?;
604        self.sync_mirror_stats();
605
606        alloc.freed = true;
607        Ok(())
608    }
609
610    /// Advance a stream to its head (model of `cuStreamSynchronize`),
611    /// completing every operation submitted on it so far and reclaiming any
612    /// completed stream-ordered frees into the pool for reuse.
613    pub fn synchronize_stream(&mut self, stream: StreamOrderId) {
614        self.model.synchronize(stream);
615        self.sync_mirror_stats();
616    }
617
618    /// Returns `true` if `alloc` is valid to read on its own ordering stream,
619    /// i.e. that stream has executed past the allocation point.
620    pub fn is_ready(&self, alloc: &StreamAllocation) -> bool {
621        let model_alloc = crate::stream_ordered_model::ModelAllocation {
622            ptr: alloc.ptr,
623            size: alloc.size,
624            capacity: alloc.size,
625            stream: alloc.stream_id(),
626            ready_seq: alloc.ready_seq,
627        };
628        self.model.is_ready_same_stream(&model_alloc)
629    }
630
631    /// Returns `true` if `alloc` (made on its own stream) is safe to use on
632    /// `consumer` given that `consumer` was ordered after `wait_seq` on the
633    /// allocation's stream (the sequence captured by an event it waited on).
634    ///
635    /// Use [`StreamMemoryPool::record_event`] on the producing stream to obtain
636    /// a `wait_seq` that captures the allocation.
637    pub fn is_ready_on(
638        &self,
639        alloc: &StreamAllocation,
640        consumer: StreamOrderId,
641        wait_seq: u64,
642    ) -> bool {
643        let model_alloc = crate::stream_ordered_model::ModelAllocation {
644            ptr: alloc.ptr,
645            size: alloc.size,
646            capacity: alloc.size,
647            stream: alloc.stream_id(),
648            ready_seq: alloc.ready_seq,
649        };
650        self.model
651            .is_ready_cross_stream(&model_alloc, consumer, wait_seq)
652    }
653
654    /// Record an event on `stream`, returning the sequence number it captures.
655    ///
656    /// A later cross-stream wait on this value orders the waiting stream after
657    /// every operation submitted on `stream` before this point.
658    pub fn record_event(&mut self, stream: StreamOrderId) -> u64 {
659        self.model.record_event(stream)
660    }
661
662    /// Trim the CPU model's pool, releasing free-list bytes above
663    /// `min_bytes_to_keep` back to the (virtual) device.
664    ///
665    /// This is the CPU-model analogue of `cuMemPoolTrimTo`; it always succeeds
666    /// and is available on every platform.  For the raw driver trim, see the
667    /// platform-gated [`StreamMemoryPool::trim`].
668    pub fn model_trim(&mut self, min_bytes_to_keep: usize) {
669        self.model.trim_to(min_bytes_to_keep);
670        self.sync_mirror_stats();
671    }
672
673    /// Trim the driver-side pool, releasing unused memory back to the OS.
674    ///
675    /// At least `min_bytes_to_keep` bytes of reserved memory remain in the
676    /// pool.  This drives the real `cuMemPoolTrimTo` binding; for the
677    /// always-available CPU-model trim, use [`StreamMemoryPool::model_trim`].
678    ///
679    /// # Errors
680    ///
681    /// * [`CudaError::NotSupported`] on macOS.
682    /// * Any [`CudaError`] from `cuMemPoolTrimTo`.
683    pub fn trim(&mut self, min_bytes_to_keep: usize) -> CudaResult<()> {
684        self.platform_trim(min_bytes_to_keep)
685    }
686
687    /// Get pool usage statistics from the CPU model.
688    ///
689    /// `reserved_*` reflects everything the pool has carved from the (virtual)
690    /// device (live + reusable + pending-free bytes), whereas `used_*` reflects
691    /// only currently-live allocations.
692    pub fn stats(&self) -> PoolUsageStats {
693        PoolUsageStats {
694            reserved_current: self.model.reserved() as u64,
695            reserved_high: self.model.reserved_high() as u64,
696            used_current: self.model.used() as u64,
697            used_high: self.model.used_high() as u64,
698            active_allocations: self.model.active(),
699            peak_allocations: self.model.peak_active(),
700        }
701    }
702
703    /// Set a pool attribute.
704    ///
705    /// Only attributes that carry a value (e.g. [`PoolAttribute::ReleaseThreshold`])
706    /// modify pool state.  Read-only attributes (e.g. `ReservedMemCurrent`)
707    /// return [`CudaError::InvalidValue`].
708    ///
709    /// The release threshold is applied to the CPU model as well as (on
710    /// non-macOS) the driver pool.
711    ///
712    /// # Errors
713    ///
714    /// * [`CudaError::InvalidValue`] for read-only attributes.
715    /// * [`CudaError::NotSupported`] on macOS for non-threshold attributes.
716    pub fn set_attribute(&mut self, attr: PoolAttribute) -> CudaResult<()> {
717        // Read-only attributes cannot be set.
718        match attr {
719            PoolAttribute::ReservedMemCurrent
720            | PoolAttribute::UsedMemCurrent
721            | PoolAttribute::ReservedMemHigh
722            | PoolAttribute::UsedMemHigh => {
723                return Err(CudaError::InvalidValue);
724            }
725            _ => {}
726        }
727
728        // Apply locally-meaningful attributes to the config and CPU model.
729        if let PoolAttribute::ReleaseThreshold(val) = attr {
730            self.config.release_threshold = val as usize;
731            self.model.set_release_threshold(val as usize);
732        }
733
734        self.platform_set_attribute(attr)
735    }
736
737    /// Enable peer access from another device to allocations in this pool.
738    ///
739    /// After this call, kernels running on `peer_device` can access memory
740    /// allocated from this pool.
741    ///
742    /// # Errors
743    ///
744    /// * [`CudaError::InvalidDevice`] if `peer_device` equals this pool's device.
745    /// * [`CudaError::NotSupported`] on macOS.
746    pub fn enable_peer_access(&self, peer_device: i32) -> CudaResult<()> {
747        if peer_device == self.device {
748            return Err(CudaError::InvalidDevice);
749        }
750
751        self.platform_enable_peer_access(peer_device)
752    }
753
754    /// Disable peer access from another device to allocations in this pool.
755    ///
756    /// # Errors
757    ///
758    /// * [`CudaError::InvalidDevice`] if `peer_device` equals this pool's device.
759    /// * [`CudaError::NotSupported`] on macOS.
760    pub fn disable_peer_access(&self, peer_device: i32) -> CudaResult<()> {
761        if peer_device == self.device {
762            return Err(CudaError::InvalidDevice);
763        }
764
765        self.platform_disable_peer_access(peer_device)
766    }
767
768    /// Reset peak statistics (peak used bytes and peak allocation count).
769    pub fn reset_peak_stats(&mut self) {
770        self.model.reset_peaks();
771        self.sync_mirror_stats();
772    }
773
774    /// Mirror the model's current/peak figures into the cheap struct fields.
775    fn sync_mirror_stats(&mut self) {
776        self.active_allocations = self.model.active();
777        self.total_allocated = self.model.used();
778        self.peak_allocated = self.model.used_high();
779        self.peak_allocation_count = self.model.peak_active();
780    }
781
782    /// Get the default memory pool for a device.
783    ///
784    /// CUDA provides a default pool per device, queried via
785    /// `cuDeviceGetDefaultMemPool`.  The returned pool is owned by the
786    /// driver and is *not* destroyed when the [`StreamMemoryPool`] wrapper
787    /// is dropped.  On macOS, this returns a local-only pool with default
788    /// configuration.  In all cases the CPU model is initialised.
789    ///
790    /// # Errors
791    ///
792    /// * [`CudaError::InvalidValue`] if `device` is negative.
793    /// * [`CudaError::NotInitialized`] if the CUDA driver is not loaded.
794    /// * Any [`CudaError`] mapped from `cuDeviceGetDefaultMemPool`.
795    pub fn default_pool(device: i32) -> CudaResult<Self> {
796        if device < 0 {
797            return Err(CudaError::InvalidValue);
798        }
799
800        let config = StreamOrderedAllocConfig::default_for_device(device);
801
802        // On macOS there is no driver — fall back to a local-only pool.
803        #[cfg(target_os = "macos")]
804        {
805            Ok(Self::with_model(config))
806        }
807
808        // On real GPU platforms, resolve the device's default pool handle.
809        #[cfg(not(target_os = "macos"))]
810        {
811            let handle = Self::gpu_default_pool(device)?;
812            let mut pool = Self::with_model(config);
813            pool.handle = handle;
814            Ok(pool)
815        }
816    }
817
818    /// Returns the raw pool handle.
819    #[inline]
820    pub fn handle(&self) -> u64 {
821        self.handle
822    }
823
824    /// Returns the device ordinal.
825    #[inline]
826    pub fn device(&self) -> i32 {
827        self.device
828    }
829
830    /// Returns the pool configuration.
831    #[inline]
832    pub fn config(&self) -> &StreamOrderedAllocConfig {
833        &self.config
834    }
835
836    // -----------------------------------------------------------------------
837    // Platform-specific helpers (driver passthrough)
838    // -----------------------------------------------------------------------
839
840    /// Trim the driver pool on the current platform.
841    fn platform_trim(&mut self, min_bytes_to_keep: usize) -> CudaResult<()> {
842        #[cfg(target_os = "macos")]
843        {
844            let _ = min_bytes_to_keep;
845            Err(CudaError::NotSupported)
846        }
847
848        #[cfg(not(target_os = "macos"))]
849        {
850            Self::gpu_trim(self.handle, min_bytes_to_keep)
851        }
852    }
853
854    /// Set attribute on the driver pool on the current platform.
855    fn platform_set_attribute(&self, attr: PoolAttribute) -> CudaResult<()> {
856        #[cfg(target_os = "macos")]
857        {
858            match attr {
859                PoolAttribute::ReleaseThreshold(_) => Ok(()),
860                _ => Err(CudaError::NotSupported),
861            }
862        }
863
864        #[cfg(not(target_os = "macos"))]
865        {
866            Self::gpu_set_attribute(self.handle, attr)
867        }
868    }
869
870    /// Enable peer access on the driver pool on the current platform.
871    fn platform_enable_peer_access(&self, peer_device: i32) -> CudaResult<()> {
872        #[cfg(target_os = "macos")]
873        {
874            let _ = peer_device;
875            Err(CudaError::NotSupported)
876        }
877
878        #[cfg(not(target_os = "macos"))]
879        {
880            Self::gpu_enable_peer_access(self.handle, peer_device)
881        }
882    }
883
884    /// Disable peer access on the driver pool on the current platform.
885    fn platform_disable_peer_access(&self, peer_device: i32) -> CudaResult<()> {
886        #[cfg(target_os = "macos")]
887        {
888            let _ = peer_device;
889            Err(CudaError::NotSupported)
890        }
891
892        #[cfg(not(target_os = "macos"))]
893        {
894            Self::gpu_disable_peer_access(self.handle, peer_device)
895        }
896    }
897
898    // -----------------------------------------------------------------------
899    // GPU-only driver bindings (compiled out on macOS)
900    //
901    // These remain available for genuine GPU use.  They are *not* on the CPU
902    // model's hot path: the model is the allocator on every platform, and these
903    // bindings are exercised directly by the `gpu_*` tests against whatever
904    // driver the host provides.
905    // -----------------------------------------------------------------------
906
907    /// Create the pool on the GPU via `cuMemPoolCreate`.
908    ///
909    /// Builds a [`CUmemPoolProps`] from the pool configuration (pinned device
910    /// memory on `config.device`, `max_size` from `config.max_pool_size`),
911    /// invokes the driver, and returns the raw `CUmemoryPool` handle encoded
912    /// as a `u64`.
913    ///
914    /// When the driver is absent, [`try_driver`](crate::loader::try_driver)
915    /// returns `Err(CudaError::NotInitialized)` and pool creation fails
916    /// cleanly.  When the driver is present but predates CUDA 11.2 (no
917    /// `cuMemPoolCreate`), [`CudaError::NotSupported`] is returned.
918    #[cfg(not(target_os = "macos"))]
919    fn gpu_create_pool(config: &StreamOrderedAllocConfig) -> CudaResult<u64> {
920        use crate::ffi::{
921            CUmemAllocationType, CUmemLocation, CUmemLocationType, CUmemPoolProps, CUmemoryPool,
922        };
923
924        let api = crate::loader::try_driver()?;
925        let create = api.cu_mem_pool_create.ok_or(CudaError::NotSupported)?;
926
927        let props = CUmemPoolProps {
928            alloc_type: CUmemAllocationType::Pinned as u32,
929            handle_types: 0,
930            location: CUmemLocation {
931                loc_type: CUmemLocationType::Device as u32,
932                id: config.device,
933            },
934            win32_security_attributes: std::ptr::null_mut(),
935            max_size: config.max_pool_size,
936            reserved: [0u8; 56],
937        };
938
939        let mut pool = CUmemoryPool::default();
940        // SAFETY: `create` was just resolved from the driver; `props` and
941        // `pool` are valid, correctly-typed local variables, and the CUDA
942        // ABI's reserved padding is zeroed.
943        let rc = unsafe { create(&mut pool, &props) };
944        crate::error::check(rc)?;
945
946        Ok(pool.0 as usize as u64)
947    }
948
949    /// Resolve a device's default memory pool via `cuDeviceGetDefaultMemPool`.
950    #[cfg(not(target_os = "macos"))]
951    fn gpu_default_pool(device: i32) -> CudaResult<u64> {
952        use crate::ffi::CUmemoryPool;
953
954        let api = crate::loader::try_driver()?;
955        let get_default = api
956            .cu_device_get_default_mem_pool
957            .ok_or(CudaError::NotSupported)?;
958
959        let mut pool = CUmemoryPool::default();
960        // SAFETY: `get_default` was just resolved from the driver; `pool` is
961        // a valid local and `device` is a plain device ordinal.
962        let rc = unsafe { get_default(&mut pool, device) };
963        crate::error::check(rc)?;
964
965        Ok(pool.0 as usize as u64)
966    }
967
968    /// Allocate stream-ordered memory.
969    ///
970    /// When `pool_handle` is non-zero, allocates from that explicit pool via
971    /// `cuMemAllocFromPoolAsync`; when it is zero (default-pool semantics),
972    /// uses the context-wide `cuMemAllocAsync`.
973    ///
974    /// The CPU model is the allocator on every platform, so this real-driver
975    /// binding has no production caller; it is retained for genuine GPU use and
976    /// exercised directly by the `gpu_*` FFI tests.  `#[cfg_attr(not(test), …)]`
977    /// keeps the production lib build warning-free without removing the
978    /// real binding.
979    #[cfg(not(target_os = "macos"))]
980    #[cfg_attr(not(test), allow(dead_code))]
981    fn gpu_alloc_async(pool_handle: u64, size: usize, stream: u64) -> CudaResult<CUdeviceptr> {
982        use crate::ffi::{CUmemoryPool, CUstream};
983
984        let api = crate::loader::try_driver()?;
985        let cu_stream = CUstream(stream as usize as *mut std::ffi::c_void);
986        let mut dptr: CUdeviceptr = 0;
987
988        if pool_handle != 0 {
989            let alloc_from_pool = api
990                .cu_mem_alloc_from_pool_async
991                .ok_or(CudaError::NotSupported)?;
992            let pool = CUmemoryPool(pool_handle as usize as *mut std::ffi::c_void);
993            // SAFETY: `alloc_from_pool` was just resolved; `dptr` is a valid
994            // out-pointer and `pool`/`cu_stream` are reconstructed handles.
995            let rc = unsafe { alloc_from_pool(&mut dptr, size, pool, cu_stream) };
996            crate::error::check(rc)?;
997        } else {
998            let alloc_async = api.cu_mem_alloc_async.ok_or(CudaError::NotSupported)?;
999            // SAFETY: `alloc_async` was just resolved; `dptr` is a valid
1000            // out-pointer and `cu_stream` is a reconstructed handle.
1001            let rc = unsafe { alloc_async(&mut dptr, size, cu_stream) };
1002            crate::error::check(rc)?;
1003        }
1004
1005        Ok(dptr)
1006    }
1007
1008    /// Free stream-ordered memory via `cuMemFreeAsync`.
1009    ///
1010    /// Retained for genuine GPU use and exercised directly by the `gpu_*` FFI
1011    /// tests; the CPU model is the allocator on the production path.
1012    #[cfg(not(target_os = "macos"))]
1013    #[cfg_attr(not(test), allow(dead_code))]
1014    fn gpu_free_async(ptr: CUdeviceptr, stream: u64) -> CudaResult<()> {
1015        use crate::ffi::CUstream;
1016
1017        let api = crate::loader::try_driver()?;
1018        let free_async = api.cu_mem_free_async.ok_or(CudaError::NotSupported)?;
1019        let cu_stream = CUstream(stream as usize as *mut std::ffi::c_void);
1020        // SAFETY: `free_async` was just resolved from the driver; `ptr` is a
1021        // device pointer previously returned by an async allocation and
1022        // `cu_stream` is a reconstructed handle.
1023        crate::error::check(unsafe { free_async(ptr, cu_stream) })
1024    }
1025
1026    /// Trim the pool via `cuMemPoolTrimTo`.
1027    #[cfg(not(target_os = "macos"))]
1028    fn gpu_trim(pool_handle: u64, min_bytes_to_keep: usize) -> CudaResult<()> {
1029        use crate::ffi::CUmemoryPool;
1030
1031        let api = crate::loader::try_driver()?;
1032        let trim = api.cu_mem_pool_trim_to.ok_or(CudaError::NotSupported)?;
1033        let pool = CUmemoryPool(pool_handle as usize as *mut std::ffi::c_void);
1034        // SAFETY: `trim` was just resolved from the driver; `pool` is a
1035        // reconstructed pool handle and `min_bytes_to_keep` is a plain count.
1036        crate::error::check(unsafe { trim(pool, min_bytes_to_keep) })
1037    }
1038
1039    /// Set a pool attribute via `cuMemPoolSetAttribute`.
1040    ///
1041    /// The reuse-policy attributes carry an `int` value; the release
1042    /// threshold carries a `cuuint64_t`.  The value buffer is sized
1043    /// accordingly and passed to the driver.
1044    #[cfg(not(target_os = "macos"))]
1045    fn gpu_set_attribute(pool_handle: u64, attr: PoolAttribute) -> CudaResult<()> {
1046        use crate::ffi::CUmemoryPool;
1047
1048        let api = crate::loader::try_driver()?;
1049        let set_attr = api
1050            .cu_mem_pool_set_attribute
1051            .ok_or(CudaError::NotSupported)?;
1052        let pool = CUmemoryPool(pool_handle as usize as *mut std::ffi::c_void);
1053        let raw_attr = Self::map_pool_attribute(attr)?;
1054
1055        // The driver dereferences `value` as either `int` or `cuuint64_t`
1056        // depending on the attribute.  Stack-allocate the correct width.
1057        match attr {
1058            PoolAttribute::ReuseFollowEventDependencies
1059            | PoolAttribute::ReuseAllowOpportunistic
1060            | PoolAttribute::ReuseAllowInternalDependencies => {
1061                // Boolean-style reuse policies: enable (1) the policy.
1062                let mut value: std::ffi::c_int = 1;
1063                // SAFETY: `set_attr` was just resolved; `pool` is a
1064                // reconstructed handle and `value` is a valid `int` matching
1065                // the attribute's documented value type.
1066                let rc = unsafe {
1067                    set_attr(pool, raw_attr, (&mut value as *mut std::ffi::c_int).cast())
1068                };
1069                crate::error::check(rc)
1070            }
1071            PoolAttribute::ReleaseThreshold(threshold) => {
1072                let mut value: u64 = threshold;
1073                // SAFETY: `set_attr` was just resolved; `pool` is a
1074                // reconstructed handle and `value` is a valid `cuuint64_t`
1075                // matching the release-threshold value type.
1076                let rc = unsafe { set_attr(pool, raw_attr, (&mut value as *mut u64).cast()) };
1077                crate::error::check(rc)
1078            }
1079            // Read-only attributes are rejected before reaching this point.
1080            PoolAttribute::ReservedMemCurrent
1081            | PoolAttribute::ReservedMemHigh
1082            | PoolAttribute::UsedMemCurrent
1083            | PoolAttribute::UsedMemHigh => Err(CudaError::InvalidValue),
1084        }
1085    }
1086
1087    /// Map a [`PoolAttribute`] to the driver's [`CUmemPoolAttribute`].
1088    #[cfg(not(target_os = "macos"))]
1089    fn map_pool_attribute(attr: PoolAttribute) -> CudaResult<crate::ffi::CUmemPoolAttribute> {
1090        use crate::ffi::CUmemPoolAttribute;
1091        Ok(match attr {
1092            PoolAttribute::ReuseFollowEventDependencies => {
1093                CUmemPoolAttribute::ReuseFollowEventDependencies
1094            }
1095            PoolAttribute::ReuseAllowOpportunistic => CUmemPoolAttribute::ReuseAllowOpportunistic,
1096            PoolAttribute::ReuseAllowInternalDependencies => {
1097                CUmemPoolAttribute::ReuseAllowInternalDependencies
1098            }
1099            PoolAttribute::ReleaseThreshold(_) => CUmemPoolAttribute::ReleaseThreshold,
1100            PoolAttribute::ReservedMemCurrent => CUmemPoolAttribute::ReservedMemCurrent,
1101            PoolAttribute::ReservedMemHigh => CUmemPoolAttribute::ReservedMemHigh,
1102            PoolAttribute::UsedMemCurrent => CUmemPoolAttribute::UsedMemCurrent,
1103            PoolAttribute::UsedMemHigh => CUmemPoolAttribute::UsedMemHigh,
1104        })
1105    }
1106
1107    /// Enable peer access from `peer_device` via `cuMemPoolSetAccess`.
1108    ///
1109    /// Builds a [`CUmemAccessDesc`] granting read-write access to the peer
1110    /// device and applies it to the pool.
1111    #[cfg(not(target_os = "macos"))]
1112    fn gpu_enable_peer_access(pool_handle: u64, peer_device: i32) -> CudaResult<()> {
1113        Self::gpu_set_pool_access(pool_handle, peer_device, true)
1114    }
1115
1116    /// Disable peer access from `peer_device` via `cuMemPoolSetAccess`.
1117    #[cfg(not(target_os = "macos"))]
1118    fn gpu_disable_peer_access(pool_handle: u64, peer_device: i32) -> CudaResult<()> {
1119        Self::gpu_set_pool_access(pool_handle, peer_device, false)
1120    }
1121
1122    /// Shared implementation for enabling / disabling pool peer access.
1123    #[cfg(not(target_os = "macos"))]
1124    fn gpu_set_pool_access(pool_handle: u64, peer_device: i32, enable: bool) -> CudaResult<()> {
1125        use crate::ffi::{
1126            CUmemAccessDesc, CUmemAccessFlags, CUmemLocation, CUmemLocationType, CUmemoryPool,
1127        };
1128
1129        let api = crate::loader::try_driver()?;
1130        let set_access = api.cu_mem_pool_set_access.ok_or(CudaError::NotSupported)?;
1131        let pool = CUmemoryPool(pool_handle as usize as *mut std::ffi::c_void);
1132
1133        let flags = if enable {
1134            CUmemAccessFlags::ReadWrite
1135        } else {
1136            CUmemAccessFlags::None
1137        };
1138        let desc = CUmemAccessDesc {
1139            location: CUmemLocation {
1140                loc_type: CUmemLocationType::Device as u32,
1141                id: peer_device,
1142            },
1143            flags: flags as u32,
1144        };
1145
1146        // SAFETY: `set_access` was just resolved from the driver; `pool` is a
1147        // reconstructed handle and `desc` is a single valid descriptor.
1148        let rc = unsafe { set_access(pool, &desc, 1) };
1149        crate::error::check(rc)
1150    }
1151}
1152
1153impl Drop for StreamMemoryPool {
1154    /// Destroy the driver-side pool if this wrapper created (and therefore
1155    /// owns) it.
1156    ///
1157    /// Only pools built by [`StreamMemoryPool::new`] via `cuMemPoolCreate` are
1158    /// owned; the CPU model owns no driver resource, and the driver-owned
1159    /// default pool must never be destroyed. Errors are logged via
1160    /// `tracing::warn!` rather than propagated (destructors must not panic),
1161    /// mirroring the [`Stream`]/`Event`/`Module` drop pattern.
1162    fn drop(&mut self) {
1163        if !self.owned || self.handle == 0 {
1164            return;
1165        }
1166        if let Ok(api) = crate::loader::try_driver() {
1167            if let Some(destroy) = api.cu_mem_pool_destroy {
1168                let pool = crate::ffi::CUmemoryPool(self.handle as usize as *mut std::ffi::c_void);
1169                // SAFETY: `destroy` was just resolved from the driver; `pool`
1170                // is the handle returned by `cuMemPoolCreate` for this wrapper
1171                // and has not been destroyed yet (owned pools are destroyed
1172                // exactly once, here).
1173                let rc = unsafe { destroy(pool) };
1174                if rc != 0 {
1175                    tracing::warn!(
1176                        cuda_error = rc,
1177                        pool = format_args!("0x{:016x}", self.handle),
1178                        "cuMemPoolDestroy failed during StreamMemoryPool drop"
1179                    );
1180                }
1181            }
1182        }
1183    }
1184}
1185
1186// ---------------------------------------------------------------------------
1187// Convenience free functions
1188// ---------------------------------------------------------------------------
1189
1190/// Process-wide default pool backing [`stream_alloc`] / [`stream_free`].
1191///
1192/// Initialised lazily from [`StreamMemoryPool::default_pool`] for device 0 so
1193/// that byte accounting and block reuse persist across calls (rather than being
1194/// discarded with a throwaway pool on every allocation). The default pool is
1195/// driver-owned, so this wrapper does not destroy it. A failed initialisation
1196/// is *not* cached (the cell stays empty), allowing a later retry once a driver
1197/// becomes available.
1198static DEFAULT_POOL: std::sync::OnceLock<std::sync::Mutex<StreamMemoryPool>> =
1199    std::sync::OnceLock::new();
1200
1201/// Obtain (initialising on first use) the shared default pool.
1202fn shared_default_pool() -> CudaResult<&'static std::sync::Mutex<StreamMemoryPool>> {
1203    if let Some(pool) = DEFAULT_POOL.get() {
1204        return Ok(pool);
1205    }
1206    // Build the pool *before* touching the cell so a transient failure (e.g. no
1207    // driver yet) is not cached permanently.
1208    let pool = StreamMemoryPool::default_pool(0)?;
1209    Ok(DEFAULT_POOL.get_or_init(|| std::sync::Mutex::new(pool)))
1210}
1211
1212/// Allocate memory on a stream using a shared, process-wide default pool for
1213/// device 0.
1214///
1215/// Unlike a throwaway pool, the returned [`StreamAllocation`] stays accounted
1216/// for and can be returned to the same pool via [`stream_free`].
1217///
1218/// # Errors
1219///
1220/// Propagates errors from pool creation and allocation.
1221pub fn stream_alloc(size: usize, stream: u64) -> CudaResult<StreamAllocation> {
1222    let pool = shared_default_pool()?;
1223    // A poisoned lock only means a previous holder panicked; the pool state is
1224    // still consistent enough to keep serving accounting, so recover the guard.
1225    let mut guard = pool.lock().unwrap_or_else(|e| e.into_inner());
1226    guard.alloc_async(size, stream)
1227}
1228
1229/// Free a stream-ordered allocation.
1230///
1231/// If the allocation came from the shared default pool (see [`stream_alloc`]),
1232/// it is returned to that pool via [`StreamMemoryPool::free_on`] so the freed
1233/// block re-enters the reuse list and accounting stays consistent. Otherwise
1234/// the allocation is simply marked freed. For allocations from an explicit
1235/// pool, prefer [`StreamMemoryPool::free_on`] directly.
1236///
1237/// # Errors
1238///
1239/// * [`CudaError::InvalidValue`] if the allocation is already freed.
1240pub fn stream_free(alloc: &mut StreamAllocation) -> CudaResult<()> {
1241    if alloc.freed {
1242        return Err(CudaError::InvalidValue);
1243    }
1244
1245    if let Some(pool) = DEFAULT_POOL.get() {
1246        let mut guard = pool.lock().unwrap_or_else(|e| e.into_inner());
1247        if guard.handle() == alloc.pool {
1248            let stream = alloc.stream_id();
1249            return guard.free_on(alloc, stream);
1250        }
1251    }
1252
1253    alloc.freed = true;
1254    Ok(())
1255}
1256
1257// ---------------------------------------------------------------------------
1258// Tests
1259// ---------------------------------------------------------------------------
1260#[cfg(test)]
1261#[path = "stream_ordered_alloc_tests.rs"]
1262mod tests;