cubecl_server/memory_management/relocation/
copy_queue.rs

1//! The device side of a relocation.
2
3use super::StorageCopy;
4use crate::server::{IoError, ServerError};
5#[cfg(feature = "storage-bytes")]
6use crate::storage::BytesStorage;
7use crate::storage::ComputeStorage;
8
9/// Where a relocation's bytes are copied, and what it waits on.
10///
11/// The one part of a relocation a runtime supplies: everything else — when to
12/// move, what moves, and when the move is committed — is the same whichever
13/// device is underneath. It holds the device queue the copies go on, never the
14/// memory the relocation is moving, so a reservation can hand it the storage
15/// it is already holding.
16pub trait CopyQueue<Storage: ComputeStorage> {
17    /// Wait until the device is done with everything that could still touch
18    /// the allocations about to move: every stream, and whatever the driver
19    /// runs outside them.
20    ///
21    /// # Errors
22    ///
23    /// The fault a wait revealed.
24    fn wait_device(&mut self) -> Result<(), ServerError>;
25
26    /// Add copying `copy.source`'s bytes into `copy.target` to the batch
27    /// [`wait_copies`](Self::wait_copies) lands. A queue may hold it until
28    /// then, so one relocation is one submission however many copies it
29    /// makes.
30    ///
31    /// # Errors
32    ///
33    /// The device's refusal to copy.
34    fn copy(&mut self, storage: &mut Storage, copy: &StorageCopy) -> Result<(), IoError>;
35
36    /// Submit the copies added so far, if the queue held them, and wait until
37    /// they have landed.
38    ///
39    /// # Errors
40    ///
41    /// The fault the wait revealed.
42    fn wait_copies(&mut self) -> Result<(), ServerError>;
43}
44
45/// Copies between the allocations of a [`BytesStorage`]: host memory, so a
46/// copy is a `memcpy` and has landed by the time it returns.
47///
48/// Waits on nothing: whoever runs work against the memory waits for it before
49/// handing the relocation over.
50#[cfg(feature = "storage-bytes")]
51#[derive(Debug, Default)]
52pub struct HostCopies;
53
54#[cfg(feature = "storage-bytes")]
55impl CopyQueue<BytesStorage> for HostCopies {
56    fn wait_device(&mut self) -> Result<(), ServerError> {
57        Ok(())
58    }
59
60    fn copy(&mut self, storage: &mut BytesStorage, copy: &StorageCopy) -> Result<(), IoError> {
61        let source = storage.get(&copy.source)?;
62        let mut target = storage.get(&copy.target)?;
63        target.write().copy_from_slice(source.read());
64        Ok(())
65    }
66
67    fn wait_copies(&mut self) -> Result<(), ServerError> {
68        Ok(())
69    }
70}
cubecl_server/memory_management/relocation/copy_queue.rs

cubecl_server/memory_management/relocation/
copy_queue.rs
