cubecl_runtime/config/memory.rs
1use super::logger::{LogLevel, LoggerConfig};
2use super::size::MemorySize;
3use alloc::vec::Vec;
4
5/// Configuration for memory settings in `CubeCL`.
6///
7/// Unknown fields are rejected so a leftover `pools` entry (now a programmatic
8/// setting, see [`MemoryPoolsConfig`]) or a misspelled option is a load error
9/// rather than a silently dropped setting.
10#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Default)]
11#[serde(deny_unknown_fields)]
12pub struct MemoryConfig {
13 /// Logger configuration for memory-related logs, using specific log levels.
14 #[serde(default)]
15 pub logger: LoggerConfig<MemoryLogLevel>,
16 /// Configuration for persistent memory pools.
17 #[serde(default)]
18 pub persistent_memory: PersistentMemory,
19}
20
21/// A pool layout override for a runtime's **main GPU** memory: a preset or an
22/// explicit list of pool entries, tried in order at allocation time (the first
23/// pool that accepts an allocation's size serves it).
24///
25/// This is a **programmatic** setting, deliberately not a config-file one —
26/// pool layouts are dynamic (e.g. resized per model just before a load) and
27/// must not freeze at startup. Apply it with
28/// [`configure_memory_pools`](crate::client::ComputeClient::configure_memory_pools):
29/// it rebuilds the calling stream's pools in place and becomes the layout for
30/// streams created afterwards. Auxiliary pools (pinned CPU, staging, uniforms)
31/// are never affected.
32#[derive(Clone, Debug, PartialEq)]
33pub enum MemoryPoolsConfig {
34 /// A named preset matching the runtime-level presets.
35 Preset(MemoryPoolsPreset),
36 /// An explicit pool list, mirroring
37 /// [`MemoryPoolOptions`](crate::memory_management::MemoryPoolOptions).
38 Explicit(Vec<MemoryPoolConfig>),
39}
40
41/// The presets of [`MemoryConfiguration`](crate::memory_management::MemoryConfiguration).
42#[derive(Clone, Copy, Debug, PartialEq, Eq)]
43pub enum MemoryPoolsPreset {
44 /// The runtime's `SubSlices` preset: a ladder of size-bucketed pools that
45 /// sub-slice large pages.
46 SubSlices,
47 /// The runtime's `ExclusivePages` preset: one page per allocation, in
48 /// exponentially spaced size buckets.
49 ExclusivePages,
50}
51
52/// One pool entry; mirrors [`MemoryPoolOptions`](crate::memory_management::MemoryPoolOptions)
53/// and [`PoolType`](crate::memory_management::PoolType).
54///
55/// Sizes are aligned up to the device alignment.
56#[derive(Clone, Debug, PartialEq)]
57pub enum MemoryPoolConfig {
58 /// Every allocation gets its own page
59 /// ([`PoolType::ExclusivePages`](crate::memory_management::PoolType::ExclusivePages)).
60 Exclusive {
61 /// Largest allocation this pool accepts. `0` is valid: it makes a pool
62 /// dedicated to zero-sized (sub-alignment) allocations.
63 max_alloc_size: MemorySize,
64 /// Period (in parent allocation count) after which unused pages are
65 /// deallocated. `None` never deallocates.
66 dealloc_period: Option<u64>,
67 },
68 /// Allocations are slices of larger pages
69 /// ([`PoolType::SlicedPages`](crate::memory_management::PoolType::SlicedPages)).
70 Sliced {
71 /// Size of each page.
72 page_size: MemorySize,
73 /// Largest slice this pool accepts. Defaults to `page_size`.
74 max_slice_size: Option<MemorySize>,
75 /// Hard cap on the pool's total reserved bytes: exceeding it is an
76 /// error instead of silent growth. `None` grows unbounded. Note:
77 /// runtimes that create one memory management per stream (CUDA, HIP)
78 /// apply the cap per stream.
79 max_pool_size: Option<MemorySize>,
80 /// Period (in parent allocation count) after which unused pages are
81 /// deallocated. `None` never deallocates.
82 dealloc_period: Option<u64>,
83 },
84}
85
86/// Configuration options for persistent memory pools in `CubeCL` runtimes.
87#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Default)]
88pub enum PersistentMemory {
89 /// Persistent memory is enabled but used only when explicitly specified.
90 #[default]
91 #[serde(rename = "enabled")]
92 Enabled,
93 /// Like `enabled`, and automatic allocations whose size matches an
94 /// existing persistent bucket are also served from the persistent pool.
95 ///
96 /// A good heuristic for training, where the recurring allocations are
97 /// weight-shaped and updated within the same pool. Less suited to
98 /// inference: activations that happen to match a weight size get pulled
99 /// into exact-sized persistent slices.
100 #[serde(rename = "size-match")]
101 SizeMatch,
102 /// Persistent memory is disabled, allowing dynamic allocations.
103 #[serde(rename = "disabled")]
104 Disabled,
105 /// Persistent memory is enforced, preventing dynamic allocations.
106 ///
107 /// # Warning
108 ///
109 /// Enforcing persistent memory may cause out-of-memory errors if tensors of varying sizes are used.
110 #[serde(rename = "enforced")]
111 Enforced,
112}
113
114/// Log levels for memory-related events in `CubeCL`.
115#[derive(Default, Clone, Copy, Debug, serde::Serialize, serde::Deserialize)]
116pub enum MemoryLogLevel {
117 /// No memory-related logging.
118 #[default]
119 #[serde(rename = "disabled")]
120 Disabled,
121 /// Logs basic memory events, such as creating memory pages and manually cleaning memory.
122 #[serde(rename = "basic")]
123 Basic,
124 /// Logs detailed memory information.
125 #[serde(rename = "full")]
126 Full,
127}
128
129impl LogLevel for MemoryLogLevel {}
130
131#[cfg(all(test, feature = "std"))]
132mod tests {
133 use super::*;
134
135 #[test]
136 fn pools_rejected_in_config_files() {
137 // Pool layouts are a programmatic setting; a leftover `pools` entry in
138 // a config file must be a load error, not a silently ignored setting.
139 assert!(toml::from_str::<MemoryConfig>("pools = \"sub-slices\"").is_err());
140 assert!(
141 toml::from_str::<MemoryConfig>("[[pools]]\ntype = \"sliced\"\npage_size = \"1MiB\"\n")
142 .is_err()
143 );
144 }
145}