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;