gpu_allocator/allocator/
mod.rs

1#[cfg(feature = "std")]
2use alloc::sync::Arc;
3use alloc::{fmt, string::String, vec::Vec};
4use core::ops::Range;
5#[cfg(feature = "std")]
6use std::backtrace::Backtrace;
7
8use log::*;
9
10use crate::result::*;
11
12pub(crate) mod dedicated_block_allocator;
13pub(crate) use dedicated_block_allocator::DedicatedBlockAllocator;
14
15pub(crate) mod free_list_allocator;
16pub(crate) use free_list_allocator::FreeListAllocator;
17
18#[derive(PartialEq, Copy, Clone, Debug)]
19#[repr(u8)]
20pub(crate) enum AllocationType {
21    Free,
22    Linear,
23    NonLinear,
24}
25
26impl AllocationType {
27    #[cfg(feature = "visualizer")]
28    pub fn as_str(self) -> &'static str {
29        match self {
30            Self::Free => "Free",
31            Self::Linear => "Linear",
32            Self::NonLinear => "Non-Linear",
33        }
34    }
35}
36
37/// Describes an allocation in the [`AllocatorReport`].
38#[derive(Clone)]
39pub struct AllocationReport {
40    /// The name provided to the `allocate()` function.
41    pub name: String,
42    /// The offset in bytes of the allocation in its memory block.
43    pub offset: u64,
44    /// The size in bytes of the allocation.
45    pub size: u64,
46    #[cfg(feature = "visualizer")]
47    pub(crate) backtrace: Arc<Backtrace>,
48}
49
50/// Describes a memory block in the [`AllocatorReport`].
51#[derive(Clone)]
52pub struct MemoryBlockReport {
53    /// The size in bytes of this memory block.
54    pub size: u64,
55    /// The range of allocations in [`AllocatorReport::allocations`] that are associated
56    /// to this memory block.
57    pub allocations: Range<usize>,
58}
59
60/// A report that can be generated for informational purposes using `Allocator::generate_report()`.
61#[derive(Clone)]
62pub struct AllocatorReport {
63    /// All live allocations, sub-allocated from memory blocks.
64    pub allocations: Vec<AllocationReport>,
65    /// All memory blocks.
66    pub blocks: Vec<MemoryBlockReport>,
67    /// Sum of the memory used by all allocations, in bytes.
68    pub total_allocated_bytes: u64,
69    /// Sum of the memory capacity of all memory blocks including unallocated regions, in bytes.
70    pub total_capacity_bytes: u64,
71}
72
73impl fmt::Debug for AllocationReport {
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        let name = if !self.name.is_empty() {
76            self.name.as_str()
77        } else {
78            "--"
79        };
80        write!(f, "{name:?}: {}", fmt_bytes(self.size))
81    }
82}
83
84impl fmt::Debug for AllocatorReport {
85    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
86        let mut allocations = self.allocations.clone();
87        allocations.sort_by_key(|alloc| core::cmp::Reverse(alloc.size));
88
89        let max_num_allocations_to_print = f.precision().unwrap_or(usize::MAX);
90        allocations.truncate(max_num_allocations_to_print);
91
92        f.debug_struct("AllocatorReport")
93            .field(
94                "summary",
95                &core::format_args!(
96                    "{} / {}",
97                    fmt_bytes(self.total_allocated_bytes),
98                    fmt_bytes(self.total_capacity_bytes)
99                ),
100            )
101            .field("blocks", &self.blocks.len())
102            .field("allocations", &self.allocations.len())
103            .field("largest", &allocations.as_slice())
104            .finish()
105    }
106}
107
108#[cfg(feature = "visualizer")]
109pub(crate) trait SubAllocatorBase: crate::visualizer::SubAllocatorVisualizer {}
110#[cfg(not(feature = "visualizer"))]
111pub(crate) trait SubAllocatorBase {}
112
113pub(crate) trait SubAllocator: SubAllocatorBase + fmt::Debug + Sync + Send {
114    fn allocate(
115        &mut self,
116        size: u64,
117        alignment: u64,
118        allocation_type: AllocationType,
119        granularity: u64,
120        name: &str,
121        #[cfg(feature = "std")] backtrace: Arc<Backtrace>,
122    ) -> Result<(u64, core::num::NonZeroU64)>;
123
124    fn free(&mut self, chunk_id: Option<core::num::NonZeroU64>) -> Result<()>;
125
126    fn rename_allocation(
127        &mut self,
128        chunk_id: Option<core::num::NonZeroU64>,
129        name: &str,
130    ) -> Result<()>;
131
132    fn report_memory_leaks(
133        &self,
134        log_level: Level,
135        memory_type_index: usize,
136        memory_block_index: usize,
137    );
138
139    fn report_allocations(&self) -> Vec<AllocationReport>;
140
141    /// Returns [`true`] if this allocator allows sub-allocating multiple allocations, [`false`] if
142    /// it is designed to only represent dedicated allocations.
143    #[must_use]
144    fn supports_general_allocations(&self) -> bool;
145    #[must_use]
146    fn allocated(&self) -> u64;
147
148    /// Helper function: reports if the suballocator is empty (meaning, having no allocations).
149    #[must_use]
150    fn is_empty(&self) -> bool {
151        self.allocated() == 0
152    }
153}
154
155pub(crate) fn fmt_bytes(mut amount: u64) -> String {
156    const SUFFIX: [&str; 5] = ["B", "KB", "MB", "GB", "TB"];
157
158    let mut idx = 0;
159    let mut print_amount = amount as f64;
160    loop {
161        if amount < 1024 {
162            return format!("{:.2} {}", print_amount, SUFFIX[idx]);
163        }
164
165        print_amount = amount as f64 / 1024.0;
166        amount /= 1024;
167        idx += 1;
168    }
169}