cubecl_server/memory_management/relocation/
streams.rs

1//! A relocation as a server runs it: across every stream, since any of them
2//! may use the memory that moves.
3
4use super::{RelocationNeed, RelocationReason};
5use crate::memory_management::ErrorGraph;
6use crate::server::ServerError;
7
8/// One stream, as relocating its memory needs it: what a runtime whose
9/// streams share one way to finish their work (a scheduler, say) supplies per
10/// stream, leaving the rest of [`RelocatingStreams`] to that shared way.
11pub trait RelocatableStream {
12    /// Whether any stream of the device records a graph.
13    fn recording(&self) -> bool;
14
15    /// Whether a growth left pages behind in the stream's memory.
16    fn has_outdated(&self) -> bool;
17
18    /// Whether the stream's memory wants a relocation, see
19    /// [`MemoryManagement::relocation_need`](crate::memory_management::MemoryManagement::relocation_need).
20    fn relocation_need(&self) -> RelocationNeed;
21
22    /// The bytes the stream's memory holds from the device.
23    fn bytes_allocated(&self) -> u64;
24
25    /// Move what the stream's memory holds on outdated pages, once every
26    /// stream's work is done.
27    fn relocate(&mut self, reason: RelocationReason, failures: &mut ErrorGraph);
28
29    /// Give back every page the stream's memory holds and nothing needs.
30    fn cleanup_memory(&mut self, failures: &mut ErrorGraph);
31}
32
33/// A runtime's streams, as relocating one stream's memory needs them.
34///
35/// A relocation moves memory any stream may still read or write, so every
36/// stream's work is done before a byte moves; and nothing moves while a graph
37/// records, since the pages it touches keep their addresses until it seals.
38/// The runtime supplies each step. The order they run in lives here, once.
39pub trait RelocatingStreams {
40    /// Whether any stream is recording a graph.
41    fn recording(&mut self) -> bool;
42
43    /// Whether a growth left pages behind in the relocating stream's memory:
44    /// without any, a relocation has nothing to move and waits on nothing.
45    fn has_outdated(&mut self) -> bool;
46
47    /// Whether the relocating stream's memory wants a relocation, before the
48    /// bytes the device holds are known.
49    fn relocation_need(&mut self) -> RelocationNeed;
50
51    /// The bytes every stream's memory holds from the device.
52    fn bytes_allocated(&mut self) -> u64;
53
54    /// Get every stream's queued and submitted work done, or leave it to the
55    /// [`CopyQueue`](super::CopyQueue) to wait for before its first copy.
56    fn finish(&mut self);
57
58    /// Move what the relocating stream's memory holds on outdated pages.
59    fn relocate_memory(&mut self, reason: RelocationReason);
60
61    /// Whether a growth left pages behind in any stream's memory.
62    fn device_has_outdated(&mut self) -> bool;
63
64    /// Move what every stream's memory holds on outdated pages.
65    fn relocate_device_memory(&mut self, reason: RelocationReason);
66
67    /// Give back every page the relocating stream's memory holds and nothing
68    /// needs.
69    fn cleanup_memory(&mut self);
70
71    /// Refuse while any stream records a graph: the pages a recording
72    /// touches keep their addresses until it seals, and nothing may wait on
73    /// the device under it.
74    ///
75    /// What [`reclaim`](Self::reclaim) checks first, for a caller with work of
76    /// its own to refuse before it.
77    ///
78    /// # Errors
79    ///
80    /// [`ServerError`] naming the recording when a stream records.
81    fn refuse_while_recording(&mut self) -> Result<(), ServerError> {
82        if self.recording() {
83            return Err(ServerError::graph_state(
84                "memory_cleanup: a stream is recording a graph, whose pages keep their addresses \
85                 until it seals",
86            ));
87        }
88        Ok(())
89    }
90
91    /// Give back everything the relocating stream's memory holds and nothing
92    /// needs: what the outdated pages still hold moves first, while the
93    /// current pages keep the room it moves into, then every page left empty
94    /// goes back.
95    ///
96    /// # Errors
97    ///
98    /// Refused while any stream records a graph, see
99    /// [`refuse_while_recording`](Self::refuse_while_recording).
100    fn reclaim(&mut self) -> Result<(), ServerError> {
101        self.refuse_while_recording()?;
102        self.relocate(RelocationReason::Explicit);
103        self.cleanup_memory();
104        Ok(())
105    }
106
107    /// Relocate when the memory asks for it: while the device still has a
108    /// page to spare, or before the pools run out of page sizes.
109    fn relocate_when_wanted(&mut self) {
110        let need = self.relocation_need();
111        if need == RelocationNeed::Nothing || self.recording() {
112            return;
113        }
114        if let Some(reason) = need.reason(|| self.bytes_allocated()) {
115            self.finish();
116            self.relocate_memory(reason);
117        }
118    }
119
120    /// Relocate for `reason`, unless a graph records or nothing is outdated.
121    fn relocate(&mut self, reason: RelocationReason) {
122        if self.recording() || !self.has_outdated() {
123            return;
124        }
125        self.finish();
126        self.relocate_memory(reason);
127    }
128
129    /// Relocate every stream's memory for `reason`, unless a graph records or
130    /// nothing is outdated anywhere.
131    ///
132    /// What a capture does before its warmup: a recording may touch memory
133    /// any stream owns, and every page it touches is guarded for the graph's
134    /// life, where no relocation can empty it.
135    fn relocate_device(&mut self, reason: RelocationReason) {
136        if self.recording() || !self.device_has_outdated() {
137            return;
138        }
139        self.finish();
140        self.relocate_device_memory(reason);
141    }
142}
cubecl_server/memory_management/relocation/streams.rs

cubecl_server/memory_management/relocation/
streams.rs
