cubecl_server/memory_management/relocation/
trigger.rs

1//! When a relocation runs, and why.
2
3use super::TargetRoom;
4
5/// What a [`RelocationTrigger`] reads of the pools a workload's allocations
6/// are carved from.
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8pub struct ArenaState {
9    /// The size the current pages are carved at, which the next page has.
10    pub page_size: u64,
11    /// Whether any pool is outdated: something a relocation could empty.
12    pub has_outdated: bool,
13    /// Whether every slot holds a pool, so the next growth has nowhere to go.
14    pub full: bool,
15    /// What a relocation plan reads, to tell whether anything changed since
16    /// the last plan found nothing to move.
17    pub shape: ArenaShape,
18}
19
20/// What a relocation plan reads of the pools: how many pages the current one
21/// has room on, and what the outdated ones hold and keep guarded.
22///
23/// Room freed inside the current pages does not show here, since slices free
24/// as their owners drop them: that room is found by the plan after the next
25/// page, growth or explicit cleanup.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub struct ArenaShape {
28    /// Pages the current pool holds.
29    pub current_pages: u64,
30    /// Pools a growth left behind.
31    pub outdated_pools: usize,
32    /// Pages those pools hold.
33    pub outdated_pages: u64,
34    /// Of those, the ones a guard keeps where they are.
35    pub outdated_guarded: usize,
36}
37
38/// Decides when a relocation is worth its copies: what the device holds, and
39/// how the last relocation went.
40#[derive(Debug, Clone)]
41pub struct RelocationTrigger {
42    /// What the device holds in total, where it says.
43    capacity: Option<u64>,
44    /// The pools as the last relocation found them with nothing to move, so
45    /// the next is not planned until something changed.
46    stalled: Option<ArenaShape>,
47}
48
49/// Whether a relocation is wanted, before the bytes the device holds are
50/// known: plain data, so the caller can gather those bytes from every stream
51/// without anything borrowed.
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub enum RelocationNeed {
54    /// Nothing is outdated, the last relocation found nothing to move and
55    /// nothing changed since, or the device reports no capacity to judge by.
56    Nothing,
57    /// Every page size the pools track is taken.
58    ArenaFull,
59    /// Wanted once the next page of `page_size` would leave a device of
60    /// `capacity` bytes less than another page of room.
61    UnderPressure {
62        /// What the device holds in total.
63        capacity: u64,
64        /// The size of the next page.
65        page_size: u64,
66    },
67}
68
69/// Why a relocation runs, which decides where it may put what it moves.
70#[derive(Debug, Clone, Copy, PartialEq, Eq)]
71pub enum RelocationReason {
72    /// An explicit cleanup, which exists to give memory back.
73    Explicit,
74    /// The next page would leave the device too little room.
75    MemoryPressure,
76    /// Every page size the memory tracks is taken, so the next growth has
77    /// nowhere to go, however much room the device has.
78    ArenaFull,
79    /// A graph capture is about to start: every page it touches is guarded
80    /// for the graph's life, and a guarded page can no longer be emptied.
81    Capture,
82}
83
84/// How close the device is to its capacity, for the next page of a given size.
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86enum MemoryPressure {
87    /// The next page leaves room for at least another one.
88    Low,
89    /// The next page would leave less than another page of room: the last
90    /// moment a relocation can still free what the device needs, since once
91    /// the device refuses a page there is nowhere left to copy to.
92    High,
93}
94
95impl RelocationTrigger {
96    /// A trigger for a device that holds `capacity` bytes, where it says.
97    pub fn new(capacity: Option<u64>) -> Self {
98        Self {
99            capacity,
100            stalled: None,
101        }
102    }
103
104    /// Whether the pools in `arena` want a relocation.
105    pub fn need(&self, arena: &ArenaState) -> RelocationNeed {
106        if !arena.has_outdated || self.stalled == Some(arena.shape) {
107            return RelocationNeed::Nothing;
108        }
109        if arena.full {
110            return RelocationNeed::ArenaFull;
111        }
112        match self.capacity {
113            Some(capacity) => RelocationNeed::UnderPressure {
114                capacity,
115                page_size: arena.page_size,
116            },
117            None => RelocationNeed::Nothing,
118        }
119    }
120
121    /// A relocation moved something: the next is planned whenever it is
122    /// wanted.
123    pub fn moved(&mut self) {
124        self.stalled = None;
125    }
126
127    /// A relocation moved nothing: none is planned again until the pools in
128    /// `arena` change.
129    pub fn stalled(&mut self, arena: &ArenaState) {
130        self.stalled = Some(arena.shape);
131    }
132}
133
134impl RelocationNeed {
135    /// Why to relocate now, if at all, on a device whose memories hold
136    /// `allocated` bytes. Only asked when the answer depends on it.
137    pub fn reason(self, allocated: impl FnOnce() -> u64) -> Option<RelocationReason> {
138        match self {
139            RelocationNeed::Nothing => None,
140            RelocationNeed::ArenaFull => Some(RelocationReason::ArenaFull),
141            RelocationNeed::UnderPressure {
142                capacity,
143                page_size,
144            } => match MemoryPressure::new(allocated(), capacity, page_size) {
145                MemoryPressure::High => Some(RelocationReason::MemoryPressure),
146                MemoryPressure::Low => None,
147            },
148        }
149    }
150}
151
152impl RelocationReason {
153    /// Where the relocation may reserve its targets: only in the room held
154    /// when memory is what is short, anywhere otherwise.
155    pub fn room(self) -> TargetRoom {
156        match self {
157            RelocationReason::Explicit | RelocationReason::MemoryPressure => TargetRoom::Held,
158            RelocationReason::ArenaFull | RelocationReason::Capture => TargetRoom::MayAllocate,
159        }
160    }
161}
162
163impl MemoryPressure {
164    /// The pressure on a device that holds `allocated` of its `capacity`
165    /// bytes, for a next page of `page_size`.
166    fn new(allocated: u64, capacity: u64, page_size: u64) -> Self {
167        match allocated + page_size > capacity.saturating_sub(page_size) {
168            true => MemoryPressure::High,
169            false => MemoryPressure::Low,
170        }
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    const MIB: u64 = 1024 * 1024;
179
180    fn arena(page_size: u64) -> ArenaState {
181        ArenaState {
182            page_size,
183            has_outdated: true,
184            full: false,
185            shape: ArenaShape {
186                current_pages: 1,
187                outdated_pools: 1,
188                outdated_pages: 1,
189                outdated_guarded: 0,
190            },
191        }
192    }
193
194    #[test]
195    fn pressure_is_high_once_the_next_page_leaves_less_than_another() {
196        let need = RelocationTrigger::new(Some(24 * MIB)).need(&arena(4 * MIB));
197        assert_eq!(need.reason(|| 10 * MIB), None);
198        assert_eq!(
199            need.reason(|| 17 * MIB),
200            Some(RelocationReason::MemoryPressure)
201        );
202    }
203
204    #[test]
205    fn nothing_outdated_needs_nothing_and_sums_nothing() {
206        let state = ArenaState {
207            has_outdated: false,
208            ..arena(4 * MIB)
209        };
210        let need = RelocationTrigger::new(Some(MIB)).need(&state);
211        assert_eq!(need, RelocationNeed::Nothing);
212        assert_eq!(need.reason(|| panic!("the bytes are never summed")), None);
213    }
214
215    #[test]
216    fn a_full_arena_relocates_whatever_the_device_holds() {
217        let state = ArenaState {
218            full: true,
219            ..arena(4 * MIB)
220        };
221        assert_eq!(
222            RelocationTrigger::new(None).need(&state).reason(|| 0),
223            Some(RelocationReason::ArenaFull)
224        );
225    }
226
227    #[test]
228    fn a_device_that_reports_no_capacity_is_not_judged() {
229        assert_eq!(
230            RelocationTrigger::new(None).need(&arena(MIB)),
231            RelocationNeed::Nothing
232        );
233    }
234
235    #[test]
236    fn a_relocation_that_moved_nothing_waits_for_the_pools_to_change() {
237        let mut trigger = RelocationTrigger::new(Some(24 * MIB));
238        let state = arena(4 * MIB);
239        trigger.stalled(&state);
240        assert_eq!(trigger.need(&state), RelocationNeed::Nothing);
241
242        let grown = ArenaState {
243            shape: ArenaShape {
244                current_pages: 2,
245                ..state.shape
246            },
247            ..state
248        };
249        assert_ne!(trigger.need(&grown), RelocationNeed::Nothing);
250    }
251}
cubecl_server/memory_management/relocation/trigger.rs

cubecl_server/memory_management/relocation/
trigger.rs
