cubecl_runtime/memory_management/
mod.rs

1pub(crate) mod memory_pool;
2
3mod base;
4
5pub use base::*;
6
7/// Dynamic memory management strategy.
8mod memory_manage;
9pub use memory_manage::*;
10
11use alloc::vec::Vec;
12
13/// The type of memory pool to use.
14#[derive(Debug, Clone)]
15pub enum PoolType {
16    /// Use a memory where every allocation is a separate page.
17    ExclusivePages {
18        /// The minimum number of bytes to allocate in this pool.
19        max_alloc_size: u64,
20    },
21    /// Use a memory where each allocation is a slice of a bigger allocation.
22    SlicedPages {
23        /// The page size to allocate.
24        page_size: u64,
25        /// The maximum size of a slice to allocate in the pool.
26        max_slice_size: u64,
27    },
28}
29
30/// Options to create a memory pool.
31#[derive(Debug, Clone)]
32pub struct MemoryPoolOptions {
33    /// What kind of pool to use.
34    pub pool_type: PoolType,
35    /// Period after which allocations are deemed unused and deallocated.
36    ///
37    /// This period is measured in the number of allocations in the parent allocator. If a page
38    /// in the pool was unused for the entire period, it will be deallocated. This period is
39    /// approximmate, as checks are only done occasionally.
40    pub dealloc_period: Option<u64>,
41}
42
43/// High level configuration of memory management.
44#[derive(Clone, Debug)]
45pub enum MemoryConfiguration {
46    /// The default preset, which uses pools that allocate sub slices.
47    #[cfg(not(exclusive_memory_only))]
48    SubSlices,
49    /// Default preset for using exclusive pages.
50    /// This can be necessary for backends don't support sub-slices.
51    ExclusivePages,
52    /// Custom settings.
53    Custom {
54        /// Options for each pool to construct. When allocating, the first
55        /// possible pool will be picked for an allocation.
56        pool_options: Vec<MemoryPoolOptions>,
57    },
58}
59
60#[allow(clippy::derivable_impls)]
61impl Default for MemoryConfiguration {
62    fn default() -> Self {
63        #[cfg(exclusive_memory_only)]
64        {
65            MemoryConfiguration::ExclusivePages
66        }
67        #[cfg(not(exclusive_memory_only))]
68        {
69            MemoryConfiguration::SubSlices
70        }
71    }
72}
cubecl_runtime/memory_management/mod.rs

cubecl_runtime/memory_management/
mod.rs
