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    /// Give every allocation its own device allocation, sized to the request
25    /// and reused by exact size.
26    ///
27    /// No carving at all, so it wastes only alignment padding. Worth it where
28    /// padding matters more than allocation count — under a
29    /// [`DryRun`](crate::dry_run::DryRun), where unresolved reservations never
30    /// reach the driver at all, or on a device the workload barely fits.
31    Direct {
32        /// Reserved bytes above which free slices are returned to the driver.
33        /// A watermark rather than a budget; `None` reclaims only on an
34        /// explicit cleanup.
35        reclaim_at: Option<u64>,
36    },
37    /// Use a memory where each allocation is a slice of a bigger allocation.
38    SlicedPages {
39        /// The page size to allocate.
40        page_size: u64,
41        /// The maximum size of a slice to allocate in the pool.
42        max_slice_size: u64,
43        /// Hard cap on the total bytes of pages this pool may hold.
44        ///
45        /// The effective cap is `floor(max_pool_size / page_size)` whole pages.
46        /// If `max_pool_size < page_size`, the page size is shrunk to the
47        /// (alignment-rounded) cap so the budget is honored with a single page.
48        /// When the cap is reached and no free slice fits after coalescing,
49        /// reserving returns [`IoError`](crate::server::IoError)
50        /// `PoolCapacityExceeded` instead of silently growing. `None` (the
51        /// previous behavior) keeps unbounded growth.
52        ///
53        /// Note: runtimes that create one memory management per stream (CUDA,
54        /// HIP) apply the cap per stream.
55        max_pool_size: Option<u64>,
56    },
57}
58
59/// Options to create a memory pool.
60#[derive(Debug, Clone)]
61pub struct MemoryPoolOptions {
62    /// What kind of pool to use.
63    pub pool_type: PoolType,
64    /// Period after which allocations are deemed unused and deallocated.
65    ///
66    /// This period is measured in the number of allocations in the parent allocator. If a page
67    /// in the pool was unused for the entire period, it will be deallocated. This period is
68    /// approximmate, as checks are only done occasionally.
69    pub dealloc_period: Option<u64>,
70}
71
72/// High level configuration of memory management.
73#[derive(Clone, Debug)]
74pub enum MemoryConfiguration {
75    /// The default preset, which uses pools that allocate sub slices.
76    #[cfg(not(exclusive_memory_only))]
77    SubSlices,
78    /// Default preset for using exclusive pages.
79    /// This can be necessary for backends don't support sub-slices.
80    ExclusivePages,
81    /// Custom settings.
82    Custom {
83        /// Options for each pool to construct. When allocating, the first
84        /// possible pool will be picked for an allocation.
85        pool_options: Vec<MemoryPoolOptions>,
86    },
87}
88
89#[allow(clippy::derivable_impls)]
90impl Default for MemoryConfiguration {
91    fn default() -> Self {
92        #[cfg(exclusive_memory_only)]
93        {
94            MemoryConfiguration::ExclusivePages
95        }
96        #[cfg(not(exclusive_memory_only))]
97        {
98            MemoryConfiguration::SubSlices
99        }
100    }
101}
cubecl_runtime/memory_management/mod.rs

cubecl_runtime/memory_management/
mod.rs
