Skip to main content

cubecl_runtime/memory_management/
guard.rs

1use super::ManagedMemoryBinding;
2use alloc::sync::Arc;
3
4/// Keeps a page as it is for as long as it lives: nothing new is carved from
5/// it, it is never released, and relocation leaves everything on it in place.
6///
7/// What a graph holds on the memory it replays against. A graph keeps raw
8/// addresses but no bindings, so the bytes behind them must neither move nor
9/// be handed to another allocation. The allocations on the page keep their
10/// handles untouched, so one can still be updated in place while the page is
11/// guarded.
12///
13/// A resolved resource needs no guard: the binding it holds keeps its own
14/// allocation in place, and leaves the rest of the page to other allocations.
15#[derive(Debug, Clone)]
16pub struct PageGuard {
17    _hold: Hold,
18}
19
20#[derive(Debug, Clone)]
21enum Hold {
22    /// A page carved into slices, which counts the guards it hands out.
23    Page { _token: Arc<()> },
24    /// A slice that is its own device allocation. Holding its binding is what
25    /// keeps it from being handed out again or released, whether an
26    /// allocation is on it or not.
27    Allocation { _binding: ManagedMemoryBinding },
28}
29
30impl PageGuard {
31    /// A guard on a page that counts its guards with `token`.
32    #[doc(hidden)]
33    pub fn page(token: &Arc<()>) -> Self {
34        Self {
35            _hold: Hold::Page {
36                _token: token.clone(),
37            },
38        }
39    }
40
41    /// A guard on a page that is one allocation, held through `binding`.
42    #[doc(hidden)]
43    pub fn allocation(binding: ManagedMemoryBinding) -> Self {
44        Self {
45            _hold: Hold::Allocation { _binding: binding },
46        }
47    }
48}