diff options
Diffstat (limited to 'third_party/rust/gpu-alloc/src/buddy.rs')
| -rw-r--r-- | third_party/rust/gpu-alloc/src/buddy.rs | 460 |
1 files changed, 460 insertions, 0 deletions
diff --git a/third_party/rust/gpu-alloc/src/buddy.rs b/third_party/rust/gpu-alloc/src/buddy.rs new file mode 100644 index 0000000000..dea303b948 --- /dev/null +++ b/third_party/rust/gpu-alloc/src/buddy.rs @@ -0,0 +1,460 @@ +use {
+ crate::{
+ align_up, error::AllocationError, heap::Heap, slab::Slab, unreachable_unchecked,
+ util::try_arc_unwrap, MemoryBounds,
+ },
+ alloc::{sync::Arc, vec::Vec},
+ core::{convert::TryFrom as _, mem::replace, ptr::NonNull},
+ gpu_alloc_types::{AllocationFlags, DeviceMapError, MemoryDevice, MemoryPropertyFlags},
+};
+
+#[derive(Debug)]
+pub(crate) struct BuddyBlock<M> {
+ pub memory: Arc<M>,
+ pub ptr: Option<NonNull<u8>>,
+ pub offset: u64,
+ pub size: u64,
+ pub chunk: usize,
+ pub index: usize,
+}
+
+unsafe impl<M> Sync for BuddyBlock<M> where M: Sync {}
+unsafe impl<M> Send for BuddyBlock<M> where M: Send {}
+
+#[derive(Clone, Copy, Debug)]
+enum PairState {
+ Exhausted,
+ Ready {
+ ready: Side,
+ next: usize,
+ prev: usize,
+ },
+}
+
+impl PairState {
+ unsafe fn replace_next(&mut self, value: usize) -> usize {
+ match self {
+ PairState::Exhausted => unreachable_unchecked(),
+ PairState::Ready { next, .. } => replace(next, value),
+ }
+ }
+
+ unsafe fn replace_prev(&mut self, value: usize) -> usize {
+ match self {
+ PairState::Exhausted => unreachable_unchecked(),
+ PairState::Ready { prev, .. } => replace(prev, value),
+ }
+ }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum Side {
+ Left,
+ Right,
+}
+use Side::*;
+
+#[derive(Debug)]
+struct PairEntry {
+ state: PairState,
+ chunk: usize,
+ offset: u64,
+ parent: Option<usize>,
+}
+
+struct SizeBlockEntry {
+ chunk: usize,
+ offset: u64,
+ index: usize,
+}
+
+#[derive(Debug)]
+struct Size {
+ next_ready: usize,
+ pairs: Slab<PairEntry>,
+}
+#[derive(Debug)]
+enum Release {
+ None,
+ Parent(usize),
+ Chunk(usize),
+}
+
+impl Size {
+ fn new() -> Self {
+ Size {
+ pairs: Slab::new(),
+ next_ready: 0,
+ }
+ }
+
+ unsafe fn add_pair_and_acquire_left(
+ &mut self,
+ chunk: usize,
+ offset: u64,
+ parent: Option<usize>,
+ ) -> SizeBlockEntry {
+ if self.next_ready < self.pairs.len() {
+ unreachable_unchecked()
+ }
+
+ let index = self.pairs.insert(PairEntry {
+ state: PairState::Exhausted,
+ chunk,
+ offset,
+ parent,
+ });
+
+ let entry = self.pairs.get_unchecked_mut(index);
+ entry.state = PairState::Ready {
+ next: index,
+ prev: index,
+ ready: Right, // Left is allocated.
+ };
+ self.next_ready = index;
+
+ SizeBlockEntry {
+ chunk,
+ offset,
+ index: index << 1,
+ }
+ }
+
+ fn acquire(&mut self, size: u64) -> Option<SizeBlockEntry> {
+ if self.next_ready >= self.pairs.len() {
+ return None;
+ }
+
+ let ready = self.next_ready;
+
+ let entry = unsafe { self.pairs.get_unchecked_mut(ready) };
+ let chunk = entry.chunk;
+ let offset = entry.offset;
+
+ let bit = match entry.state {
+ PairState::Exhausted => unsafe { unreachable_unchecked() },
+ PairState::Ready { ready, next, prev } => {
+ entry.state = PairState::Exhausted;
+
+ if prev == self.next_ready {
+ // The only ready entry.
+ debug_assert_eq!(next, self.next_ready);
+ self.next_ready = self.pairs.len();
+ } else {
+ let prev_entry = unsafe { self.pairs.get_unchecked_mut(prev) };
+ let prev_next = unsafe { prev_entry.state.replace_next(next) };
+ debug_assert_eq!(prev_next, self.next_ready);
+
+ let next_entry = unsafe { self.pairs.get_unchecked_mut(next) };
+ let next_prev = unsafe { next_entry.state.replace_prev(prev) };
+ debug_assert_eq!(next_prev, self.next_ready);
+
+ self.next_ready = next;
+ }
+
+ match ready {
+ Left => 0,
+ Right => 1,
+ }
+ }
+ };
+
+ Some(SizeBlockEntry {
+ chunk,
+ offset: offset + bit as u64 * size,
+ index: (ready << 1) | bit as usize,
+ })
+ }
+
+ fn release(&mut self, index: usize) -> Release {
+ let side = match index & 1 {
+ 0 => Side::Left,
+ 1 => Side::Right,
+ _ => unsafe { unreachable_unchecked() },
+ };
+ let entry_index = index >> 1;
+
+ let len = self.pairs.len();
+
+ let entry = self.pairs.get_mut(entry_index);
+
+ let chunk = entry.chunk;
+ let offset = entry.offset;
+ let parent = entry.parent;
+
+ match entry.state {
+ PairState::Exhausted => {
+ if self.next_ready == len {
+ entry.state = PairState::Ready {
+ ready: side,
+ next: entry_index,
+ prev: entry_index,
+ };
+ self.next_ready = entry_index;
+ } else {
+ debug_assert!(self.next_ready < len);
+
+ let next = self.next_ready;
+ let next_entry = unsafe { self.pairs.get_unchecked_mut(next) };
+ let prev = unsafe { next_entry.state.replace_prev(entry_index) };
+
+ let prev_entry = unsafe { self.pairs.get_unchecked_mut(prev) };
+ let prev_next = unsafe { prev_entry.state.replace_next(entry_index) };
+ debug_assert_eq!(prev_next, next);
+
+ let entry = unsafe { self.pairs.get_unchecked_mut(entry_index) };
+ entry.state = PairState::Ready {
+ ready: side,
+ next,
+ prev,
+ };
+ }
+ Release::None
+ }
+
+ PairState::Ready { ready, .. } if ready == side => {
+ panic!("Attempt to dealloate already free block")
+ }
+
+ PairState::Ready { next, prev, .. } => {
+ unsafe {
+ self.pairs.remove_unchecked(entry_index);
+ }
+
+ if prev == entry_index {
+ debug_assert_eq!(next, entry_index);
+ self.next_ready = self.pairs.len();
+ } else {
+ let prev_entry = unsafe { self.pairs.get_unchecked_mut(prev) };
+ let prev_next = unsafe { prev_entry.state.replace_next(next) };
+ debug_assert_eq!(prev_next, entry_index);
+
+ let next_entry = unsafe { self.pairs.get_unchecked_mut(next) };
+ let next_prev = unsafe { next_entry.state.replace_prev(prev) };
+ debug_assert_eq!(next_prev, entry_index);
+
+ self.next_ready = next;
+ }
+
+ match parent {
+ Some(parent) => Release::Parent(parent),
+ None => {
+ debug_assert_eq!(offset, 0);
+ Release::Chunk(chunk)
+ }
+ }
+ }
+ }
+ }
+}
+
+#[derive(Debug)]
+struct Chunk<M> {
+ memory: Arc<M>,
+ ptr: Option<NonNull<u8>>,
+ size: u64,
+}
+
+#[derive(Debug)]
+pub(crate) struct BuddyAllocator<M> {
+ minimal_size: u64,
+ chunks: Slab<Chunk<M>>,
+ sizes: Vec<Size>,
+ memory_type: u32,
+ props: MemoryPropertyFlags,
+ atom_mask: u64,
+}
+
+unsafe impl<M> Sync for BuddyAllocator<M> where M: Sync {}
+unsafe impl<M> Send for BuddyAllocator<M> where M: Send {}
+
+impl<M> BuddyAllocator<M>
+where
+ M: MemoryBounds + 'static,
+{
+ pub fn new(
+ minimal_size: u64,
+ initial_dedicated_size: u64,
+ memory_type: u32,
+ props: MemoryPropertyFlags,
+ atom_mask: u64,
+ ) -> Self {
+ assert!(
+ minimal_size.is_power_of_two(),
+ "Minimal allocation size of buddy allocator must be power of two"
+ );
+ assert!(
+ initial_dedicated_size.is_power_of_two(),
+ "Dedicated allocation size of buddy allocator must be power of two"
+ );
+
+ let initial_sizes = (initial_dedicated_size
+ .trailing_zeros()
+ .saturating_sub(minimal_size.trailing_zeros())) as usize;
+
+ BuddyAllocator {
+ minimal_size,
+ chunks: Slab::new(),
+ sizes: (0..initial_sizes).map(|_| Size::new()).collect(),
+ memory_type,
+ props,
+ atom_mask: atom_mask | (minimal_size - 1),
+ }
+ }
+
+ #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, device)))]
+ pub unsafe fn alloc(
+ &mut self,
+ device: &impl MemoryDevice<M>,
+ size: u64,
+ align_mask: u64,
+ flags: AllocationFlags,
+ heap: &mut Heap,
+ allocations_remains: &mut u32,
+ ) -> Result<BuddyBlock<M>, AllocationError> {
+ let align_mask = align_mask | self.atom_mask;
+
+ let size = align_up(size, align_mask)
+ .and_then(|size| size.checked_next_power_of_two())
+ .ok_or(AllocationError::OutOfDeviceMemory)?;
+
+ let size = size.max(self.minimal_size);
+
+ let size_index = size.trailing_zeros() - self.minimal_size.trailing_zeros();
+ let size_index =
+ usize::try_from(size_index).map_err(|_| AllocationError::OutOfDeviceMemory)?;
+
+ while self.sizes.len() <= size_index {
+ self.sizes.push(Size::new());
+ }
+
+ let host_visible = self.host_visible();
+
+ let mut candidate_size_index = size_index;
+
+ let (mut entry, entry_size_index) = loop {
+ let sizes_len = self.sizes.len();
+
+ let candidate_size_entry = &mut self.sizes[candidate_size_index];
+ let candidate_size = self.minimal_size << candidate_size_index;
+
+ if let Some(entry) = candidate_size_entry.acquire(candidate_size) {
+ break (entry, candidate_size_index);
+ }
+
+ if sizes_len == candidate_size_index + 1 {
+ // That's size of device allocation.
+ if *allocations_remains == 0 {
+ return Err(AllocationError::TooManyObjects);
+ }
+
+ let chunk_size = self.minimal_size << (candidate_size_index + 1);
+ let mut memory = device.allocate_memory(chunk_size, self.memory_type, flags)?;
+ *allocations_remains -= 1;
+ heap.alloc(chunk_size);
+
+ let ptr = if host_visible {
+ match device.map_memory(&mut memory, 0, chunk_size) {
+ Ok(ptr) => Some(ptr),
+ Err(DeviceMapError::OutOfDeviceMemory) => {
+ return Err(AllocationError::OutOfDeviceMemory)
+ }
+ Err(DeviceMapError::MapFailed) | Err(DeviceMapError::OutOfHostMemory) => {
+ return Err(AllocationError::OutOfHostMemory)
+ }
+ }
+ } else {
+ None
+ };
+
+ let chunk = self.chunks.insert(Chunk {
+ memory: Arc::new(memory),
+ ptr,
+ size: chunk_size,
+ });
+
+ let entry = candidate_size_entry.add_pair_and_acquire_left(chunk, 0, None);
+
+ break (entry, candidate_size_index);
+ }
+
+ candidate_size_index += 1;
+ };
+
+ for size_index in (size_index..entry_size_index).rev() {
+ let size_entry = &mut self.sizes[size_index];
+ entry =
+ size_entry.add_pair_and_acquire_left(entry.chunk, entry.offset, Some(entry.index));
+ }
+
+ let chunk_entry = self.chunks.get_unchecked(entry.chunk);
+
+ debug_assert!(
+ entry
+ .offset
+ .checked_add(size)
+ .map_or(false, |end| end <= chunk_entry.size),
+ "Offset + size is not in chunk bounds"
+ );
+
+ Ok(BuddyBlock {
+ memory: chunk_entry.memory.clone(),
+ ptr: chunk_entry
+ .ptr
+ .map(|ptr| NonNull::new_unchecked(ptr.as_ptr().add(entry.offset as usize))),
+ offset: entry.offset,
+ size,
+ chunk: entry.chunk,
+ index: entry.index,
+ })
+ }
+
+ #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, device)))]
+ pub unsafe fn dealloc(
+ &mut self,
+ device: &impl MemoryDevice<M>,
+ block: BuddyBlock<M>,
+ heap: &mut Heap,
+ allocations_remains: &mut u32,
+ ) {
+ debug_assert!(block.size.is_power_of_two());
+
+ let size_index =
+ (block.size.trailing_zeros() - self.minimal_size.trailing_zeros()) as usize;
+
+ let mut release_index = block.index;
+ let mut release_size_index = size_index;
+
+ loop {
+ match self.sizes[release_size_index].release(release_index) {
+ Release::Parent(parent) => {
+ release_size_index += 1;
+ release_index = parent;
+ }
+ Release::Chunk(chunk) => {
+ debug_assert_eq!(chunk, block.chunk);
+ debug_assert_eq!(
+ self.chunks.get(chunk).size,
+ self.minimal_size << (release_size_index + 1)
+ );
+ let chunk = self.chunks.remove(chunk);
+ drop(block);
+
+ let memory = try_arc_unwrap(chunk.memory)
+ .expect("Memory shared after last block deallocated");
+
+ device.deallocate_memory(memory);
+ *allocations_remains += 1;
+ heap.dealloc(chunk.size);
+
+ return;
+ }
+ Release::None => return,
+ }
+ }
+ }
+
+ fn host_visible(&self) -> bool {
+ self.props.contains(MemoryPropertyFlags::HOST_VISIBLE)
+ }
+}
|