burn_backend/backend/
memory_pools.rs

1//! The layout of a device's dynamic memory pools, and what one reports.
2//!
3//! An allocator that only grows keeps whatever page it ever needed, so a
4//! long-running workload reserves its worst moment for life. These types let a
5//! caller install a layout of its own — a pool per size class, each capped at a
6//! number of pages — and read back what the workload actually held, so the caps
7//! are a measurement rather than a guess.
8//!
9//! The vocabulary is the backend's own rather than any runtime's, since
10//! [`Backend`](super::Backend) is also implemented by backends with no pools at
11//! all.
12
13use alloc::string::String;
14use alloc::vec::Vec;
15
16/// A layout for a device's dynamic memory pools, applied with
17/// [`Backend::memory_install_pools`](super::Backend::memory_install_pools).
18#[derive(Clone, Debug, PartialEq, Eq)]
19pub enum MemoryPoolLayout {
20    /// An ordered list of pools that sub-slice fixed-size pages. An allocation
21    /// lands in the first pool whose [`max_slice`](SlicedPool::max_slice)
22    /// accepts its size, so a small pool listed before a large one captures the
23    /// small-allocation traffic. Pages are allocated on first use.
24    Sliced(Vec<SlicedPool>),
25    /// One direct pool for everything: every allocation is its own, sized to
26    /// the request and reused by exact size — no pages, no sub-slicing, no
27    /// padding beyond alignment. With no page size chosen in advance, this is
28    /// what a workload runs on for its largest allocation to be read back as it
29    /// was asked for.
30    Direct,
31    /// The runtime's default: a ladder of size-bucketed pools that sub-slice
32    /// large pages.
33    SubSlices,
34    /// One page per allocation, in exponentially spaced size buckets.
35    ExclusivePages,
36}
37
38/// One pool of a [`MemoryPoolLayout::Sliced`] layout: allocations are slices of
39/// fixed-size pages, capped at `pages` pages. An allocation that no longer fits
40/// goes to the next pool that accepts it, and fails when none does.
41///
42/// Sizes are rounded up to the device's alignment, so a pool holds at least
43/// what it was asked for. A layout that cannot be honoured at all — a zero
44/// size, `max_slice` past `page_size`, `pages` outside `1..=65535` — is refused
45/// with [`InvalidLayout`](InstallMemoryPoolsError::InvalidLayout).
46#[derive(Clone, Copy, Debug, PartialEq, Eq)]
47pub struct SlicedPool {
48    /// Size of each page in bytes; also the largest single allocation this pool
49    /// can serve. Rounded up to the device's alignment.
50    pub page_size: u64,
51    /// How many pages the pool may hold. `None` grows without a cap, which is
52    /// what a workload is measured on before its caps are known.
53    pub pages: Option<u64>,
54    /// Largest allocation routed to this pool; `None` accepts anything up to
55    /// `page_size`. Later pools see only what this one declines.
56    pub max_slice: Option<u64>,
57}
58
59/// One dynamic pool's measured state, in the order the pools were installed.
60///
61/// The read side of a measured layout: install a growable one, run the
62/// workload, read these, re-install capped at `pages_peak`. Pool placement is
63/// deterministic, so the same allocations fit the capped layout by
64/// construction.
65#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
66pub struct SlicedPoolReport {
67    /// Size of each page in bytes; `0` for a [direct](MemoryPoolLayout::Direct)
68    /// pool, which has no pages to size.
69    pub page_size: u64,
70    /// Pages currently held — for a direct pool, live allocations.
71    pub pages: u64,
72    /// The most pages ever held at once.
73    pub pages_peak: u64,
74    /// The largest single allocation served, in requested bytes.
75    pub largest_alloc: u64,
76}
77
78/// One reading of a device allocator's state, across the runtime's streams and
79/// pools.
80#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
81pub struct MemoryPoolUsage {
82    /// Live allocations, not pages.
83    pub number_allocs: u64,
84    /// Bytes those allocations use, excluding padding.
85    pub bytes_in_use: u64,
86    /// Bytes of padding inside them.
87    pub bytes_padding: u64,
88    /// Total bytes reserved on the device: at least `bytes_in_use`, plus pages
89    /// held for reuse.
90    pub bytes_reserved: u64,
91}
92
93/// Why installing a pool layout did not take effect.
94///
95/// The distinction a caller needs is **transient or permanent**: a refused
96/// rebuild is worth retrying once whatever holds the pools drains, while a
97/// backend with no configurable pools refuses forever. Treating the two alike
98/// either gives up on a layout that would have installed, or repeats an
99/// expensive measurement that can never succeed.
100#[derive(Clone, Debug, PartialEq, Eq)]
101pub enum InstallMemoryPoolsError {
102    /// The pools being rebuilt still hold live allocations, so the previous
103    /// layout was kept. Transient.
104    PoolsInUse {
105        /// Bytes still live in those pools.
106        bytes_in_use: u64,
107    },
108    /// The calling stream is already in an error state, so its pools were not
109    /// rebuilt. The layout still applies to streams created afterwards, and the
110    /// underlying failure surfaces at the next flush or sync.
111    StreamUnavailable,
112    /// This backend has no configurable dynamic pools. Permanent.
113    Unsupported,
114    /// The layout itself cannot be honoured, so the previous one was kept: an
115    /// empty pool list, a zero size, a slice larger than its page, an
116    /// unusable cap, or a pool shape this build has none of. Permanent — what
117    /// has to change is the layout, not the moment it is installed at.
118    InvalidLayout {
119        /// What the backend objected to, in its own words.
120        reason: String,
121    },
122}
123
124impl core::fmt::Display for InstallMemoryPoolsError {
125    fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
126        match self {
127            Self::PoolsInUse { bytes_in_use } => {
128                write!(formatter, "{bytes_in_use} B still live in the pools")
129            }
130            Self::StreamUnavailable => {
131                write!(formatter, "the calling stream is in an error state")
132            }
133            Self::Unsupported => {
134                write!(formatter, "this backend has no configurable memory pools")
135            }
136            Self::InvalidLayout { reason } => {
137                write!(formatter, "the pool layout cannot be honoured: {reason}")
138            }
139        }
140    }
141}
142
143impl core::error::Error for InstallMemoryPoolsError {}
burn_backend/backend/memory_pools.rs

burn_backend/backend/
memory_pools.rs
