cubecl_runtime/memory_management/
mod.rs

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

cubecl_runtime/memory_management/
mod.rs
