cubecl_runtime/memory_management/base.rs
1use alloc::string::{String, ToString};
2use alloc::vec::Vec;
3
4/// Amount of memory in use by this allocator
5/// and statistics on how much memory is reserved and
6/// wasted in total.
7#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
8pub struct MemoryUsage {
9 /// The number of allocations currently active.
10 ///
11 /// This is not the number of times an actual allocation happens to create a new memory page,
12 /// but really the number of active slices.
13 pub number_allocs: u64,
14 /// The number of bytes that are currently actually in use.
15 ///
16 /// This doesn't include any padding or other memory that needs to be
17 /// reserved, and is the minimum amount of memory that could possible
18 /// be allocated.
19 pub bytes_in_use: u64,
20 /// The amount of bytes used for padding memory in currently active allocations.
21 pub bytes_padding: u64,
22 /// The total amount of memory reserved on the device.
23 ///
24 /// This will be at least as much as `bytes_in_use` but in practice will
25 /// be higher, as allocations reserve memory for future allocations
26 /// and for padding.
27 pub bytes_reserved: u64,
28}
29
30impl MemoryUsage {
31 /// Calculate the combined memory usage of two reports (summing them).
32 pub fn combine(&self, other: MemoryUsage) -> MemoryUsage {
33 MemoryUsage {
34 number_allocs: self.number_allocs + other.number_allocs,
35 bytes_in_use: self.bytes_in_use + other.bytes_in_use,
36 bytes_padding: self.bytes_padding + other.bytes_padding,
37 bytes_reserved: self.bytes_reserved + other.bytes_reserved,
38 }
39 }
40}
41
42#[derive(new)]
43#[doc(hidden)]
44pub struct BytesFormat {
45 bytes: u64,
46}
47
48impl core::fmt::Display for BytesFormat {
49 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
50 let unit = 1000;
51
52 if self.bytes < unit {
53 f.write_fmt(format_args!("{} B", self.bytes))
54 } else {
55 let size = self.bytes as f64;
56 let exp = match size.log(1000.0).floor() as usize {
57 0 => 1,
58 e => e,
59 };
60 let unit_prefix = "KMGTPEZY".as_bytes();
61 f.write_fmt(format_args!(
62 "{:.2} {}B",
63 (size / unit.pow(exp as u32) as f64),
64 unit_prefix[exp - 1] as char,
65 ))
66 }
67 }
68}
69
70fn bytes_format(bytes: u64) -> String {
71 BytesFormat::new(bytes).to_string()
72}
73
74impl core::fmt::Display for MemoryUsage {
75 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
76 // In the future it'd be nice if MemoryUsage also held some stats about say,
77 // the 5 biggest allocations, to show when you an OOM.
78 let usage_percentage = (self.bytes_in_use as f32 / self.bytes_reserved as f32) * 100.0;
79 let padding_percentage = (self.bytes_padding as f32 / self.bytes_in_use as f32) * 100.0;
80 writeln!(f, "Memory Usage Report:")?;
81 writeln!(f, " Number of allocations: {}", self.number_allocs)?;
82 writeln!(f, " Bytes in use: {}", bytes_format(self.bytes_in_use))?;
83 writeln!(
84 f,
85 " Bytes used for padding: {}",
86 bytes_format(self.bytes_padding)
87 )?;
88 writeln!(
89 f,
90 " Total bytes reserved: {}",
91 bytes_format(self.bytes_reserved)
92 )?;
93 writeln!(f, " Usage efficiency: {usage_percentage:.2}%")?;
94 writeln!(f, " Padding overhead: {padding_percentage:.2}%")
95 }
96}
97
98/// The pool shape a [`MemoryPoolReport`] describes, carrying the pool's
99/// effective configuration (after alignment rounding and page-size shrinking).
100#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
101pub enum MemoryPoolKind {
102 /// Allocations are slices carved from shared pages.
103 Sliced {
104 /// The size of each device page.
105 page_size: u64,
106 /// The largest allocation the pool accepts.
107 max_slice_size: u64,
108 /// The pool's byte cap (`None` grows unbounded).
109 max_pool_size: Option<u64>,
110 },
111 /// Every allocation is its own device page.
112 Exclusive {
113 /// The largest allocation the pool accepts.
114 max_alloc_size: u64,
115 },
116 /// One device allocation per reservation, sized to the request, reused by
117 /// exact size and returned to the driver only under memory pressure.
118 /// Wastes only alignment padding, and pays a driver allocation per
119 /// distinct size rather than per page.
120 Direct,
121 /// Exact-fit slices that are reused only by identical size.
122 Persistent,
123}
124
125/// A structured snapshot of one memory pool: its shape, its current usage, and
126/// the high-water marks a memory plan is derived from.
127#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
128pub struct MemoryPoolReport {
129 /// The pool's shape and effective configuration.
130 pub kind: MemoryPoolKind,
131 /// The pool's current usage.
132 pub usage: MemoryUsage,
133 /// Device allocations (pages) currently held.
134 pub pages: u64,
135 /// The most device allocations ever held at once.
136 ///
137 /// For a sliced pool this is the number a capped layout needs:
138 /// pages are carved by a deterministic first-fit policy, so replaying the
139 /// same allocation stream against `pages_peak * page_size` fits by
140 /// construction.
141 pub pages_peak: u64,
142 /// How many of the current pages have no device backing yet — carved
143 /// under a dry run and never resolved into anything that executes. They
144 /// count toward `pages`/`pages_peak` (the plan is the *reserved* stream)
145 /// while costing no device memory; `pages - pages_unmapped` is the dry
146 /// run's actual footprint in this pool.
147 pub pages_unmapped: u64,
148 /// The largest single allocation this pool ever served, in requested
149 /// (pre-padding) bytes.
150 pub largest_alloc: u64,
151}
152
153/// A per-pool report of one `MemoryManagement` (in `cubecl-server`)
154/// instance — the read side of a measured memory plan.
155///
156/// The intended cycle: install a growable layout, run the workload once under
157/// a [`DryRun`](crate::dry_run::DryRun) (same allocation stream, no compute),
158/// read this report, and re-install the same layout capped at the observed
159/// `pages_peak`. Padding then comes only from alignment and the first-fit
160/// remainders the dry run already measured.
161///
162/// A tuning pass inside the measured run allocates too, and its scratch counts
163/// toward these marks like anything else. Warming the tune caches in an
164/// earlier pass and rebuilding the pools
165/// (`MemoryManagement::install_pools`, which resets the
166/// marks)
167/// before the measured one leaves the peaks to the workload alone.
168#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
169pub struct MemoryReport {
170 /// One entry per dynamic pool, in allocation-routing order — the same
171 /// order the layout was configured with.
172 pub dynamic: Vec<MemoryPoolReport>,
173 /// The persistent pool (weights, caches; explicit persistent windows).
174 pub persistent: MemoryPoolReport,
175}
176
177/// A [`MemoryReport`] as the environment records it: a snapshot of one
178/// stream's pools at a moment the caller named, written by
179/// [`Client::record_memory`](crate::client::Client::record_memory).
180#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
181pub struct MemoryRecord {
182 /// What the caller was doing: `model loaded`, `after the dry run`.
183 pub label: alloc::string::String,
184 /// The pools at that moment.
185 pub report: MemoryReport,
186}
187
188impl cubecl_environment::records::Record for MemoryRecord {
189 const KIND: &'static str = "memory";
190}
191
192/// The managed tensor buffer handle that points to some memory segment.
193/// It should not contain actual data.
194pub trait MemoryHandle<Binding>: Clone + core::fmt::Debug {
195 /// Checks if the underlying memory can be safely mutated.
196 fn can_mut(&self) -> bool;
197 /// Get the binding associated to the current handle.
198 fn binding(self) -> Binding;
199}