burn_compute/memory_management/
base.rs

1use crate::storage::ComputeStorage;
2
3/// The MemoryHandle trait is an abstract way to refer to some memory segment.
4/// It should not contain actual references to data.
5///
6/// It is responsible for determining if the memory segment can be mutated,
7/// for instance by keeping track of a reference count
8pub trait MemoryHandle: Clone + Send + Sync + core::fmt::Debug {
9    /// Checks if the underlying memory can be safely mutated.
10    fn can_mut(&self) -> bool;
11}
12
13/// The MemoryManagement trait encapsulates strategies for (de)allocating memory.
14/// It is bound to the ComputeStorage trait, which does the actual (de)allocations.
15///
16/// The MemoryManagement can only reserve memory space or get the resource located at a space.
17/// Modification of the resource data should be done directly on the resource.
18pub trait MemoryManagement<Storage: ComputeStorage>: Send + core::fmt::Debug {
19    /// The associated type Handle must implement MemoryHandle
20    type Handle: MemoryHandle;
21
22    /// Returns the resource from the storage at the specified handle
23    fn get(&mut self, handle: &Self::Handle) -> Storage::Resource;
24
25    /// Finds a spot in memory for a resource with the given size in bytes, and returns a handle to it
26    fn reserve(&mut self, size: usize) -> Self::Handle;
27
28    /// Bypass the memory allocation algorithm to allocate data directly.
29    ///
30    /// # Notes
31    ///
32    /// Can be useful for servers that want specific control over memory.
33    fn alloc(&mut self, size: usize) -> Self::Handle;
34
35    /// Bypass the memory allocation algorithm to deallocate data directly.
36    ///
37    /// # Notes
38    ///
39    /// Can be useful for servers that want specific control over memory.
40    fn dealloc(&mut self, handle: &Self::Handle);
41
42    /// Fetch the storage used by the memory manager.
43    ///
44    /// # Notes
45    ///
46    /// The storage should probably not be used for allocations since the handles won't be
47    /// compatible with the ones provided by the current trait. Prefer using the
48    /// [alloc](MemoryManagement::alloc) and [dealloc](MemoryManagement::dealloc) functions.
49    ///
50    /// This is useful if you need to time the deallocations based on async computation, or to
51    /// change the mode of storage for different reasons.
52    fn storage(&mut self) -> &mut Storage;
53}
burn_compute/memory_management/base.rs

burn_compute/memory_management/
base.rs
