molgfx_gpu/residency/
arena.rs

1//! Deterministic page-granular suballocation.
2//!
3//! Allocation is first-fit over address-sorted free runs. Allocation is
4//! `O(free runs + allocated pages)` and release is
5//! `O(free runs + released pages)`. All metadata is reserved at construction,
6//! so neither operation grows the heap after warm-up.
7
8use thiserror::Error;
9
10#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11struct FreeRun {
12    start: u32,
13    pages: u32,
14}
15
16#[derive(Clone, Copy, Debug, PartialEq, Eq)]
17enum PageOwner {
18    Free,
19    Head { generation: u64, pages: u32 },
20    Tail { head: u32, generation: u64 },
21}
22
23/// A generation-checked allocation within a [`PagedArena`].
24#[derive(Clone, Copy, Debug, PartialEq, Eq)]
25pub struct ArenaAllocation {
26    start_page: u32,
27    page_count: u32,
28    generation: u64,
29    byte_offset: u64,
30    byte_len: u64,
31}
32
33impl ArenaAllocation {
34    /// First page occupied by this allocation.
35    #[must_use]
36    pub const fn start_page(self) -> u32 {
37        self.start_page
38    }
39
40    /// Number of whole pages reserved.
41    #[must_use]
42    pub const fn page_count(self) -> u32 {
43        self.page_count
44    }
45
46    /// Byte offset suitable for binding or copy commands.
47    #[must_use]
48    pub const fn byte_offset(self) -> u64 {
49        self.byte_offset
50    }
51
52    /// Requested byte length, excluding page padding.
53    #[must_use]
54    pub const fn byte_len(self) -> u64 {
55        self.byte_len
56    }
57}
58
59/// Observable arena counters. Values are cumulative except resident bytes.
60#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
61pub struct ArenaMetrics {
62    /// Bytes currently covered by live pages.
63    pub resident_bytes: u64,
64    /// Highest live page footprint observed.
65    pub peak_resident_bytes: u64,
66    /// Sum of requested payload bytes accepted.
67    pub requested_bytes: u64,
68    /// Sum of page-rounded bytes accepted.
69    pub allocated_bytes: u64,
70    /// Number of successful allocations.
71    pub allocations: u64,
72    /// Number of requests rejected for capacity.
73    pub allocation_stalls: u64,
74    /// Payload bytes rejected for capacity.
75    pub stalled_bytes: u64,
76    /// Host storage allocations performed by this primitive.
77    pub host_allocation_events: u64,
78}
79
80/// Invalid requests and stale allocation handles.
81#[derive(Clone, Copy, Debug, Error, PartialEq, Eq)]
82pub enum ArenaError {
83    /// Page size and page count must both be non-zero.
84    #[error("page size and page count must both be non-zero")]
85    InvalidConfiguration,
86    /// Zero-byte allocations are not meaningful.
87    #[error("an arena allocation must contain at least one byte")]
88    EmptyAllocation,
89    /// The requested byte count cannot be represented by this arena.
90    #[error("arena allocation size overflow")]
91    SizeOverflow,
92    /// No contiguous free run can satisfy the request.
93    #[error("no contiguous page run can satisfy {requested_bytes} bytes")]
94    OutOfMemory {
95        /// Payload size that could not be placed.
96        requested_bytes: u64,
97    },
98    /// The handle was already freed or belongs to an older allocation.
99    #[error("arena allocation is stale")]
100    StaleAllocation,
101    /// A fixed-capacity arena cannot be extended by zero pages.
102    #[error("arena growth must add at least one page")]
103    EmptyGrowth,
104    /// The enlarged page count cannot be represented by the arena.
105    #[error("arena page count overflow")]
106    PageCountOverflow,
107}
108
109/// A fixed-capacity page allocator with deterministic first-fit placement.
110#[derive(Debug)]
111pub struct PagedArena {
112    page_size: u64,
113    page_count: u32,
114    owners: Vec<PageOwner>,
115    free_runs: Vec<FreeRun>,
116    next_generation: u64,
117    metrics: ArenaMetrics,
118}
119
120impl PagedArena {
121    /// Reserves all allocator metadata up front.
122    ///
123    /// # Errors
124    ///
125    /// Returns [`ArenaError::InvalidConfiguration`] for an empty arena.
126    pub fn new(page_size: u64, page_count: u32) -> Result<Self, ArenaError> {
127        if page_size == 0 || page_count == 0 {
128            return Err(ArenaError::InvalidConfiguration);
129        }
130        let owner_count = page_count as usize;
131        let owners = vec![PageOwner::Free; owner_count];
132        let mut free_runs = Vec::with_capacity(owner_count);
133        free_runs.push(FreeRun {
134            start: 0,
135            pages: page_count,
136        });
137        Ok(Self {
138            page_size,
139            page_count,
140            owners,
141            free_runs,
142            next_generation: 1,
143            metrics: ArenaMetrics {
144                host_allocation_events: 2,
145                ..ArenaMetrics::default()
146            },
147        })
148    }
149
150    /// Reserves the first address-ordered run large enough for `byte_len`.
151    ///
152    /// # Errors
153    ///
154    /// Empty, overflowing and unsatisfied requests are returned as values.
155    pub fn allocate(&mut self, byte_len: u64) -> Result<ArenaAllocation, ArenaError> {
156        if byte_len == 0 {
157            return Err(ArenaError::EmptyAllocation);
158        }
159        let rounded = byte_len
160            .checked_add(self.page_size - 1)
161            .ok_or(ArenaError::SizeOverflow)?;
162        let page_count_u64 = rounded / self.page_size;
163        let page_count = u32::try_from(page_count_u64).map_err(|_| ArenaError::SizeOverflow)?;
164        let Some(run_index) = self
165            .free_runs
166            .iter()
167            .position(|run| run.pages >= page_count)
168        else {
169            self.metrics.allocation_stalls = self.metrics.allocation_stalls.saturating_add(1);
170            self.metrics.stalled_bytes = self.metrics.stalled_bytes.saturating_add(byte_len);
171            return Err(ArenaError::OutOfMemory {
172                requested_bytes: byte_len,
173            });
174        };
175        let start_page = self.take_pages(run_index, page_count);
176        let generation = self.next_generation;
177        self.next_generation = self.next_generation.wrapping_add(1).max(1);
178        self.mark_owned(start_page, page_count, generation);
179        let allocated = u64::from(page_count) * self.page_size;
180        self.metrics.resident_bytes = self.metrics.resident_bytes.saturating_add(allocated);
181        self.metrics.peak_resident_bytes = self
182            .metrics
183            .peak_resident_bytes
184            .max(self.metrics.resident_bytes);
185        self.metrics.requested_bytes = self.metrics.requested_bytes.saturating_add(byte_len);
186        self.metrics.allocated_bytes = self.metrics.allocated_bytes.saturating_add(allocated);
187        self.metrics.allocations = self.metrics.allocations.saturating_add(1);
188        Ok(ArenaAllocation {
189            start_page,
190            page_count,
191            generation,
192            byte_offset: u64::from(start_page) * self.page_size,
193            byte_len,
194        })
195    }
196
197    /// Releases a live allocation and coalesces adjacent free runs.
198    ///
199    /// # Errors
200    ///
201    /// Returns [`ArenaError::StaleAllocation`] when generation or extent no
202    /// longer matches the arena.
203    pub fn release(&mut self, allocation: ArenaAllocation) -> Result<(), ArenaError> {
204        if !self.is_live(allocation) {
205            return Err(ArenaError::StaleAllocation);
206        }
207        let start = allocation.start_page as usize;
208        let end = start + allocation.page_count as usize;
209        self.owners[start..end].fill(PageOwner::Free);
210        self.insert_free_run(FreeRun {
211            start: allocation.start_page,
212            pages: allocation.page_count,
213        });
214        self.metrics.resident_bytes -= u64::from(allocation.page_count) * self.page_size;
215        Ok(())
216    }
217
218    /// Appends free pages without moving any live allocation.
219    ///
220    /// Existing byte offsets remain stable, allowing a physical owner to copy
221    /// the old address space verbatim into a larger GPU buffer.
222    ///
223    /// # Errors
224    ///
225    /// Zero growth and page-count overflow are rejected as typed values.
226    pub fn grow(&mut self, additional_pages: u32) -> Result<(), ArenaError> {
227        if additional_pages == 0 {
228            return Err(ArenaError::EmptyGrowth);
229        }
230        let old_pages = self.page_count;
231        self.page_count = old_pages
232            .checked_add(additional_pages)
233            .ok_or(ArenaError::PageCountOverflow)?;
234        self.owners.reserve(additional_pages as usize);
235        self.owners.extend(std::iter::repeat_n(
236            PageOwner::Free,
237            additional_pages as usize,
238        ));
239        self.metrics.host_allocation_events = self.metrics.host_allocation_events.saturating_add(1);
240        self.insert_free_run(FreeRun {
241            start: old_pages,
242            pages: additional_pages,
243        });
244        Ok(())
245    }
246
247    /// Current cumulative counters.
248    #[must_use]
249    pub const fn metrics(&self) -> ArenaMetrics {
250        self.metrics
251    }
252
253    /// Total bytes represented by the arena.
254    #[must_use]
255    pub fn capacity_bytes(&self) -> u64 {
256        self.owners.len() as u64 * self.page_size
257    }
258
259    /// Number of addressable pages.
260    #[must_use]
261    pub fn page_count(&self) -> u32 {
262        self.page_count
263    }
264
265    fn take_pages(&mut self, index: usize, pages: u32) -> u32 {
266        let start = self.free_runs[index].start;
267        if self.free_runs[index].pages == pages {
268            self.free_runs.remove(index);
269        } else {
270            self.free_runs[index].start += pages;
271            self.free_runs[index].pages -= pages;
272        }
273        start
274    }
275
276    fn mark_owned(&mut self, start: u32, pages: u32, generation: u64) {
277        let first = start as usize;
278        self.owners[first] = PageOwner::Head { generation, pages };
279        let end = first + pages as usize;
280        for owner in &mut self.owners[first + 1..end] {
281            *owner = PageOwner::Tail {
282                head: start,
283                generation,
284            };
285        }
286    }
287
288    fn is_live(&self, allocation: ArenaAllocation) -> bool {
289        let Some(owner) = self.owners.get(allocation.start_page as usize) else {
290            return false;
291        };
292        matches!(
293            owner,
294            PageOwner::Head { generation, pages }
295                if *generation == allocation.generation && *pages == allocation.page_count
296        )
297    }
298
299    fn insert_free_run(&mut self, run: FreeRun) {
300        let index = self
301            .free_runs
302            .partition_point(|candidate| candidate.start < run.start);
303        self.free_runs.insert(index, run);
304        let merge_index = index.saturating_sub(1);
305        self.coalesce_from(merge_index);
306    }
307
308    fn coalesce_from(&mut self, index: usize) {
309        let cursor = index;
310        while cursor + 1 < self.free_runs.len() {
311            let current = self.free_runs[cursor];
312            let next = self.free_runs[cursor + 1];
313            if current.start + current.pages != next.start {
314                break;
315            }
316            self.free_runs[cursor].pages += next.pages;
317            self.free_runs.remove(cursor + 1);
318        }
319    }
320}
321
322#[cfg(test)]
323#[path = "arena_tests.rs"]
324mod tests;
molgfx_gpu/residency/arena.rs

molgfx_gpu/residency/
arena.rs
