alex/volcanic-mem

root / src / allocator.rs

use crate::device::DeviceFns;
use crate::error::{Error, Result, check};
use crate::tlsf::Tlsf;
use core::ffi::c_void;
use core::ptr::{NonNull, null, null_mut};
use std::sync::{Mutex, MutexGuard, PoisonError};
use volcanic::loader::{CmdsDevice, CmdsInstance};
use volcanic::*;

const DEDICATED: u32 = u32::MAX;
const SMALL_HEAP: u64 = 1 << 30;
const BLOCK_ALIGNMENT: u64 = 1 << 12;
const NEW_BLOCK_SHIFTS: u32 = 3;
const UNUSABLE: VkMemoryPropertyFlags = VkMemoryPropertyFlags::from_bits_retain(
    VkMemoryPropertyFlags::LAZILY_ALLOCATED.bits() | (1 << 5) | (1 << 6) | (1 << 7),
);

#[inline]
const fn align_up(value: u64, align: u64) -> u64 {
    (value + align - 1) & !(align - 1)
}

// ---- Public types -------------------------------------------

/// Where an allocation should live, from the point of view of who accesses it.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub enum MemoryLocation {
    /// Device-local memory, not host visible when avoidable.
    #[default]
    GpuOnly,
    /// Host-visible memory written sequentially by the CPU. Resources the GPU reads directly
    /// go to device-local host-visible memory when available, while transfer-only (staging)
    /// buffers go to system memory.
    CpuToGpu,
    /// Host-visible, preferably cached memory for reading GPU results back on the CPU.
    GpuToCpu,
}

/// Distinguishes resources for `bufferImageGranularity`: buffers and linear images are
/// `Linear`, optimal-tiling images are `Optimal`.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub enum ResourceKind {
    #[default]
    Linear,
    Optimal,
}

#[derive(Clone, Copy)]
pub struct AllocationDesc {
    pub requirements: VkMemoryRequirements,
    pub location: MemoryLocation,
    pub kind: ResourceKind,
    /// Gives the allocation its own `VkDeviceMemory`.
    pub dedicated: bool,
}

#[derive(Clone, Copy, Debug)]
pub struct AllocatorConfig {
    /// Block size for heaps larger than 1 GiB. Smaller heaps use 1/8 of their size.
    pub preferred_block_size: u64,
    /// Set when `VK_EXT_memory_budget` is enabled, so [`Allocator::budgets`] reports driver numbers.
    pub memory_budget: bool,
}

impl Default for AllocatorConfig {
    fn default() -> Self {
        Self {
            preferred_block_size: 256 << 20,
            memory_budget: false,
        }
    }
}

/// A sub-range of a `VkDeviceMemory`. Must be returned with [`Allocator::free`].
#[derive(Debug)]
pub struct Allocation {
    memory: VkDeviceMemory,
    offset: u64,
    size: u64,
    mapped: *mut u8,
    memory_type: u32,
    block: u32,
    node: u32,
}

unsafe impl Send for Allocation {}
unsafe impl Sync for Allocation {}

impl Allocation {
    #[inline]
    pub fn memory(&self) -> VkDeviceMemory {
        self.memory
    }

    #[inline]
    pub fn offset(&self) -> u64 {
        self.offset
    }

    #[inline]
    pub fn size(&self) -> u64 {
        self.size
    }

    #[inline]
    pub fn memory_type(&self) -> u32 {
        self.memory_type
    }

    #[inline]
    pub fn is_dedicated(&self) -> bool {
        self.block == DEDICATED
    }

    /// Host pointer to the start of the allocation, if its memory is host visible.
    #[inline]
    pub fn mapped_ptr(&self) -> Option<NonNull<u8>> {
        NonNull::new(self.mapped)
    }

    #[inline]
    pub fn mapped_slice(&self) -> Option<&[u8]> {
        self.mapped_ptr()
            .map(|ptr| unsafe { core::slice::from_raw_parts(ptr.as_ptr(), self.size as usize) })
    }

    #[inline]
    pub fn mapped_slice_mut(&mut self) -> Option<&mut [u8]> {
        self.mapped_ptr()
            .map(|ptr| unsafe { core::slice::from_raw_parts_mut(ptr.as_ptr(), self.size as usize) })
    }
}

/// Usage counters of one memory heap.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Statistics {
    /// `VkDeviceMemory` objects, including dedicated allocations.
    pub block_count: u32,
    pub block_bytes: u64,
    pub allocation_count: u32,
    pub allocation_bytes: u64,
}

impl Statistics {
    fn add(&mut self, other: &Self) {
        self.block_count += other.block_count;
        self.block_bytes += other.block_bytes;
        self.allocation_count += other.allocation_count;
        self.allocation_bytes += other.allocation_bytes;
    }
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct HeapBudget {
    pub statistics: Statistics,
    /// Bytes the whole process uses in this heap; estimated from `block_bytes` without
    /// `VK_EXT_memory_budget`.
    pub usage: u64,
    /// Bytes the process can use before running into trouble; estimated as 80% of the heap
    /// size without `VK_EXT_memory_budget`.
    pub budget: u64,
}

// ---- Internals ----------------------------------------------

#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) enum Dedicated {
    No,
    Prefer,
    Require,
}

#[derive(Clone, Copy)]
pub(crate) enum Target {
    None,
    Buffer(VkBuffer),
    Image(VkImage),
}

pub(crate) struct Request {
    pub size: u64,
    pub alignment: u64,
    pub type_bits: u32,
    pub location: MemoryLocation,
    pub kind: ResourceKind,
    pub device_access: bool,
    pub dedicated: Dedicated,
    pub target: Target,
}

struct Block {
    memory: VkDeviceMemory,
    mapped: *mut u8,
    tlsf: Tlsf,
}

struct Pool {
    blocks: Vec<Option<Block>>,
    /// Live block slots, kept roughly sorted by ascending free space.
    order: Vec<u32>,
    vacant: Vec<u32>,
    statistics: Statistics,
}

unsafe impl Send for Pool {}

impl Pool {
    /// One bubble step towards ascending free space, keeping allocations packed into the fullest blocks.
    #[inline]
    fn sort_step(&mut self) {
        for i in 1..self.order.len() {
            let free = |slot: u32| {
                self.blocks[slot as usize]
                    .as_ref()
                    .map_or(0, |b| b.tlsf.free_bytes())
            };
            if free(self.order[i - 1]) > free(self.order[i]) {
                self.order.swap(i - 1, i);
                return;
            }
        }
    }
}

#[derive(Clone, Copy)]
struct MemoryType {
    flags: VkMemoryPropertyFlags,
    heap: u32,
}

/// Thread-safe TLSF-based Vulkan memory allocator.
pub struct Allocator {
    device: VkDevice,
    fns: DeviceFns,
    physical_device: VkPhysicalDevice,
    get_memory_properties2: Option<PFN_vkGetPhysicalDeviceMemoryProperties2>,
    types: Box<[MemoryType]>,
    heaps: Box<[VkMemoryHeap]>,
    pools: Box<[Mutex<Pool>]>,
    block_sizes: Box<[u64]>,
    /// Memory types ordered by preference, indexed by `location * 2 + device_access`.
    preferences: [Box<[u8]>; 6],
    granularity: u64,
    atom: u64,
    pub(crate) heap_limits: crate::resource::HeapLimits,
}

unsafe impl Send for Allocator {}
unsafe impl Sync for Allocator {}

impl Allocator {
    /// # Safety
    /// The command tables must belong to `physical_device` and `device`. The device must be a
    /// Vulkan 1.4 device with `VK_EXT_descriptor_heap` and the `bufferDeviceAddress`,
    /// `maintenance5` and `descriptorHeap` features enabled, and it must outlive the allocator.
    pub unsafe fn new(
        instance: &CmdsInstance,
        physical_device: VkPhysicalDevice,
        device_cmds: &CmdsDevice,
        device: VkDevice,
        config: &AllocatorConfig,
    ) -> Result<Self> {
        let fns = DeviceFns::load(device_cmds)?;
        let get_properties2 = instance
            .vkGetPhysicalDeviceProperties2
            .ok_or(Error::MissingCommand("vkGetPhysicalDeviceProperties2"))?;
        let get_memory_properties2 =
            instance
                .vkGetPhysicalDeviceMemoryProperties2
                .ok_or(Error::MissingCommand(
                    "vkGetPhysicalDeviceMemoryProperties2",
                ))?;

        let mut heap_properties: VkPhysicalDeviceDescriptorHeapPropertiesEXT =
            unsafe { core::mem::zeroed() };
        heap_properties.sType = VkStructureType::PhysicalDeviceDescriptorHeapPropertiesExt;
        let mut properties: VkPhysicalDeviceProperties2 = unsafe { core::mem::zeroed() };
        properties.sType = VkStructureType::PhysicalDeviceProperties2;
        properties.pNext = (&raw mut heap_properties).cast();
        unsafe { get_properties2(physical_device, &mut properties) };
        let limits = &properties.properties.limits;

        let mut memory: VkPhysicalDeviceMemoryProperties2 = unsafe { core::mem::zeroed() };
        memory.sType = VkStructureType::PhysicalDeviceMemoryProperties2;
        unsafe { get_memory_properties2(physical_device, &mut memory) };
        let memory = &memory.memoryProperties;

        let heaps: Box<[_]> = memory.memoryHeaps[..memory.memoryHeapCount as usize].into();
        let types: Box<[_]> = memory.memoryTypes[..memory.memoryTypeCount as usize]
            .iter()
            .map(|t| MemoryType {
                flags: t.propertyFlags,
                heap: t.heapIndex,
            })
            .collect();
        let block_sizes = types
            .iter()
            .map(|t| {
                let heap_size = heaps[t.heap as usize].size;
                let size = if heap_size <= SMALL_HEAP {
                    heap_size / 8
                } else {
                    config.preferred_block_size
                };
                align_up(size.max(BLOCK_ALIGNMENT), BLOCK_ALIGNMENT)
            })
            .collect();
        let pools = types
            .iter()
            .map(|_| {
                Mutex::new(Pool {
                    blocks: Vec::new(),
                    order: Vec::new(),
                    vacant: Vec::new(),
                    statistics: Statistics::default(),
                })
            })
            .collect();
        let preferences = core::array::from_fn(|i| {
            let location = [
                MemoryLocation::GpuOnly,
                MemoryLocation::CpuToGpu,
                MemoryLocation::GpuToCpu,
            ][i / 2];
            preference_order(&types, location, i % 2 == 1)
        });

        Ok(Self {
            device,
            fns,
            physical_device,
            get_memory_properties2: config.memory_budget.then_some(get_memory_properties2),
            types,
            heaps,
            pools,
            block_sizes,
            preferences,
            granularity: limits.bufferImageGranularity.max(1),
            atom: limits.nonCoherentAtomSize.max(1),
            heap_limits: crate::resource::HeapLimits::new(&heap_properties),
        })
    }

    #[inline]
    pub fn device(&self) -> VkDevice {
        self.device
    }

    #[inline]
    pub(crate) fn fns(&self) -> &DeviceFns {
        &self.fns
    }

    /// Allocates memory for a resource the caller creates and binds itself.
    pub fn allocate(&self, desc: &AllocationDesc) -> Result<Allocation> {
        self.allocate_request(&Request {
            size: desc.requirements.size,
            alignment: desc.requirements.alignment,
            type_bits: desc.requirements.memoryTypeBits,
            location: desc.location,
            kind: desc.kind,
            device_access: true,
            dedicated: if desc.dedicated {
                Dedicated::Require
            } else {
                Dedicated::No
            },
            target: Target::None,
        })
    }

    pub fn free(&self, allocation: Allocation) {
        let size = allocation.size;
        if allocation.block == DEDICATED {
            unsafe { (self.fns.free_memory)(self.device, allocation.memory, null()) };
            let stats = &mut self.lock(allocation.memory_type).statistics;
            stats.block_count -= 1;
            stats.block_bytes -= size;
            stats.allocation_count -= 1;
            stats.allocation_bytes -= size;
            return;
        }
        let released = {
            let mut guard = self.lock(allocation.memory_type);
            let pool = &mut *guard;
            let slot = allocation.block as usize;
            let tlsf = &mut pool.blocks[slot]
                .as_mut()
                .expect("allocation from a freed block")
                .tlsf;
            tlsf.free(allocation.node);
            pool.statistics.allocation_count -= 1;
            pool.statistics.allocation_bytes -= size;

            let empty = tlsf.is_empty();
            let keep = !empty
                || pool
                    .order
                    .iter()
                    .filter(|&&s| {
                        pool.blocks[s as usize]
                            .as_ref()
                            .is_some_and(|b| b.tlsf.is_empty())
                    })
                    .count()
                    == 1;
            if keep {
                pool.sort_step();
                None
            } else {
                let block = pool.blocks[slot].take().expect("live block");
                pool.order.retain(|&s| s as usize != slot);
                pool.vacant.push(slot as u32);
                pool.statistics.block_count -= 1;
                pool.statistics.block_bytes -= block.tlsf.size();
                Some(block.memory)
            }
        };
        if let Some(memory) = released {
            unsafe { (self.fns.free_memory)(self.device, memory, null()) };
        }
    }

    /// Makes host writes in `offset..offset + size` of the allocation visible to the device.
    /// Pass `VK_WHOLE_SIZE` to cover the rest of the allocation. No-op on coherent memory.
    pub fn flush(&self, allocation: &Allocation, offset: u64, size: u64) -> Result<()> {
        match self.mapped_range(allocation, offset, size) {
            Some(range) => check("vkFlushMappedMemoryRanges", unsafe {
                (self.fns.flush_mapped_memory_ranges)(self.device, 1, &range)
            }),
            None => Ok(()),
        }
    }

    /// Makes device writes in `offset..offset + size` of the allocation visible to the host.
    /// Pass `VK_WHOLE_SIZE` to cover the rest of the allocation. No-op on coherent memory.
    pub fn invalidate(&self, allocation: &Allocation, offset: u64, size: u64) -> Result<()> {
        match self.mapped_range(allocation, offset, size) {
            Some(range) => check("vkInvalidateMappedMemoryRanges", unsafe {
                (self.fns.invalidate_mapped_memory_ranges)(self.device, 1, &range)
            }),
            None => Ok(()),
        }
    }

    /// Per-heap usage, indexed like `VkPhysicalDeviceMemoryProperties::memoryHeaps`.
    pub fn budgets(&self) -> Vec<HeapBudget> {
        let mut budgets: Vec<HeapBudget> = self
            .heaps
            .iter()
            .map(|heap| HeapBudget {
                budget: heap.size / 10 * 8,
                ..Default::default()
            })
            .collect();
        for (pool, ty) in self.pools.iter().zip(&self.types) {
            let stats = pool
                .lock()
                .unwrap_or_else(PoisonError::into_inner)
                .statistics;
            budgets[ty.heap as usize].statistics.add(&stats);
        }
        for budget in &mut budgets {
            budget.usage = budget.statistics.block_bytes;
        }
        if let Some(driver) = self.driver_budget() {
            for (i, budget) in budgets.iter_mut().enumerate() {
                budget.usage = driver.heapUsage[i];
                budget.budget = driver.heapBudget[i];
            }
        }
        budgets
    }

    /// Statistics summed over all heaps.
    pub fn statistics(&self) -> Statistics {
        let mut total = Statistics::default();
        for pool in &self.pools {
            total.add(
                &pool
                    .lock()
                    .unwrap_or_else(PoisonError::into_inner)
                    .statistics,
            );
        }
        total
    }

    // ---- Allocation paths ---------------------------------------

    pub(crate) fn allocate_request(&self, request: &Request) -> Result<Allocation> {
        if request.size == 0 {
            return Err(Error::ZeroSize);
        }
        if !request.alignment.is_power_of_two() {
            return Err(Error::InvalidAlignment(request.alignment));
        }
        let order =
            &self.preferences[request.location as usize * 2 + request.device_access as usize];
        let mut candidates = order
            .iter()
            .map(|&ty| ty as u32)
            .filter(|ty| request.type_bits & (1 << ty) != 0)
            .peekable();
        if candidates.peek().is_none() {
            return Err(Error::NoMemoryType);
        }
        while let Some(ty) = candidates.next() {
            // Over-budget memory types are skipped while a fallback remains.
            let within_budget = candidates.peek().is_some();
            match self.allocate_from_type(ty, request, within_budget) {
                Err(Error::OutOfMemory) => continue,
                result => return result,
            }
        }
        Err(Error::OutOfMemory)
    }

    fn allocate_from_type(
        &self,
        ty: u32,
        request: &Request,
        within_budget: bool,
    ) -> Result<Allocation> {
        let block_size = self.block_sizes[ty as usize];
        let dedicated = match request.dedicated {
            Dedicated::Require => return self.allocate_dedicated(ty, request, within_budget),
            Dedicated::Prefer => true,
            Dedicated::No => request.size > block_size / 2,
        };
        if dedicated {
            match self.allocate_dedicated(ty, request, within_budget) {
                Err(Error::OutOfMemory) if request.size <= block_size => {}
                result => return result,
            }
        }
        self.allocate_in_block(ty, request, within_budget)
    }

    fn allocate_dedicated(
        &self,
        ty: u32,
        request: &Request,
        within_budget: bool,
    ) -> Result<Allocation> {
        if within_budget && !self.fits_budget(ty, request.size) {
            return Err(Error::OutOfMemory);
        }
        let (memory, mapped) = self.allocate_memory(ty, request.size, request.target)?;
        let stats = &mut self.lock(ty).statistics;
        stats.block_count += 1;
        stats.block_bytes += request.size;
        stats.allocation_count += 1;
        stats.allocation_bytes += request.size;
        Ok(Allocation {
            memory,
            offset: 0,
            size: request.size,
            mapped,
            memory_type: ty,
            block: DEDICATED,
            node: 0,
        })
    }

    fn allocate_in_block(
        &self,
        ty: u32,
        request: &Request,
        within_budget: bool,
    ) -> Result<Allocation> {
        let (size, alignment, kind) = (request.size, request.alignment, request.kind as u32);
        if size > self.block_sizes[ty as usize] {
            return Err(Error::OutOfMemory);
        }
        let mut guard = self.lock(ty);
        let pool = &mut *guard;
        // Close fits in any block beat splitting a large hole in a fuller one.
        for tight in [true, false] {
            for i in 0..pool.order.len() {
                let slot = pool.order[i];
                let block = pool.blocks[slot as usize]
                    .as_mut()
                    .expect("ordered block is live");
                let hit = if tight {
                    block.tlsf.allocate_tight(size, alignment, kind)
                } else {
                    block.tlsf.allocate_loose(size, alignment, kind)
                };
                if let Some((node, offset)) = hit {
                    let allocation = suballocation(block, ty, slot, node, offset, size);
                    pool.statistics.allocation_count += 1;
                    pool.statistics.allocation_bytes += size;
                    pool.sort_step();
                    return Ok(allocation);
                }
            }
        }

        let slot = self.create_block(ty, pool, size, within_budget)?;
        let block = pool.blocks[slot as usize]
            .as_mut()
            .expect("new block is live");
        let (node, offset) = block
            .tlsf
            .allocate(size, alignment, kind)
            .expect("an empty block larger than the request always fits");
        let allocation = suballocation(block, ty, slot, node, offset, size);
        pool.statistics.allocation_count += 1;
        pool.statistics.allocation_bytes += size;
        Ok(allocation)
    }

    /// Starts small (1/8 of the preferred size) while the pool is young, and halves on failure.
    fn create_block(
        &self,
        ty: u32,
        pool: &mut Pool,
        min_size: u64,
        within_budget: bool,
    ) -> Result<u32> {
        let largest = pool
            .blocks
            .iter()
            .flatten()
            .map(|b| b.tlsf.size())
            .max()
            .unwrap_or(0);
        let mut size = self.block_sizes[ty as usize];
        let mut shifts = 0;
        while shifts < NEW_BLOCK_SHIFTS && size / 2 > largest && size / 2 >= min_size * 2 {
            size /= 2;
            shifts += 1;
        }
        if within_budget && self.get_memory_properties2.is_some() {
            let available = self.available(ty);
            while shifts < NEW_BLOCK_SHIFTS && size > available && size / 2 >= min_size {
                size /= 2;
                shifts += 1;
            }
            if size > available {
                return Err(Error::OutOfMemory);
            }
        }
        let (memory, mapped) = loop {
            match self.allocate_memory(ty, size, Target::None) {
                Err(Error::OutOfMemory) if shifts < NEW_BLOCK_SHIFTS && size / 2 >= min_size => {
                    size /= 2;
                    shifts += 1;
                }
                result => break result?,
            }
        };
        let block = Block {
            memory,
            mapped,
            tlsf: Tlsf::new(size, self.granularity),
        };
        let slot = match pool.vacant.pop() {
            Some(slot) => {
                pool.blocks[slot as usize] = Some(block);
                slot
            }
            None => {
                pool.blocks.push(Some(block));
                (pool.blocks.len() - 1) as u32
            }
        };
        pool.order.push(slot);
        pool.statistics.block_count += 1;
        pool.statistics.block_bytes += size;
        Ok(slot)
    }

    fn allocate_memory(
        &self,
        ty: u32,
        size: u64,
        target: Target,
    ) -> Result<(VkDeviceMemory, *mut u8)> {
        let flags_info = VkMemoryAllocateFlagsInfo {
            sType: VkStructureType::MemoryAllocateFlagsInfo,
            pNext: null(),
            flags: VkMemoryAllocateFlags::DEVICE_ADDRESS,
            deviceMask: 0,
        };
        let (buffer, image) = match target {
            Target::None => (0, 0),
            Target::Buffer(buffer) => (buffer, 0),
            Target::Image(image) => (0, image),
        };
        let dedicated_info = VkMemoryDedicatedAllocateInfo {
            sType: VkStructureType::MemoryDedicatedAllocateInfo,
            pNext: (&raw const flags_info).cast(),
            image,
            buffer,
        };
        let info = VkMemoryAllocateInfo {
            sType: VkStructureType::MemoryAllocateInfo,
            pNext: if matches!(target, Target::None) {
                (&raw const flags_info).cast()
            } else {
                (&raw const dedicated_info).cast()
            },
            allocationSize: size,
            memoryTypeIndex: ty,
        };
        let mut memory = 0;
        check("vkAllocateMemory", unsafe {
            (self.fns.allocate_memory)(self.device, &info, null(), &mut memory)
        })?;

        let mut mapped: *mut c_void = null_mut();
        if self.types[ty as usize]
            .flags
            .contains(VkMemoryPropertyFlags::HOST_VISIBLE)
        {
            let result = unsafe {
                (self.fns.map_memory)(
                    self.device,
                    memory,
                    0,
                    VK_WHOLE_SIZE,
                    VkMemoryMapFlags::empty(),
                    &mut mapped,
                )
            };
            if let Err(error) = check("vkMapMemory", result) {
                unsafe { (self.fns.free_memory)(self.device, memory, null()) };
                return Err(error);
            }
        }
        Ok((memory, mapped.cast()))
    }

    #[inline]
    fn fits_budget(&self, ty: u32, size: u64) -> bool {
        self.get_memory_properties2.is_none() || size <= self.available(ty)
    }

    /// Driver-reported headroom of the heap backing memory type `ty`.
    fn available(&self, ty: u32) -> u64 {
        let heap = self.types[ty as usize].heap as usize;
        self.driver_budget().map_or(u64::MAX, |driver| {
            driver.heapBudget[heap].saturating_sub(driver.heapUsage[heap])
        })
    }

    fn driver_budget(&self) -> Option<VkPhysicalDeviceMemoryBudgetPropertiesEXT> {
        let get_memory_properties2 = self.get_memory_properties2?;
        let mut driver: VkPhysicalDeviceMemoryBudgetPropertiesEXT = unsafe { core::mem::zeroed() };
        driver.sType = VkStructureType::PhysicalDeviceMemoryBudgetPropertiesExt;
        let mut memory: VkPhysicalDeviceMemoryProperties2 = unsafe { core::mem::zeroed() };
        memory.sType = VkStructureType::PhysicalDeviceMemoryProperties2;
        memory.pNext = (&raw mut driver).cast();
        unsafe { get_memory_properties2(self.physical_device, &mut memory) };
        Some(driver)
    }

    fn mapped_range(
        &self,
        allocation: &Allocation,
        offset: u64,
        size: u64,
    ) -> Option<VkMappedMemoryRange> {
        let flags = self.types[allocation.memory_type as usize].flags;
        if !flags.contains(VkMemoryPropertyFlags::HOST_VISIBLE)
            || flags.contains(VkMemoryPropertyFlags::HOST_COHERENT)
        {
            return None;
        }
        debug_assert!(offset <= allocation.size);
        let size = if size == VK_WHOLE_SIZE {
            allocation.size - offset
        } else {
            size
        };
        let start = (allocation.offset + offset) & !(self.atom - 1);
        let end = align_up(allocation.offset + offset + size, self.atom);
        // Block sizes are multiples of the atom size; dedicated allocations may not be.
        let size = if allocation.is_dedicated() && end > allocation.size {
            VK_WHOLE_SIZE
        } else {
            end - start
        };
        Some(VkMappedMemoryRange {
            sType: VkStructureType::MappedMemoryRange,
            pNext: null(),
            memory: allocation.memory,
            offset: start,
            size,
        })
    }

    #[inline]
    fn lock(&self, ty: u32) -> MutexGuard<'_, Pool> {
        self.pools[ty as usize]
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
    }
}

impl Drop for Allocator {
    fn drop(&mut self) {
        for pool in &mut self.pools {
            let pool = pool.get_mut().unwrap_or_else(PoisonError::into_inner);
            debug_assert!(
                std::thread::panicking() || pool.statistics.allocation_count == 0,
                "allocator dropped with {} live allocations",
                pool.statistics.allocation_count
            );
            for block in pool.blocks.drain(..).flatten() {
                unsafe { (self.fns.free_memory)(self.device, block.memory, null()) };
            }
        }
    }
}

#[inline]
fn suballocation(
    block: &Block,
    ty: u32,
    slot: u32,
    node: u32,
    offset: u64,
    size: u64,
) -> Allocation {
    Allocation {
        memory: block.memory,
        offset,
        size,
        mapped: if block.mapped.is_null() {
            null_mut()
        } else {
            unsafe { block.mapped.add(offset as usize) }
        },
        memory_type: ty,
        block: slot,
        node,
    }
}

/// Memory types usable for `location`, cheapest first. Cost counts missing preferred flags
/// plus present unwanted flags, mirroring VMA's `VMA_MEMORY_USAGE_AUTO` heuristics.
fn preference_order(
    types: &[MemoryType],
    location: MemoryLocation,
    device_access: bool,
) -> Box<[u8]> {
    use VkMemoryPropertyFlags as F;
    let none = F::empty();
    let (required, preferred, unwanted) = match location {
        MemoryLocation::GpuOnly => (none, F::DEVICE_LOCAL, F::HOST_VISIBLE | F::HOST_CACHED),
        MemoryLocation::CpuToGpu if device_access => (
            F::HOST_VISIBLE,
            F::DEVICE_LOCAL | F::HOST_COHERENT,
            F::HOST_CACHED,
        ),
        MemoryLocation::CpuToGpu => (
            F::HOST_VISIBLE,
            F::HOST_COHERENT,
            F::DEVICE_LOCAL | F::HOST_CACHED,
        ),
        MemoryLocation::GpuToCpu => (F::HOST_VISIBLE, F::HOST_CACHED | F::HOST_COHERENT, none),
    };
    let mut order: Vec<(u32, u8)> = types
        .iter()
        .enumerate()
        .filter(|(_, t)| t.flags.contains(required) && !t.flags.intersects(UNUSABLE))
        .map(|(i, t)| {
            (
                (preferred - t.flags).bits().count_ones()
                    + (t.flags & unwanted).bits().count_ones(),
                i as u8,
            )
        })
        .collect();
    order.sort_unstable();
    order.into_iter().map(|(_, i)| i).collect()
}