cubecl_server/memory_management/relocation/
base.rs

1//! Moving live allocations off outdated pages, so those pages go back to the
2//! driver now rather than when their longest-lived allocation ends.
3//!
4//! A [`Relocation`] is planned by the pools it empties, copied on the device
5//! through a [`CopyQueue`], and committed by those pools again from the
6//! [`Landed`] relocation the copy step hands back: an allocation cannot move
7//! before its bytes are where it moves to.
8
9use super::CopyQueue;
10use crate::memory_management::{ManagedMemoryHandle, release_or_leak};
11use crate::server::ServerError;
12use crate::storage::ComputeStorage;
13use crate::storage::StorageHandle;
14use alloc::vec::Vec;
15
16/// Live allocations on outdated pages, each with a slice reserved on a current
17/// page to receive it.
18///
19/// Nothing has moved yet: dropping it abandons the relocation with nothing
20/// lost, the reserved slices freed with it.
21#[derive(Debug)]
22pub struct Relocation {
23    moves: Vec<Move>,
24}
25
26/// Where a relocation may reserve the slices it moves allocations into.
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum TargetRoom {
29    /// Only the room the current pages already have. Memory is what is short,
30    /// and a relocation that allocated a page would spend more than it frees.
31    Held,
32    /// New pages of the current size too, when the room held runs out. Memory
33    /// is not what is short, so emptying every outdated page is worth a page.
34    MayAllocate,
35}
36
37/// A [`Relocation`] whose copies have landed: what the pools commit, and only
38/// [`Relocation::copy`] makes.
39#[derive(Debug)]
40pub struct Landed {
41    moves: Vec<Move>,
42}
43
44/// One device copy a relocation needs, between two resolved storages of the
45/// same size.
46#[derive(Debug, Clone)]
47pub struct StorageCopy {
48    /// Where the bytes are now.
49    pub source: StorageHandle,
50    /// Where they go.
51    pub target: StorageHandle,
52}
53
54/// One live allocation moving to a reserved slice: `allocation`'s handle is
55/// handed over to the target once `copy` has landed.
56#[derive(Debug)]
57pub(crate) struct Move {
58    /// The allocation being moved, as its owners hold it.
59    pub allocation: ManagedMemoryHandle,
60    /// The slice reserved to receive it.
61    pub target: ManagedMemoryHandle,
62    /// The bytes to copy, or `None` when the source was carved under a process mode
63    /// that discards launches and never resolved: there is nothing behind it to
64    /// copy.
65    pub copy: Option<StorageCopy>,
66}
67
68impl Relocation {
69    /// A relocation of `moves`, each with its target already reserved.
70    pub(crate) fn new(moves: Vec<Move>) -> Self {
71        Self { moves }
72    }
73
74    /// How many allocations move.
75    pub fn len(&self) -> usize {
76        self.moves.len()
77    }
78
79    /// Whether there is nothing to move.
80    pub fn is_empty(&self) -> bool {
81        self.moves.is_empty()
82    }
83
84    /// Copy every allocation's bytes on `queue`, waiting for the device
85    /// before the first copy and for the copies after the last.
86    ///
87    /// The wait for the copies happens whether every one of them was enqueued
88    /// or not, so a relocation abandoned here never frees a target a copy is
89    /// still writing.
90    ///
91    /// # Errors
92    ///
93    /// The first copy the device refused, or the fault a wait revealed. The
94    /// relocation is abandoned either way.
95    pub fn copy<Storage: ComputeStorage>(
96        self,
97        storage: &mut Storage,
98        queue: &mut dyn CopyQueue<Storage>,
99    ) -> Result<Landed, ServerError> {
100        queue.wait_device()?;
101        let enqueued = self
102            .moves
103            .iter()
104            .filter_map(|relocated| relocated.copy.as_ref())
105            .try_for_each(|copy| queue.copy(storage, copy));
106        if let Err(err) = queue.wait_copies() {
107            release_or_leak(self, &err);
108            enqueued?;
109            return Err(err);
110        }
111        enqueued?;
112        Ok(Landed { moves: self.moves })
113    }
114}
115
116impl Landed {
117    /// The moves, for the pools that planned them to commit.
118    pub(crate) fn into_moves(self) -> Vec<Move> {
119        self.moves
120    }
121}
122
123#[cfg(all(test, feature = "storage-bytes"))]
124mod tests {
125    use super::*;
126    use crate::{driver::DriverError, server::IoError, storage::BytesStorage};
127
128    struct FailedWait;
129
130    impl CopyQueue<BytesStorage> for FailedWait {
131        fn wait_device(&mut self) -> Result<(), ServerError> {
132            Ok(())
133        }
134
135        fn copy(&mut self, _: &mut BytesStorage, _: &StorageCopy) -> Result<(), IoError> {
136            Ok(())
137        }
138
139        fn wait_copies(&mut self) -> Result<(), ServerError> {
140            Err(DriverError::new("wait", 201).into())
141        }
142    }
143
144    #[test]
145    fn failed_wait_keeps_pending_relocation_reserved() {
146        let mut storage = BytesStorage::default();
147        let source = ManagedMemoryHandle::new();
148        let target = ManagedMemoryHandle::new();
149        let relocation = Relocation::new(alloc::vec![Move {
150            allocation: source.clone(),
151            target: target.clone(),
152            copy: Some(StorageCopy {
153                source: storage.alloc(4).unwrap(),
154                target: storage.alloc(4).unwrap(),
155            }),
156        }]);
157        assert!(relocation.copy(&mut storage, &mut FailedWait).is_err());
158        assert!(!source.is_free());
159        assert!(!target.is_free());
160    }
161}
cubecl_server/memory_management/relocation/base.rs

cubecl_server/memory_management/relocation/
base.rs
