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}