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