alex/volcanic-mem

root / src / resource.rs

use crate::allocator::{
    Allocation, Allocator, Dedicated, MemoryLocation, Request, ResourceKind, Target,
};
use crate::error::{Error, Result, check};
use core::ptr::{null, null_mut};
use volcanic::*;

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

// ---- Buffers ------------------------------------------------

#[derive(Clone, Copy, Debug)]
pub struct BufferDesc {
    pub size: u64,
    /// `SHADER_DEVICE_ADDRESS` is always added, since descriptor heaps reference buffers by address.
    pub usage: VkBufferUsageFlags2,
    pub location: MemoryLocation,
}

/// A bound buffer with its device address resolved.
#[derive(Debug)]
pub struct Buffer {
    handle: VkBuffer,
    address: VkDeviceAddress,
    size: u64,
    allocation: Allocation,
}

impl Buffer {
    #[inline]
    pub fn handle(&self) -> VkBuffer {
        self.handle
    }

    #[inline]
    pub fn address(&self) -> VkDeviceAddress {
        self.address
    }

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

    /// Whole-buffer address range, as used by `VkResourceDescriptorInfoEXT`.
    #[inline]
    pub fn range(&self) -> VkDeviceAddressRangeEXT {
        VkDeviceAddressRangeEXT {
            address: self.address,
            size: self.size,
        }
    }

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

    #[inline]
    pub fn mapped_slice(&self) -> Option<&[u8]> {
        self.allocation
            .mapped_slice()
            .map(|bytes| &bytes[..self.size as usize])
    }

    #[inline]
    pub fn mapped_slice_mut(&mut self) -> Option<&mut [u8]> {
        let size = self.size as usize;
        self.allocation
            .mapped_slice_mut()
            .map(|bytes| &mut bytes[..size])
    }
}

// ---- Images -------------------------------------------------

#[derive(Debug)]
pub struct Image {
    handle: VkImage,
    allocation: Allocation,
}

impl Image {
    #[inline]
    pub fn handle(&self) -> VkImage {
        self.handle
    }

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

// ---- Descriptor heaps ---------------------------------------

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum DescriptorHeapKind {
    Resource,
    Sampler,
    /// A sampler heap bound while pipelines with embedded samplers are in use.
    SamplerWithEmbedded,
}

pub(crate) struct HeapLimits {
    sampler_alignment: u64,
    resource_alignment: u64,
    max_sampler_size: u64,
    max_resource_size: u64,
    sampler_reserved: u64,
    sampler_reserved_embedded: u64,
    resource_reserved: u64,
    sampler_descriptor_alignment: u64,
    resource_descriptor_alignment: u64,
}

impl HeapLimits {
    pub fn new(p: &VkPhysicalDeviceDescriptorHeapPropertiesEXT) -> Self {
        Self {
            sampler_alignment: p.samplerHeapAlignment.max(1),
            resource_alignment: p.resourceHeapAlignment.max(1),
            max_sampler_size: p.maxSamplerHeapSize,
            max_resource_size: p.maxResourceHeapSize,
            sampler_reserved: p.minSamplerHeapReservedRange,
            sampler_reserved_embedded: p.minSamplerHeapReservedRangeWithEmbedded,
            resource_reserved: p.minResourceHeapReservedRange,
            sampler_descriptor_alignment: p.samplerDescriptorAlignment.max(1),
            resource_descriptor_alignment: p
                .bufferDescriptorAlignment
                .max(p.imageDescriptorAlignment)
                .max(1),
        }
    }
}

/// A host-visible descriptor heap buffer. The application range comes first, followed by the
/// implementation-reserved range.
#[derive(Debug)]
pub struct DescriptorHeap {
    buffer: Buffer,
    kind: DescriptorHeapKind,
    reserved_offset: u64,
    reserved_size: u64,
}

impl DescriptorHeap {
    #[inline]
    pub fn buffer(&self) -> &Buffer {
        &self.buffer
    }

    #[inline]
    pub fn kind(&self) -> DescriptorHeapKind {
        self.kind
    }

    /// Bytes available for application descriptors.
    #[inline]
    pub fn capacity(&self) -> u64 {
        self.reserved_offset
    }

    /// Ready-to-use info for `vkCmdBindResourceHeapEXT` or `vkCmdBindSamplerHeapEXT`.
    #[inline]
    pub fn bind_info(&self) -> VkBindHeapInfoEXT {
        VkBindHeapInfoEXT {
            sType: VkStructureType::BindHeapInfoExt,
            pNext: null(),
            heapRange: self.buffer.range(),
            reservedRangeOffset: self.reserved_offset,
            reservedRangeSize: self.reserved_size,
        }
    }

    /// The application descriptor range, for `vkWrite*DescriptorsEXT` destinations or plain copies.
    #[inline]
    pub fn descriptors_mut(&mut self) -> &mut [u8] {
        let capacity = self.reserved_offset as usize;
        let bytes = self
            .buffer
            .allocation
            .mapped_slice_mut()
            .expect("descriptor heaps are host visible");
        &mut bytes[..capacity]
    }

    /// Host address range of `size` bytes at `offset`, for `vkWrite*DescriptorsEXT`.
    #[inline]
    pub fn host_range(&mut self, offset: u64, size: u64) -> VkHostAddressRangeEXT {
        let bytes = &mut self.descriptors_mut()[offset as usize..(offset + size) as usize];
        VkHostAddressRangeEXT {
            address: bytes.as_mut_ptr().cast(),
            size: bytes.len(),
        }
    }
}

// ---- Allocator integration ----------------------------------

impl Allocator {
    /// Creates a buffer, binds memory to it and resolves its device address.
    pub fn create_buffer(&self, desc: &BufferDesc) -> Result<Buffer> {
        self.create_buffer_aligned(desc, 1)
    }

    pub fn destroy_buffer(&self, buffer: Buffer) {
        unsafe { (self.fns().destroy_buffer)(self.device(), buffer.handle, null()) };
        self.free(buffer.allocation);
    }

    /// Creates an image from an explicit create info and binds memory to it.
    pub fn create_image(
        &self,
        info: &VkImageCreateInfo,
        location: MemoryLocation,
    ) -> Result<Image> {
        let fns = self.fns();
        let mut handle = 0;
        check("vkCreateImage", unsafe {
            (fns.create_image)(self.device(), info, null(), &mut handle)
        })?;

        let requirements_info = VkImageMemoryRequirementsInfo2 {
            sType: VkStructureType::ImageMemoryRequirementsInfo2,
            pNext: null(),
            image: handle,
        };
        let (requirements, dedicated) = query_requirements(|out| unsafe {
            (fns.get_image_memory_requirements2)(self.device(), &requirements_info, out)
        });
        let request = Request {
            size: requirements.size,
            alignment: requirements.alignment,
            type_bits: requirements.memoryTypeBits,
            location,
            kind: if info.tiling == VkImageTiling::Linear {
                ResourceKind::Linear
            } else {
                ResourceKind::Optimal
            },
            device_access: true,
            dedicated,
            target: Target::Image(handle),
        };
        let bound = self.allocate_request(&request).and_then(|allocation| {
            match check("vkBindImageMemory", unsafe {
                (fns.bind_image_memory)(
                    self.device(),
                    handle,
                    allocation.memory(),
                    allocation.offset(),
                )
            }) {
                Ok(()) => Ok(allocation),
                Err(error) => {
                    self.free(allocation);
                    Err(error)
                }
            }
        });
        match bound {
            Ok(allocation) => Ok(Image { handle, allocation }),
            Err(error) => {
                unsafe { (fns.destroy_image)(self.device(), handle, null()) };
                Err(error)
            }
        }
    }

    pub fn destroy_image(&self, image: Image) {
        unsafe { (self.fns().destroy_image)(self.device(), image.handle, null()) };
        self.free(image.allocation);
    }

    /// Creates a host-visible descriptor heap with `capacity` bytes for application descriptors.
    /// The implementation-reserved range is appended and all alignment limits are honored.
    pub fn create_descriptor_heap(
        &self,
        kind: DescriptorHeapKind,
        capacity: u64,
    ) -> Result<DescriptorHeap> {
        let l = &self.heap_limits;
        let (heap_alignment, descriptor_alignment, reserved_size, max) = match kind {
            DescriptorHeapKind::Resource => (
                l.resource_alignment,
                l.resource_descriptor_alignment,
                l.resource_reserved,
                l.max_resource_size,
            ),
            DescriptorHeapKind::Sampler => (
                l.sampler_alignment,
                l.sampler_descriptor_alignment,
                l.sampler_reserved,
                l.max_sampler_size,
            ),
            DescriptorHeapKind::SamplerWithEmbedded => (
                l.sampler_alignment,
                l.sampler_descriptor_alignment,
                l.sampler_reserved_embedded,
                l.max_sampler_size,
            ),
        };
        let reserved_offset = align_up(capacity, descriptor_alignment);
        let size = reserved_offset + reserved_size;
        if size > max {
            return Err(Error::HeapTooLarge { size, max });
        }
        let buffer = self.create_buffer_aligned(
            &BufferDesc {
                size,
                usage: VkBufferUsageFlags2::DESCRIPTOR_HEAP_EXT,
                location: MemoryLocation::CpuToGpu,
            },
            heap_alignment,
        )?;
        if buffer.address % heap_alignment != 0 {
            let address = buffer.address;
            self.destroy_buffer(buffer);
            return Err(Error::MisalignedHeap {
                address,
                alignment: heap_alignment,
            });
        }
        Ok(DescriptorHeap {
            buffer,
            kind,
            reserved_offset,
            reserved_size,
        })
    }

    pub fn destroy_descriptor_heap(&self, heap: DescriptorHeap) {
        self.destroy_buffer(heap.buffer);
    }

    fn create_buffer_aligned(&self, desc: &BufferDesc, min_alignment: u64) -> Result<Buffer> {
        if desc.size == 0 {
            return Err(Error::ZeroSize);
        }
        let fns = self.fns();
        let usage_info = VkBufferUsageFlags2CreateInfo {
            sType: VkStructureType::BufferUsageFlags2CreateInfo,
            pNext: null(),
            usage: desc.usage | VkBufferUsageFlags2::SHADER_DEVICE_ADDRESS,
        };
        let info = VkBufferCreateInfo {
            sType: VkStructureType::BufferCreateInfo,
            pNext: (&raw const usage_info).cast(),
            flags: VkBufferCreateFlags::empty(),
            size: desc.size,
            usage: VkBufferUsageFlags::empty(),
            sharingMode: VkSharingMode::Exclusive,
            queueFamilyIndexCount: 0,
            pQueueFamilyIndices: null(),
        };
        let mut handle = 0;
        check("vkCreateBuffer", unsafe {
            (fns.create_buffer)(self.device(), &info, null(), &mut handle)
        })?;

        let requirements_info = VkBufferMemoryRequirementsInfo2 {
            sType: VkStructureType::BufferMemoryRequirementsInfo2,
            pNext: null(),
            buffer: handle,
        };
        let (requirements, dedicated) = query_requirements(|out| unsafe {
            (fns.get_buffer_memory_requirements2)(self.device(), &requirements_info, out)
        });
        let transfer = VkBufferUsageFlags2::TRANSFER_SRC | VkBufferUsageFlags2::TRANSFER_DST;
        let request = Request {
            size: requirements.size,
            alignment: requirements.alignment.max(min_alignment),
            type_bits: requirements.memoryTypeBits,
            location: desc.location,
            kind: ResourceKind::Linear,
            device_access: !(desc.usage - transfer).is_empty(),
            dedicated,
            target: Target::Buffer(handle),
        };
        let bound = self.allocate_request(&request).and_then(|allocation| {
            match check("vkBindBufferMemory", unsafe {
                (fns.bind_buffer_memory)(
                    self.device(),
                    handle,
                    allocation.memory(),
                    allocation.offset(),
                )
            }) {
                Ok(()) => Ok(allocation),
                Err(error) => {
                    self.free(allocation);
                    Err(error)
                }
            }
        });
        let allocation = match bound {
            Ok(allocation) => allocation,
            Err(error) => {
                unsafe { (fns.destroy_buffer)(self.device(), handle, null()) };
                return Err(error);
            }
        };
        let address_info = VkBufferDeviceAddressInfo {
            sType: VkStructureType::BufferDeviceAddressInfo,
            pNext: null(),
            buffer: handle,
        };
        Ok(Buffer {
            handle,
            address: unsafe { (fns.get_buffer_device_address)(self.device(), &address_info) },
            size: desc.size,
            allocation,
        })
    }
}

fn query_requirements(
    query: impl FnOnce(*mut VkMemoryRequirements2),
) -> (VkMemoryRequirements, Dedicated) {
    let mut dedicated = VkMemoryDedicatedRequirements {
        sType: VkStructureType::MemoryDedicatedRequirements,
        pNext: null_mut(),
        prefersDedicatedAllocation: 0,
        requiresDedicatedAllocation: 0,
    };
    let mut requirements = VkMemoryRequirements2 {
        sType: VkStructureType::MemoryRequirements2,
        pNext: (&raw mut dedicated).cast(),
        memoryRequirements: VkMemoryRequirements {
            size: 0,
            alignment: 0,
            memoryTypeBits: 0,
        },
    };
    query(&mut requirements);
    let dedicated = if dedicated.requiresDedicatedAllocation == VK_TRUE {
        Dedicated::Require
    } else if dedicated.prefersDedicatedAllocation == VK_TRUE {
        Dedicated::Prefer
    } else {
        Dedicated::No
    };
    (requirements.memoryRequirements, dedicated)
}