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 dry run
63    /// and never resolved: there is nothing behind it to copy.
64    pub copy: Option<StorageCopy>,
65}
66
67impl Relocation {
68    /// A relocation of `moves`, each with its target already reserved.
69    pub(crate) fn new(moves: Vec<Move>) -> Self {
70        Self { moves }
71    }
72
73    /// How many allocations move.
74    pub fn len(&self) -> usize {
75        self.moves.len()
76    }
77
78    /// Whether there is nothing to move.
79    pub fn is_empty(&self) -> bool {
80        self.moves.is_empty()
81    }
82
83    /// Copy every allocation's bytes on `queue`, waiting for the device
84    /// before the first copy and for the copies after the last.
85    ///
86    /// The wait for the copies happens whether every one of them was enqueued
87    /// or not, so a relocation abandoned here never frees a target a copy is
88    /// still writing.
89    ///
90    /// # Errors
91    ///
92    /// The first copy the device refused, or the fault a wait revealed. The
93    /// relocation is abandoned either way.
94    pub fn copy<Storage: ComputeStorage>(
95        self,
96        storage: &mut Storage,
97        queue: &mut dyn CopyQueue<Storage>,
98    ) -> Result<Landed, ServerError> {
99        queue.wait_device()?;
100        let enqueued = self
101            .moves
102            .iter()
103            .filter_map(|relocated| relocated.copy.as_ref())
104            .try_for_each(|copy| queue.copy(storage, copy));
105        if let Err(err) = queue.wait_copies() {
106            release_or_leak(self, &err);
107            enqueued?;
108            return Err(err);
109        }
110        enqueued?;
111        Ok(Landed { moves: self.moves })
112    }
113}
114
115impl Landed {
116    /// The moves, for the pools that planned them to commit.
117    pub(crate) fn into_moves(self) -> Vec<Move> {
118        self.moves
119    }
120}
121
122#[cfg(all(test, feature = "storage-bytes"))]
123mod tests {
124    use super::*;
125    use crate::{driver::DriverError, server::IoError, storage::BytesStorage};
126
127    struct FailedWait;
128
129    impl CopyQueue<BytesStorage> for FailedWait {
130        fn wait_device(&mut self) -> Result<(), ServerError> {
131            Ok(())
132        }
133
134        fn copy(&mut self, _: &mut BytesStorage, _: &StorageCopy) -> Result<(), IoError> {
135            Ok(())
136        }
137
138        fn wait_copies(&mut self) -> Result<(), ServerError> {
139            Err(DriverError::new("wait", 201).into())
140        }
141    }
142
143    #[test]
144    fn failed_wait_keeps_pending_relocation_reserved() {
145        let mut storage = BytesStorage::default();
146        let source = ManagedMemoryHandle::new();
147        let target = ManagedMemoryHandle::new();
148        let relocation = Relocation::new(alloc::vec![Move {
149            allocation: source.clone(),
150            target: target.clone(),
151            copy: Some(StorageCopy {
152                source: storage.alloc(4).unwrap(),
153                target: storage.alloc(4).unwrap(),
154            }),
155        }]);
156        assert!(relocation.copy(&mut storage, &mut FailedWait).is_err());
157        assert!(!source.is_free());
158        assert!(!target.is_free());
159    }
160}
cubecl_server/memory_management/relocation/base.rs

cubecl_server/memory_management/relocation/
base.rs
