cubecl_runtime/memory_management/
base.rs

1#[cfg(not(feature = "std"))]
2use alloc::string::{String, ToString};
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
39#[derive(new)]
40pub(crate) struct BytesFormat {
41    bytes: u64,
42}
43
44impl core::fmt::Display for BytesFormat {
45    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
46        let unit = 1000;
47
48        if self.bytes < unit {
49            f.write_fmt(format_args!("{} B", self.bytes))
50        } else {
51            let size = self.bytes as f64;
52            let exp = match size.log(1000.0).floor() as usize {
53                0 => 1,
54                e => e,
55            };
56            let unit_prefix = "KMGTPEZY".as_bytes();
57            f.write_fmt(format_args!(
58                "{:.2} {}B",
59                (size / unit.pow(exp as u32) as f64),
60                unit_prefix[exp - 1] as char,
61            ))
62        }
63    }
64}
65
66fn bytes_format(bytes: u64) -> String {
67    BytesFormat::new(bytes).to_string()
68}
69
70impl core::fmt::Display for MemoryUsage {
71    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
72        // In the future it'd be nice if MemoryUsage also held some stats about say,
73        // the 5 biggest allocations, to show when you an OOM.
74        let usage_percentage = (self.bytes_in_use as f32 / self.bytes_reserved as f32) * 100.0;
75        let padding_percentage = (self.bytes_padding as f32 / self.bytes_in_use as f32) * 100.0;
76        writeln!(f, "Memory Usage Report:")?;
77        writeln!(f, "  Number of allocations: {}", self.number_allocs)?;
78        writeln!(f, "  Bytes in use: {}", bytes_format(self.bytes_in_use))?;
79        writeln!(
80            f,
81            "  Bytes used for padding: {}",
82            bytes_format(self.bytes_padding)
83        )?;
84        writeln!(
85            f,
86            "  Total bytes reserved: {}",
87            bytes_format(self.bytes_reserved)
88        )?;
89        writeln!(f, "  Usage efficiency: {usage_percentage:.2}%")?;
90        writeln!(f, "  Padding overhead: {padding_percentage:.2}%")
91    }
92}
93
94/// The managed tensor buffer handle that points to some memory segment.
95/// It should not contain actual data.
96pub trait MemoryHandle<Binding>: Clone + Send + Sync + core::fmt::Debug {
97    /// Checks if the underlying memory can be safely mutated.
98    fn can_mut(&self) -> bool;
99    /// Get the binding associated to the current handle.
100    fn binding(self) -> Binding;
101}
cubecl_runtime/memory_management/base.rs

cubecl_runtime/memory_management/
base.rs
