1use {
2 crate::{
3 align_up, error::AllocationError, heap::Heap, slab::Slab, unreachable_unchecked,
4 util::try_arc_unwrap, MemoryBounds,
5 },
6 alloc::{sync::Arc, vec::Vec},
7 core::{convert::TryFrom as _, mem::replace, ptr::NonNull},
8 gpu_alloc_types::{AllocationFlags, DeviceMapError, MemoryDevice, MemoryPropertyFlags},
9};
10
11#[derive(Debug)]
12pub(crate) struct BuddyBlock<M> {
13 pub memory: Arc<M>,
14 pub ptr: Option<NonNull<u8>>,
15 pub offset: u64,
16 pub size: u64,
17 pub chunk: usize,
18 pub index: usize,
19}
20
21unsafe impl<M> Sync for BuddyBlock<M> where M: Sync {}
22unsafe impl<M> Send for BuddyBlock<M> where M: Send {}
23
24#[derive(Clone, Copy, Debug)]
25enum PairState {
26 Exhausted,
27 Ready {
28 ready: Side,
29 next: usize,
30 prev: usize,
31 },
32}
33
34impl PairState {
35 unsafe fn replace_next(&mut self, value: usize) -> usize {
36 match self {
37 PairState::Exhausted => unreachable_unchecked(),
38 PairState::Ready { next, .. } => replace(next, value),
39 }
40 }
41
42 unsafe fn replace_prev(&mut self, value: usize) -> usize {
43 match self {
44 PairState::Exhausted => unreachable_unchecked(),
45 PairState::Ready { prev, .. } => replace(prev, value),
46 }
47 }
48}
49
50#[derive(Clone, Copy, Debug, PartialEq, Eq)]
51enum Side {
52 Left,
53 Right,
54}
55use Side::*;
56
57#[derive(Debug)]
58struct PairEntry {
59 state: PairState,
60 chunk: usize,
61 offset: u64,
62 parent: Option<usize>,
63}
64
65struct SizeBlockEntry {
66 chunk: usize,
67 offset: u64,
68 index: usize,
69}
70
71#[derive(Debug)]
72struct Size {
73 next_ready: usize,
74 pairs: Slab<PairEntry>,
75}
76#[derive(Debug)]
77enum Release {
78 None,
79 Parent(usize),
80 Chunk(usize),
81}
82
83impl Size {
84 fn new() -> Self {
85 Size {
86 pairs: Slab::new(),
87 next_ready: 0,
88 }
89 }
90
91 unsafe fn add_pair_and_acquire_left(
92 &mut self,
93 chunk: usize,
94 offset: u64,
95 parent: Option<usize>,
96 ) -> SizeBlockEntry {
97 if self.next_ready < self.pairs.len() {
98 unreachable_unchecked()
99 }
100
101 let index = self.pairs.insert(PairEntry {
102 state: PairState::Exhausted,
103 chunk,
104 offset,
105 parent,
106 });
107
108 let entry = self.pairs.get_unchecked_mut(index);
109 entry.state = PairState::Ready {
110 next: index,
111 prev: index,
112 ready: Right, };
114 self.next_ready = index;
115
116 SizeBlockEntry {
117 chunk,
118 offset,
119 index: index << 1,
120 }
121 }
122
123 fn acquire(&mut self, size: u64) -> Option<SizeBlockEntry> {
124 if self.next_ready >= self.pairs.len() {
125 return None;
126 }
127
128 let ready = self.next_ready;
129
130 let entry = unsafe { self.pairs.get_unchecked_mut(ready) };
131 let chunk = entry.chunk;
132 let offset = entry.offset;
133
134 let bit = match entry.state {
135 PairState::Exhausted => unsafe { unreachable_unchecked() },
136 PairState::Ready { ready, next, prev } => {
137 entry.state = PairState::Exhausted;
138
139 if prev == self.next_ready {
140 debug_assert_eq!(next, self.next_ready);
142 self.next_ready = self.pairs.len();
143 } else {
144 let prev_entry = unsafe { self.pairs.get_unchecked_mut(prev) };
145 let prev_next = unsafe { prev_entry.state.replace_next(next) };
146 debug_assert_eq!(prev_next, self.next_ready);
147
148 let next_entry = unsafe { self.pairs.get_unchecked_mut(next) };
149 let next_prev = unsafe { next_entry.state.replace_prev(prev) };
150 debug_assert_eq!(next_prev, self.next_ready);
151
152 self.next_ready = next;
153 }
154
155 match ready {
156 Left => 0,
157 Right => 1,
158 }
159 }
160 };
161
162 Some(SizeBlockEntry {
163 chunk,
164 offset: offset + bit as u64 * size,
165 index: (ready << 1) | bit as usize,
166 })
167 }
168
169 fn release(&mut self, index: usize) -> Release {
170 let side = match index & 1 {
171 0 => Side::Left,
172 1 => Side::Right,
173 _ => unsafe { unreachable_unchecked() },
174 };
175 let entry_index = index >> 1;
176
177 let len = self.pairs.len();
178
179 let entry = self.pairs.get_mut(entry_index);
180
181 let chunk = entry.chunk;
182 let offset = entry.offset;
183 let parent = entry.parent;
184
185 match entry.state {
186 PairState::Exhausted => {
187 if self.next_ready == len {
188 entry.state = PairState::Ready {
189 ready: side,
190 next: entry_index,
191 prev: entry_index,
192 };
193 self.next_ready = entry_index;
194 } else {
195 debug_assert!(self.next_ready < len);
196
197 let next = self.next_ready;
198 let next_entry = unsafe { self.pairs.get_unchecked_mut(next) };
199 let prev = unsafe { next_entry.state.replace_prev(entry_index) };
200
201 let prev_entry = unsafe { self.pairs.get_unchecked_mut(prev) };
202 let prev_next = unsafe { prev_entry.state.replace_next(entry_index) };
203 debug_assert_eq!(prev_next, next);
204
205 let entry = unsafe { self.pairs.get_unchecked_mut(entry_index) };
206 entry.state = PairState::Ready {
207 ready: side,
208 next,
209 prev,
210 };
211 }
212 Release::None
213 }
214
215 PairState::Ready { ready, .. } if ready == side => {
216 panic!("Attempt to dealloate already free block")
217 }
218
219 PairState::Ready { next, prev, .. } => {
220 unsafe {
221 self.pairs.remove_unchecked(entry_index);
222 }
223
224 if prev == entry_index {
225 debug_assert_eq!(next, entry_index);
226 self.next_ready = self.pairs.len();
227 } else {
228 let prev_entry = unsafe { self.pairs.get_unchecked_mut(prev) };
229 let prev_next = unsafe { prev_entry.state.replace_next(next) };
230 debug_assert_eq!(prev_next, entry_index);
231
232 let next_entry = unsafe { self.pairs.get_unchecked_mut(next) };
233 let next_prev = unsafe { next_entry.state.replace_prev(prev) };
234 debug_assert_eq!(next_prev, entry_index);
235
236 self.next_ready = next;
237 }
238
239 match parent {
240 Some(parent) => Release::Parent(parent),
241 None => {
242 debug_assert_eq!(offset, 0);
243 Release::Chunk(chunk)
244 }
245 }
246 }
247 }
248 }
249}
250
251#[derive(Debug)]
252struct Chunk<M> {
253 memory: Arc<M>,
254 ptr: Option<NonNull<u8>>,
255 size: u64,
256}
257
258#[derive(Debug)]
259pub(crate) struct BuddyAllocator<M> {
260 minimal_size: u64,
261 chunks: Slab<Chunk<M>>,
262 sizes: Vec<Size>,
263 memory_type: u32,
264 props: MemoryPropertyFlags,
265 atom_mask: u64,
266}
267
268unsafe impl<M> Sync for BuddyAllocator<M> where M: Sync {}
269unsafe impl<M> Send for BuddyAllocator<M> where M: Send {}
270
271impl<M> BuddyAllocator<M>
272where
273 M: MemoryBounds + 'static,
274{
275 pub fn new(
276 minimal_size: u64,
277 initial_dedicated_size: u64,
278 memory_type: u32,
279 props: MemoryPropertyFlags,
280 atom_mask: u64,
281 ) -> Self {
282 assert!(
283 minimal_size.is_power_of_two(),
284 "Minimal allocation size of buddy allocator must be power of two"
285 );
286 assert!(
287 initial_dedicated_size.is_power_of_two(),
288 "Dedicated allocation size of buddy allocator must be power of two"
289 );
290
291 let initial_sizes = (initial_dedicated_size
292 .trailing_zeros()
293 .saturating_sub(minimal_size.trailing_zeros())) as usize;
294
295 BuddyAllocator {
296 minimal_size,
297 chunks: Slab::new(),
298 sizes: (0..initial_sizes).map(|_| Size::new()).collect(),
299 memory_type,
300 props,
301 atom_mask: atom_mask | (minimal_size - 1),
302 }
303 }
304
305 #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, device)))]
306 pub unsafe fn alloc(
307 &mut self,
308 device: &impl MemoryDevice<M>,
309 size: u64,
310 align_mask: u64,
311 flags: AllocationFlags,
312 heap: &mut Heap,
313 allocations_remains: &mut u32,
314 ) -> Result<BuddyBlock<M>, AllocationError> {
315 let align_mask = align_mask | self.atom_mask;
316
317 let size = align_up(size, align_mask)
318 .and_then(|size| size.checked_next_power_of_two())
319 .ok_or(AllocationError::OutOfDeviceMemory)?;
320
321 let size = size.max(self.minimal_size);
322
323 let size_index = size.trailing_zeros() - self.minimal_size.trailing_zeros();
324 let size_index =
325 usize::try_from(size_index).map_err(|_| AllocationError::OutOfDeviceMemory)?;
326
327 while self.sizes.len() <= size_index {
328 self.sizes.push(Size::new());
329 }
330
331 let host_visible = self.host_visible();
332
333 let mut candidate_size_index = size_index;
334
335 let (mut entry, entry_size_index) = loop {
336 let sizes_len = self.sizes.len();
337
338 let candidate_size_entry = &mut self.sizes[candidate_size_index];
339 let candidate_size = self.minimal_size << candidate_size_index;
340
341 if let Some(entry) = candidate_size_entry.acquire(candidate_size) {
342 break (entry, candidate_size_index);
343 }
344
345 if sizes_len == candidate_size_index + 1 {
346 if *allocations_remains == 0 {
348 return Err(AllocationError::TooManyObjects);
349 }
350
351 let chunk_size = self.minimal_size << (candidate_size_index + 1);
352 let mut memory = device.allocate_memory(chunk_size, self.memory_type, flags)?;
353 *allocations_remains -= 1;
354 heap.alloc(chunk_size);
355
356 let ptr = if host_visible {
357 match device.map_memory(&mut memory, 0, chunk_size) {
358 Ok(ptr) => Some(ptr),
359 Err(DeviceMapError::OutOfDeviceMemory) => {
360 return Err(AllocationError::OutOfDeviceMemory)
361 }
362 Err(DeviceMapError::MapFailed) | Err(DeviceMapError::OutOfHostMemory) => {
363 return Err(AllocationError::OutOfHostMemory)
364 }
365 }
366 } else {
367 None
368 };
369
370 let chunk = self.chunks.insert(Chunk {
371 memory: Arc::new(memory),
372 ptr,
373 size: chunk_size,
374 });
375
376 let entry = candidate_size_entry.add_pair_and_acquire_left(chunk, 0, None);
377
378 break (entry, candidate_size_index);
379 }
380
381 candidate_size_index += 1;
382 };
383
384 for size_index in (size_index..entry_size_index).rev() {
385 let size_entry = &mut self.sizes[size_index];
386 entry =
387 size_entry.add_pair_and_acquire_left(entry.chunk, entry.offset, Some(entry.index));
388 }
389
390 let chunk_entry = self.chunks.get_unchecked(entry.chunk);
391
392 debug_assert!(
393 entry
394 .offset
395 .checked_add(size)
396 .map_or(false, |end| end <= chunk_entry.size),
397 "Offset + size is not in chunk bounds"
398 );
399
400 Ok(BuddyBlock {
401 memory: chunk_entry.memory.clone(),
402 ptr: chunk_entry
403 .ptr
404 .map(|ptr| NonNull::new_unchecked(ptr.as_ptr().add(entry.offset as usize))),
405 offset: entry.offset,
406 size,
407 chunk: entry.chunk,
408 index: entry.index,
409 })
410 }
411
412 #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, device)))]
413 pub unsafe fn dealloc(
414 &mut self,
415 device: &impl MemoryDevice<M>,
416 block: BuddyBlock<M>,
417 heap: &mut Heap,
418 allocations_remains: &mut u32,
419 ) {
420 debug_assert!(block.size.is_power_of_two());
421
422 let size_index =
423 (block.size.trailing_zeros() - self.minimal_size.trailing_zeros()) as usize;
424
425 let mut release_index = block.index;
426 let mut release_size_index = size_index;
427
428 loop {
429 match self.sizes[release_size_index].release(release_index) {
430 Release::Parent(parent) => {
431 release_size_index += 1;
432 release_index = parent;
433 }
434 Release::Chunk(chunk) => {
435 debug_assert_eq!(chunk, block.chunk);
436 debug_assert_eq!(
437 self.chunks.get(chunk).size,
438 self.minimal_size << (release_size_index + 1)
439 );
440 let chunk = self.chunks.remove(chunk);
441 drop(block);
442
443 let memory = try_arc_unwrap(chunk.memory)
444 .expect("Memory shared after last block deallocated");
445
446 device.deallocate_memory(memory);
447 *allocations_remains += 1;
448 heap.dealloc(chunk.size);
449
450 return;
451 }
452 Release::None => return,
453 }
454 }
455 }
456
457 fn host_visible(&self) -> bool {
458 self.props.contains(MemoryPropertyFlags::HOST_VISIBLE)
459 }
460}