cubecl_runtime/memory_management/
policy.rs

1//! How memory is managed, and where the pools that are not the dynamic ones
2//! live.
3
4/// Whether this build refuses pools that share a page between allocations —
5/// the `exclusive-memory-only` feature, or a wasm target.
6///
7/// Decided by this crate alone, since [`MemoryConfiguration::Adaptive`] only
8/// exists under it; a crate that needs to know reads it here instead of
9/// deciding again.
10pub const EXCLUSIVE_MEMORY_ONLY: bool = cfg!(exclusive_memory_only);
11
12/// High level configuration of memory management.
13#[derive(Clone, Debug)]
14pub enum MemoryConfiguration {
15    /// One page per allocation, in exponentially spaced size buckets: what a
16    /// device that cannot sub-slice gets, and what a staging or uniform pool
17    /// wants whatever the device.
18    ExclusivePages,
19    /// Small allocations in a sliced pool of their own, everything else carved
20    /// from pages sized to the largest allocation served so far — nothing to
21    /// measure or configure per workload. When an allocation outgrows the
22    /// pages, pages of the new size take over and the old ones are returned
23    /// as they empty, or sooner when what lives on them is relocated. The
24    /// default where sub-slicing is available.
25    #[cfg(not(exclusive_memory_only))]
26    Adaptive,
27}
28
29#[allow(clippy::derivable_impls)]
30impl Default for MemoryConfiguration {
31    fn default() -> Self {
32        #[cfg(exclusive_memory_only)]
33        {
34            MemoryConfiguration::ExclusivePages
35        }
36        #[cfg(not(exclusive_memory_only))]
37        {
38            MemoryConfiguration::Adaptive
39        }
40    }
41}
42
43#[derive(Default, Clone, Copy, Debug)]
44/// The mode of allocation used.
45pub enum MemoryAllocationMode {
46    /// Use the automatic memory management strategy for allocation.
47    #[default]
48    Auto,
49    /// Use a persistent memory management strategy, meaning that all allocations are for data that is
50    /// likely never going to be freed.
51    Persistent,
52    /// Give every allocation its own device allocation, returned to the driver
53    /// once it is freed, outside every pool: for a buffer that exists only for
54    /// one measurement (the memory-bandwidth probe) and must neither stay
55    /// reserved nor shape the pools' sizing.
56    Dedicated,
57}
58
59/// The pool position stamped on persistent-pool slices, routing their binds
60/// and lookups to the persistent pool. A fixed sentinel (rather than "one past
61/// the dynamic pools") so a persistent slice stays routable however many
62/// dynamic pools a workload leaves behind.
63#[doc(hidden)]
64pub const PERSISTENT_POOL_POS: u8 = u8::MAX;
65
66/// The pool position stamped on dedicated allocations
67/// ([`MemoryAllocationMode::Dedicated`]),
68/// a fixed sentinel for the same reason as [`PERSISTENT_POOL_POS`].
69#[doc(hidden)]
70pub const DEDICATED_POOL_POS: u8 = u8::MAX - 1;
cubecl_runtime/memory_management/policy.rs

cubecl_runtime/memory_management/
policy.rs
