cubecl_runtime/memory_management/
base.rs

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

cubecl_runtime/memory_management/
base.rs
