cubecl_runtime/memory_management/
mod.rs

1mod base;
2mod config;
3mod handle;
4mod layout;
5
6pub use base::*;
7pub use config::*;
8pub use handle::*;
9pub use layout::*;
10
11use alloc::vec::Vec;
12
13/// The type of memory pool to use.
14#[derive(Debug, Clone)]
15pub enum PoolType {
16    /// Use a memory where every allocation is a separate page.
17    ExclusivePages {
18        /// The minimum number of bytes to allocate in this pool.
19        max_alloc_size: u64,
20    },
21    /// Give every allocation its own device allocation, sized to the request
22    /// and reused by exact size.
23    ///
24    /// No carving at all, so it wastes only alignment padding. Worth it where
25    /// padding matters more than allocation count — under a
26    /// [`DryRun`](crate::dry_run::DryRun), where unresolved reservations never
27    /// reach the driver at all, or on a device the workload barely fits.
28    Direct {
29        /// Reserved bytes above which free slices are returned to the driver.
30        /// A watermark rather than a budget; `None` reclaims only on an
31        /// explicit cleanup.
32        reclaim_at: Option<u64>,
33    },
34    /// Use a memory where each allocation is a slice of a bigger allocation.
35    SlicedPages {
36        /// The page size to allocate.
37        page_size: u64,
38        /// The maximum size of a slice to allocate in the pool.
39        max_slice_size: u64,
40        /// Hard cap on the total bytes of pages this pool may hold.
41        ///
42        /// The effective cap is `floor(max_pool_size / page_size)` whole pages.
43        /// If `max_pool_size < page_size`, the page size is shrunk to the
44        /// (alignment-rounded) cap so the budget is honored with a single page.
45        /// When the cap is reached and no free slice fits after coalescing,
46        /// reserving returns [`IoError`](crate::server::IoError)
47        /// `PoolCapacityExceeded` instead of silently growing. `None` (the
48        /// previous behavior) keeps unbounded growth.
49        ///
50        /// Note: runtimes that create one memory management per stream (CUDA,
51        /// HIP) apply the cap per stream.
52        max_pool_size: Option<u64>,
53    },
54}
55
56/// Options to create a memory pool.
57#[derive(Debug, Clone)]
58pub struct MemoryPoolOptions {
59    /// What kind of pool to use.
60    pub pool_type: PoolType,
61    /// Period after which allocations are deemed unused and deallocated.
62    ///
63    /// This period is measured in the number of allocations in the parent allocator. If a page
64    /// in the pool was unused for the entire period, it will be deallocated. This period is
65    /// approximmate, as checks are only done occasionally.
66    pub dealloc_period: Option<u64>,
67}
68
69/// High level configuration of memory management.
70#[derive(Clone, Debug)]
71pub enum MemoryConfiguration {
72    /// The default preset, which uses pools that allocate sub slices.
73    #[cfg(not(exclusive_memory_only))]
74    SubSlices,
75    /// Default preset for using exclusive pages.
76    /// This can be necessary for backends don't support sub-slices.
77    ExclusivePages,
78    /// Custom settings.
79    Custom {
80        /// Options for each pool to construct. When allocating, the first
81        /// possible pool will be picked for an allocation.
82        pool_options: Vec<MemoryPoolOptions>,
83    },
84}
85
86#[allow(clippy::derivable_impls)]
87impl Default for MemoryConfiguration {
88    fn default() -> Self {
89        #[cfg(exclusive_memory_only)]
90        {
91            MemoryConfiguration::ExclusivePages
92        }
93        #[cfg(not(exclusive_memory_only))]
94        {
95            MemoryConfiguration::SubSlices
96        }
97    }
98}
99
100#[derive(Default, Clone, Copy, Debug)]
101/// The mode of allocation used.
102pub enum MemoryAllocationMode {
103    /// Use the automatic memory management strategy for allocation.
104    #[default]
105    Auto,
106    /// Use a persistent memory management strategy, meaning that all allocations are for data that is
107    /// likely never going to be freed.
108    Persistent,
109}
110
111/// Why installing a dynamic pool layout did not take effect.
112///
113/// The layout itself was already valid — that is
114/// [`PoolConfigError`](PoolConfigError), reported when the configuration is
115/// resolved. This is about the pools' *state* at the moment of the swap.
116#[derive(Debug, Clone, Copy, PartialEq, Eq)]
117pub enum InstallMemoryPoolsError {
118    /// The dynamic pools still hold live allocations, so the old layout was
119    /// kept. A live slice carries its pool position, and swapping the pool
120    /// list under it would leave that position pointing at a different pool.
121    ///
122    /// Transient: retry once whatever holds them drains. A cleanup that does
123    /// not clear it usually means a cache is holding slices (the metadata
124    /// info cache) or a captured graph is pinning them.
125    PoolsInUse {
126        /// Bytes still live in the dynamic pools.
127        bytes_in_use: u64,
128    },
129    /// This server has no configurable dynamic pools. Permanent — unlike
130    /// [`PoolsInUse`](Self::PoolsInUse), retrying will never succeed.
131    Unsupported,
132}
133
134impl core::fmt::Display for InstallMemoryPoolsError {
135    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
136        match self {
137            InstallMemoryPoolsError::PoolsInUse { bytes_in_use } => write!(
138                f,
139                "the dynamic pools kept their layout: {bytes_in_use} bytes are still live in them"
140            ),
141            InstallMemoryPoolsError::Unsupported => {
142                write!(f, "this server has no configurable dynamic memory pools")
143            }
144        }
145    }
146}
147
148impl core::error::Error for InstallMemoryPoolsError {}
cubecl_runtime/memory_management/mod.rs

cubecl_runtime/memory_management/
mod.rs
