Adding upstream version 110.0.1.upstream/110.0.1upstream

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

1 files changed, 458 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..5691f4d96d
--- /dev/null
+++ b/third_party/rust/gpu-alloc/src/buddy.rs
@@ -0,0 +1,458 @@
+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_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)
+    }
+}