Vendor things

third-party/vendor/gpu-alloc/src/allocator.rs

@@ -0,0 +1,611 @@
+use {
+    crate::{
+        align_down,
+        block::{MemoryBlock, MemoryBlockFlavor},
+        buddy::{BuddyAllocator, BuddyBlock},
+        config::Config,
+        error::AllocationError,
+        freelist::{FreeListAllocator, FreeListBlock},
+        heap::Heap,
+        usage::{MemoryForUsage, UsageFlags},
+        MemoryBounds, Request,
+    },
+    alloc::boxed::Box,
+    core::convert::TryFrom as _,
+    gpu_alloc_types::{
+        AllocationFlags, DeviceProperties, MemoryDevice, MemoryPropertyFlags, MemoryType,
+        OutOfMemory,
+    },
+};
+
+/// Memory allocator for Vulkan-like APIs.
+#[derive(Debug)]
+pub struct GpuAllocator<M> {
+    dedicated_threshold: u64,
+    preferred_dedicated_threshold: u64,
+    transient_dedicated_threshold: u64,
+    max_memory_allocation_size: u64,
+    memory_for_usage: MemoryForUsage,
+    memory_types: Box<[MemoryType]>,
+    memory_heaps: Box<[Heap]>,
+    allocations_remains: u32,
+    non_coherent_atom_mask: u64,
+    starting_free_list_chunk: u64,
+    final_free_list_chunk: u64,
+    minimal_buddy_size: u64,
+    initial_buddy_dedicated_size: u64,
+    buffer_device_address: bool,
+
+    buddy_allocators: Box<[Option<BuddyAllocator<M>>]>,
+    freelist_allocators: Box<[Option<FreeListAllocator<M>>]>,
+}
+
+/// Hints for allocator to decide on allocation strategy.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+#[non_exhaustive]
+pub enum Dedicated {
+    /// Allocation directly from device.\
+    /// Very slow.
+    /// Count of allocations is limited.\
+    /// Use with caution.\
+    /// Must be used if resource has to be bound to dedicated memory object.
+    Required,
+
+    /// Hint for allocator that dedicated memory object is preferred.\
+    /// Should be used if it is known that resource placed in dedicated memory object
+    /// would allow for better performance.\
+    /// Implementation is allowed to return block to shared memory object.
+    Preferred,
+}
+
+impl<M> GpuAllocator<M>
+where
+    M: MemoryBounds + 'static,
+{
+    /// Creates  new instance of `GpuAllocator`.
+    /// Provided `DeviceProperties` should match properties of `MemoryDevice` that will be used
+    /// with created `GpuAllocator` instance.
+    #[cfg_attr(feature = "tracing", tracing::instrument)]
+    pub fn new(config: Config, props: DeviceProperties<'_>) -> Self {
+        assert!(
+            props.non_coherent_atom_size.is_power_of_two(),
+            "`non_coherent_atom_size` must be power of two"
+        );
+
+        assert!(
+            isize::try_from(props.non_coherent_atom_size).is_ok(),
+            "`non_coherent_atom_size` must fit host address space"
+        );
+
+        GpuAllocator {
+            dedicated_threshold: config.dedicated_threshold,
+            preferred_dedicated_threshold: config
+                .preferred_dedicated_threshold
+                .min(config.dedicated_threshold),
+
+            transient_dedicated_threshold: config
+                .transient_dedicated_threshold
+                .max(config.dedicated_threshold),
+
+            max_memory_allocation_size: props.max_memory_allocation_size,
+
+            memory_for_usage: MemoryForUsage::new(props.memory_types.as_ref()),
+
+            memory_types: props.memory_types.as_ref().iter().copied().collect(),
+            memory_heaps: props
+                .memory_heaps
+                .as_ref()
+                .iter()
+                .map(|heap| Heap::new(heap.size))
+                .collect(),
+
+            buffer_device_address: props.buffer_device_address,
+
+            allocations_remains: props.max_memory_allocation_count,
+            non_coherent_atom_mask: props.non_coherent_atom_size - 1,
+
+            starting_free_list_chunk: config.starting_free_list_chunk,
+            final_free_list_chunk: config.final_free_list_chunk,
+            minimal_buddy_size: config.minimal_buddy_size,
+            initial_buddy_dedicated_size: config.initial_buddy_dedicated_size,
+
+            buddy_allocators: props.memory_types.as_ref().iter().map(|_| None).collect(),
+            freelist_allocators: props.memory_types.as_ref().iter().map(|_| None).collect(),
+        }
+    }
+
+    /// Allocates memory block from specified `device` according to the `request`.
+    ///
+    /// # Safety
+    ///
+    /// * `device` must be one with `DeviceProperties` that were provided to create this `GpuAllocator` instance.
+    /// * Same `device` instance must be used for all interactions with one `GpuAllocator` instance
+    ///   and memory blocks allocated from it.
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, device)))]
+    pub unsafe fn alloc(
+        &mut self,
+        device: &impl MemoryDevice<M>,
+        request: Request,
+    ) -> Result<MemoryBlock<M>, AllocationError> {
+        self.alloc_internal(device, request, None)
+    }
+
+    /// Allocates memory block from specified `device` according to the `request`.
+    /// This function allows user to force specific allocation strategy.
+    /// Improper use can lead to suboptimal performance or too large overhead.
+    /// Prefer `GpuAllocator::alloc` if doubt.
+    ///
+    /// # Safety
+    ///
+    /// * `device` must be one with `DeviceProperties` that were provided to create this `GpuAllocator` instance.
+    /// * Same `device` instance must be used for all interactions with one `GpuAllocator` instance
+    ///   and memory blocks allocated from it.
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, device)))]
+    pub unsafe fn alloc_with_dedicated(
+        &mut self,
+        device: &impl MemoryDevice<M>,
+        request: Request,
+        dedicated: Dedicated,
+    ) -> Result<MemoryBlock<M>, AllocationError> {
+        self.alloc_internal(device, request, Some(dedicated))
+    }
+
+    unsafe fn alloc_internal(
+        &mut self,
+        device: &impl MemoryDevice<M>,
+        mut request: Request,
+        dedicated: Option<Dedicated>,
+    ) -> Result<MemoryBlock<M>, AllocationError> {
+        enum Strategy {
+            Buddy,
+            Dedicated,
+            FreeList,
+        }
+
+        request.usage = with_implicit_usage_flags(request.usage);
+
+        if request.usage.contains(UsageFlags::DEVICE_ADDRESS) {
+            assert!(self.buffer_device_address, "`DEVICE_ADDRESS` cannot be requested when `DeviceProperties::buffer_device_address` is false");
+        }
+
+        if request.size > self.max_memory_allocation_size {
+            return Err(AllocationError::OutOfDeviceMemory);
+        }
+
+        if let Some(Dedicated::Required) = dedicated {
+            if self.allocations_remains == 0 {
+                return Err(AllocationError::TooManyObjects);
+            }
+        }
+
+        if 0 == self.memory_for_usage.mask(request.usage) & request.memory_types {
+            #[cfg(feature = "tracing")]
+            tracing::error!(
+                "Cannot serve request {:?}, no memory among bitset `{}` support usage {:?}",
+                request,
+                request.memory_types,
+                request.usage
+            );
+
+            return Err(AllocationError::NoCompatibleMemoryTypes);
+        }
+
+        let transient = request.usage.contains(UsageFlags::TRANSIENT);
+
+        for &index in self.memory_for_usage.types(request.usage) {
+            if 0 == request.memory_types & (1 << index) {
+                // Skip memory type incompatible with the request.
+                continue;
+            }
+
+            let memory_type = &self.memory_types[index as usize];
+            let heap = memory_type.heap;
+            let heap = &mut self.memory_heaps[heap as usize];
+
+            if request.size > heap.size() {
+                // Impossible to use memory type from this heap.
+                continue;
+            }
+
+            let atom_mask = if host_visible_non_coherent(memory_type.props) {
+                self.non_coherent_atom_mask
+            } else {
+                0
+            };
+
+            let flags = if self.buffer_device_address {
+                AllocationFlags::DEVICE_ADDRESS
+            } else {
+                AllocationFlags::empty()
+            };
+
+            let strategy = match (dedicated, transient) {
+                (Some(Dedicated::Required), _) => Strategy::Dedicated,
+                (Some(Dedicated::Preferred), _)
+                    if request.size >= self.preferred_dedicated_threshold =>
+                {
+                    Strategy::Dedicated
+                }
+                (_, true) => {
+                    let threshold = self.transient_dedicated_threshold.min(heap.size() / 32);
+
+                    if request.size < threshold {
+                        Strategy::FreeList
+                    } else {
+                        Strategy::Dedicated
+                    }
+                }
+                (_, false) => {
+                    let threshold = self.dedicated_threshold.min(heap.size() / 32);
+
+                    if request.size < threshold {
+                        Strategy::Buddy
+                    } else {
+                        Strategy::Dedicated
+                    }
+                }
+            };
+
+            match strategy {
+                Strategy::Dedicated => {
+                    #[cfg(feature = "tracing")]
+                    tracing::debug!(
+                        "Allocating memory object `{}@{:?}`",
+                        request.size,
+                        memory_type
+                    );
+
+                    match device.allocate_memory(request.size, index, flags) {
+                        Ok(memory) => {
+                            self.allocations_remains -= 1;
+                            heap.alloc(request.size);
+
+                            return Ok(MemoryBlock::new(
+                                index,
+                                memory_type.props,
+                                0,
+                                request.size,
+                                atom_mask,
+                                MemoryBlockFlavor::Dedicated { memory },
+                            ));
+                        }
+                        Err(OutOfMemory::OutOfDeviceMemory) => continue,
+                        Err(OutOfMemory::OutOfHostMemory) => {
+                            return Err(AllocationError::OutOfHostMemory)
+                        }
+                    }
+                }
+                Strategy::FreeList => {
+                    let allocator = match &mut self.freelist_allocators[index as usize] {
+                        Some(allocator) => allocator,
+                        slot => {
+                            let starting_free_list_chunk = match align_down(
+                                self.starting_free_list_chunk.min(heap.size() / 32),
+                                atom_mask,
+                            ) {
+                                0 => atom_mask,
+                                other => other,
+                            };
+
+                            let final_free_list_chunk = match align_down(
+                                self.final_free_list_chunk
+                                    .max(self.starting_free_list_chunk)
+                                    .max(self.transient_dedicated_threshold)
+                                    .min(heap.size() / 32),
+                                atom_mask,
+                            ) {
+                                0 => atom_mask,
+                                other => other,
+                            };
+
+                            slot.get_or_insert(FreeListAllocator::new(
+                                starting_free_list_chunk,
+                                final_free_list_chunk,
+                                index,
+                                memory_type.props,
+                                if host_visible_non_coherent(memory_type.props) {
+                                    self.non_coherent_atom_mask
+                                } else {
+                                    0
+                                },
+                            ))
+                        }
+                    };
+                    let result = allocator.alloc(
+                        device,
+                        request.size,
+                        request.align_mask,
+                        flags,
+                        heap,
+                        &mut self.allocations_remains,
+                    );
+
+                    match result {
+                        Ok(block) => {
+                            return Ok(MemoryBlock::new(
+                                index,
+                                memory_type.props,
+                                block.offset,
+                                block.size,
+                                atom_mask,
+                                MemoryBlockFlavor::FreeList {
+                                    chunk: block.chunk,
+                                    ptr: block.ptr,
+                                    memory: block.memory,
+                                },
+                            ))
+                        }
+                        Err(AllocationError::OutOfDeviceMemory) => continue,
+                        Err(err) => return Err(err),
+                    }
+                }
+
+                Strategy::Buddy => {
+                    let allocator = match &mut self.buddy_allocators[index as usize] {
+                        Some(allocator) => allocator,
+                        slot => {
+                            let minimal_buddy_size = self
+                                .minimal_buddy_size
+                                .min(heap.size() / 1024)
+                                .next_power_of_two();
+
+                            let initial_buddy_dedicated_size = self
+                                .initial_buddy_dedicated_size
+                                .min(heap.size() / 32)
+                                .next_power_of_two();
+
+                            slot.get_or_insert(BuddyAllocator::new(
+                                minimal_buddy_size,
+                                initial_buddy_dedicated_size,
+                                index,
+                                memory_type.props,
+                                if host_visible_non_coherent(memory_type.props) {
+                                    self.non_coherent_atom_mask
+                                } else {
+                                    0
+                                },
+                            ))
+                        }
+                    };
+                    let result = allocator.alloc(
+                        device,
+                        request.size,
+                        request.align_mask,
+                        flags,
+                        heap,
+                        &mut self.allocations_remains,
+                    );
+
+                    match result {
+                        Ok(block) => {
+                            return Ok(MemoryBlock::new(
+                                index,
+                                memory_type.props,
+                                block.offset,
+                                block.size,
+                                atom_mask,
+                                MemoryBlockFlavor::Buddy {
+                                    chunk: block.chunk,
+                                    ptr: block.ptr,
+                                    index: block.index,
+                                    memory: block.memory,
+                                },
+                            ))
+                        }
+                        Err(AllocationError::OutOfDeviceMemory) => continue,
+                        Err(err) => return Err(err),
+                    }
+                }
+            }
+        }
+
+        Err(AllocationError::OutOfDeviceMemory)
+    }
+
+    /// Creates a memory block from an existing memory allocation, transferring ownership to the allocator.
+    ///
+    /// This function allows the [`GpuAllocator`] to manage memory allocated outside of the typical
+    /// [`GpuAllocator::alloc`] family of functions.
+    ///
+    /// # Usage
+    ///
+    /// If you need to import external memory, such as a Win32 `HANDLE` or a Linux `dmabuf`, import the device
+    /// memory using the graphics api and platform dependent functions. Once that is done, call this function
+    /// to make the [`GpuAllocator`] take ownership of the imported memory.
+    ///
+    /// When calling this function, you **must** ensure there are [enough remaining allocations](GpuAllocator::remaining_allocations).
+    ///
+    /// # Safety
+    ///
+    /// - The `memory` must be allocated with the same device that was provided to create this [`GpuAllocator`]
+    ///   instance.
+    /// - The `memory` must be valid.
+    /// - The `props`, `offset` and `size` must match the properties, offset and size of the memory allocation.
+    /// - The memory must have been allocated with the specified `memory_type`.
+    /// - There must be enough remaining allocations.
+    /// - The memory allocation must not come from an existing memory block created by this allocator.
+    /// - The underlying memory object must be deallocated using the returned [`MemoryBlock`] with
+    ///   [`GpuAllocator::dealloc`].
+    pub unsafe fn import_memory(
+        &mut self,
+        memory: M,
+        memory_type: u32,
+        props: MemoryPropertyFlags,
+        offset: u64,
+        size: u64,
+    ) -> MemoryBlock<M> {
+        // Get the heap which the imported memory is from.
+        let heap = self
+            .memory_types
+            .get(memory_type as usize)
+            .expect("Invalid memory type specified when importing memory")
+            .heap;
+        let heap = &mut self.memory_heaps[heap as usize];
+
+        #[cfg(feature = "tracing")]
+        tracing::debug!(
+            "Importing memory object {:?} `{}@{:?}`",
+            memory,
+            size,
+            memory_type
+        );
+
+        assert_ne!(
+            self.allocations_remains, 0,
+            "Out of allocations when importing a memory block. Ensure you check GpuAllocator::remaining_allocations before import."
+        );
+        self.allocations_remains -= 1;
+
+        let atom_mask = if host_visible_non_coherent(props) {
+            self.non_coherent_atom_mask
+        } else {
+            0
+        };
+
+        heap.alloc(size);
+
+        MemoryBlock::new(
+            memory_type,
+            props,
+            offset,
+            size,
+            atom_mask,
+            MemoryBlockFlavor::Dedicated { memory },
+        )
+    }
+
+    /// Deallocates memory block previously allocated from this `GpuAllocator` instance.
+    ///
+    /// # Safety
+    ///
+    /// * Memory block must have been allocated by this `GpuAllocator` instance
+    /// * `device` must be one with `DeviceProperties` that were provided to create this `GpuAllocator` instance
+    /// * Same `device` instance must be used for all interactions with one `GpuAllocator` instance
+    ///   and memory blocks allocated from it
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, device)))]
+    pub unsafe fn dealloc(&mut self, device: &impl MemoryDevice<M>, block: MemoryBlock<M>) {
+        let memory_type = block.memory_type();
+        let offset = block.offset();
+        let size = block.size();
+        let flavor = block.deallocate();
+        match flavor {
+            MemoryBlockFlavor::Dedicated { memory } => {
+                let heap = self.memory_types[memory_type as usize].heap;
+                device.deallocate_memory(memory);
+                self.allocations_remains += 1;
+                self.memory_heaps[heap as usize].dealloc(size);
+            }
+            MemoryBlockFlavor::Buddy {
+                chunk,
+                ptr,
+                index,
+                memory,
+            } => {
+                let heap = self.memory_types[memory_type as usize].heap;
+                let heap = &mut self.memory_heaps[heap as usize];
+
+                let allocator = self.buddy_allocators[memory_type as usize]
+                    .as_mut()
+                    .expect("Allocator should exist");
+
+                allocator.dealloc(
+                    device,
+                    BuddyBlock {
+                        memory,
+                        ptr,
+                        offset,
+                        size,
+                        chunk,
+                        index,
+                    },
+                    heap,
+                    &mut self.allocations_remains,
+                );
+            }
+            MemoryBlockFlavor::FreeList { chunk, ptr, memory } => {
+                let heap = self.memory_types[memory_type as usize].heap;
+                let heap = &mut self.memory_heaps[heap as usize];
+
+                let allocator = self.freelist_allocators[memory_type as usize]
+                    .as_mut()
+                    .expect("Allocator should exist");
+
+                allocator.dealloc(
+                    device,
+                    FreeListBlock {
+                        memory,
+                        ptr,
+                        chunk,
+                        offset,
+                        size,
+                    },
+                    heap,
+                    &mut self.allocations_remains,
+                );
+            }
+        }
+    }
+
+    /// Returns the maximum allocation size supported.
+    pub fn max_allocation_size(&self) -> u64 {
+        self.max_memory_allocation_size
+    }
+
+    /// Returns the number of remaining available allocations.
+    ///
+    /// This may be useful if you need know if the allocator can allocate a number of allocations ahead of
+    /// time. This function is also useful for ensuring you do not allocate too much memory outside allocator
+    /// (such as external memory).
+    pub fn remaining_allocations(&self) -> u32 {
+        self.allocations_remains
+    }
+
+    /// Sets the number of remaining available allocations.
+    ///
+    /// # Safety
+    ///
+    /// The caller is responsible for ensuring the number of remaining allocations does not exceed how many
+    /// remaining allocations there actually are on the memory device.
+    pub unsafe fn set_remaining_allocations(&mut self, remaining: u32) {
+        self.allocations_remains = remaining;
+    }
+
+    /// Deallocates leftover memory objects.
+    /// Should be used before dropping.
+    ///
+    /// # Safety
+    ///
+    /// * `device` must be one with `DeviceProperties` that were provided to create this `GpuAllocator` instance
+    /// * Same `device` instance must be used for all interactions with one `GpuAllocator` instance
+    ///   and memory blocks allocated from it
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, device)))]
+    pub unsafe fn cleanup(&mut self, device: &impl MemoryDevice<M>) {
+        for (index, allocator) in self
+            .freelist_allocators
+            .iter_mut()
+            .enumerate()
+            .filter_map(|(index, allocator)| Some((index, allocator.as_mut()?)))
+        {
+            let memory_type = &self.memory_types[index];
+            let heap = memory_type.heap;
+            let heap = &mut self.memory_heaps[heap as usize];
+
+            allocator.cleanup(device, heap, &mut self.allocations_remains);
+        }
+    }
+}
+
+fn host_visible_non_coherent(props: MemoryPropertyFlags) -> bool {
+    (props & (MemoryPropertyFlags::HOST_COHERENT | MemoryPropertyFlags::HOST_VISIBLE))
+        == MemoryPropertyFlags::HOST_VISIBLE
+}
+
+fn with_implicit_usage_flags(usage: UsageFlags) -> UsageFlags {
+    if usage.is_empty() {
+        UsageFlags::FAST_DEVICE_ACCESS
+    } else if usage.intersects(UsageFlags::DOWNLOAD | UsageFlags::UPLOAD) {
+        usage | UsageFlags::HOST_ACCESS
+    } else {
+        usage
+    }
+}