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

cubecl_runtime/memory_management/
base.rs
