cubecl_runtime/memory_management/
base.rs

1#[cfg(not(feature = "std"))]
2use alloc::{format, string::String};
3
4/// Amount of memory in use by this allocator
5/// and statistics on how much memory is reserved and
6/// wasted in total.
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct MemoryUsage {
9    /// The number of allocations currently active.
10    pub number_allocs: u64,
11    /// The number of bytes that are currently actually in use.
12    ///
13    /// This doesn't include any padding or other memory that needs to be
14    /// reserved, and is the minimum amount of memory that could possible
15    /// be allocated.
16    pub bytes_in_use: u64,
17    /// The amount of bytes used for padding memory in currently active allocations.
18    pub bytes_padding: u64,
19    /// The total amount of memory reserved on the device.
20    ///
21    /// This will be at least as much as bytes_in_use but in practice will
22    /// be higher, as allocations reserve memory for future allocations
23    /// and for padding.
24    pub bytes_reserved: u64,
25}
26
27impl MemoryUsage {
28    /// Calculate the combined memory usage of two reports (summing them).
29    pub fn combine(&self, other: MemoryUsage) -> MemoryUsage {
30        MemoryUsage {
31            number_allocs: self.number_allocs + other.number_allocs,
32            bytes_in_use: self.bytes_in_use + other.bytes_in_use,
33            bytes_padding: self.bytes_padding + other.bytes_padding,
34            bytes_reserved: self.bytes_reserved + other.bytes_reserved,
35        }
36    }
37}
38
39fn bytes_format(bytes: u64) -> String {
40    let unit = 1000;
41
42    if bytes < unit {
43        format!("{bytes} B")
44    } else {
45        let size = bytes as f64;
46        let exp = match size.log(1000.0).floor() as usize {
47            0 => 1,
48            e => e,
49        };
50        let unit_prefix = "KMGTPEZY".as_bytes();
51        format!(
52            "{:.2} {}B",
53            (size / unit.pow(exp as u32) as f64),
54            unit_prefix[exp - 1] as char,
55        )
56    }
57}
58
59impl core::fmt::Display for MemoryUsage {
60    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
61        // In the future it'd be nice if MemoryUsage also held some stats about say,
62        // the 5 biggest allocations, to show when you an OOM.
63        let usage_percentage = (self.bytes_in_use as f32 / self.bytes_reserved as f32) * 100.0;
64        let padding_percentage = (self.bytes_padding as f32 / self.bytes_in_use as f32) * 100.0;
65        writeln!(f, "Memory Usage Report:")?;
66        writeln!(f, "  Number of allocations: {}", self.number_allocs)?;
67        writeln!(f, "  Bytes in use: {}", bytes_format(self.bytes_in_use))?;
68        writeln!(
69            f,
70            "  Bytes used for padding: {}",
71            bytes_format(self.bytes_padding)
72        )?;
73        writeln!(
74            f,
75            "  Total bytes reserved: {}",
76            bytes_format(self.bytes_reserved)
77        )?;
78        writeln!(f, "  Usage efficiency: {usage_percentage:.2}%")?;
79        writeln!(f, "  Padding overhead: {padding_percentage:.2}%")
80    }
81}
82
83/// The managed tensor buffer handle that points to some memory segment.
84/// It should not contain actual data.
85pub trait MemoryHandle<Binding>: Clone + Send + Sync + core::fmt::Debug {
86    /// Checks if the underlying memory can be safely mutated.
87    fn can_mut(&self) -> bool;
88    /// Get the binding associated to the current handle.
89    fn binding(self) -> Binding;
90}
cubecl_runtime/memory_management/base.rs

cubecl_runtime/memory_management/
base.rs
