alex/volcanic-mem

root / tests / gpu.rs

mod common;

use common::{Context, VALIDATION_ERRORS, check};
use core::ptr::null;
use std::collections::HashMap;
use std::sync::atomic::Ordering;
use volcanic::*;
use volcanic_mem::*;

struct Rng(u64);

impl Rng {
    fn next(&mut self) -> u64 {
        self.0 ^= self.0 << 13;
        self.0 ^= self.0 >> 7;
        self.0 ^= self.0 << 17;
        self.0
    }

    fn below(&mut self, n: u64) -> u64 {
        self.next() % n
    }

    fn log_uniform(&mut self, min: u64, max: u64) -> u64 {
        let bits = self.below(64 - max.leading_zeros() as u64) + 1;
        (self.next() & ((1 << bits) - 1)).clamp(min, max)
    }
}

fn pattern(id: usize, i: usize) -> u8 {
    ((id as u64).wrapping_mul(0x9e37_79b9_7f4a_7c15)
        ^ (i as u64).wrapping_mul(0xbf58_476d_1ce4_e5b9) >> 29) as u8
}

fn image_info(width: u32, height: u32) -> VkImageCreateInfo {
    VkImageCreateInfo {
        sType: VkStructureType::ImageCreateInfo,
        pNext: null(),
        flags: VkImageCreateFlags::empty(),
        imageType: VkImageType::_2D,
        format: VkFormat::R8G8B8A8Uint,
        extent: VkExtent3D {
            width,
            height,
            depth: 1,
        },
        mipLevels: 1,
        arrayLayers: 1,
        samples: VkSampleCountFlagBits::_1,
        tiling: VkImageTiling::Optimal,
        usage: VkImageUsageFlags::TRANSFER_SRC | VkImageUsageFlags::TRANSFER_DST,
        sharingMode: VkSharingMode::Exclusive,
        queueFamilyIndexCount: 0,
        pQueueFamilyIndices: null(),
        initialLayout: VkImageLayout::Undefined,
    }
}

fn color_layers() -> VkImageSubresourceLayers {
    VkImageSubresourceLayers {
        aspectMask: VkImageAspectFlags::COLOR,
        mipLevel: 0,
        baseArrayLayer: 0,
        layerCount: 1,
    }
}

fn layout_barrier(image: VkImage, old: VkImageLayout, new: VkImageLayout) -> VkImageMemoryBarrier {
    VkImageMemoryBarrier {
        sType: VkStructureType::ImageMemoryBarrier,
        pNext: null(),
        srcAccessMask: if old == VkImageLayout::Undefined {
            VkAccessFlags::empty()
        } else {
            VkAccessFlags::TRANSFER_WRITE
        },
        dstAccessMask: VkAccessFlags::TRANSFER_READ | VkAccessFlags::TRANSFER_WRITE,
        oldLayout: old,
        newLayout: new,
        srcQueueFamilyIndex: VK_QUEUE_FAMILY_IGNORED,
        dstQueueFamilyIndex: VK_QUEUE_FAMILY_IGNORED,
        image,
        subresourceRange: VkImageSubresourceRange {
            aspectMask: VkImageAspectFlags::COLOR,
            baseMipLevel: 0,
            levelCount: 1,
            baseArrayLayer: 0,
            layerCount: 1,
        },
    }
}

fn assert_no_validation_errors() {
    assert_eq!(
        VALIDATION_ERRORS.load(Ordering::Relaxed),
        0,
        "validation layer reported errors"
    );
}

/// Fills many interleaved buffers and optimal images with unique patterns via the GPU and reads
/// them back. Any overlapping or granularity-violating placement corrupts some pattern.
#[test]
fn gpu_round_trip_detects_no_overlap() {
    let ctx = Context::new();
    let allocator = ctx.allocator();
    let mut rng = Rng(0xdead_beef);
    let usage = VkBufferUsageFlags2::TRANSFER_SRC
        | VkBufferUsageFlags2::TRANSFER_DST
        | VkBufferUsageFlags2::STORAGE_BUFFER;

    let mut buffers = Vec::new();
    let mut images = Vec::new();
    for round in 0..4 {
        for _ in 0..1500 {
            let desc = BufferDesc {
                size: rng.log_uniform(4, 4 << 20),
                usage,
                location: MemoryLocation::GpuOnly,
            };
            buffers.push(allocator.create_buffer(&desc).unwrap());
            if rng.below(8) == 0 {
                let (w, h) = (
                    rng.log_uniform(1, 256) as u32,
                    rng.log_uniform(1, 256) as u32,
                );
                images.push((
                    allocator
                        .create_image(&image_info(w, h), MemoryLocation::GpuOnly)
                        .unwrap(),
                    w,
                    h,
                ));
            }
        }
        if round < 3 {
            for _ in 0..buffers.len() / 2 {
                let buffer = buffers.swap_remove(rng.below(buffers.len() as u64) as usize);
                allocator.destroy_buffer(buffer);
            }
            for _ in 0..images.len() / 2 {
                let (image, ..) = images.swap_remove(rng.below(images.len() as u64) as usize);
                allocator.destroy_image(image);
            }
        }
    }

    let mut offsets = Vec::new();
    let mut total = 0u64;
    for size in buffers
        .iter()
        .map(Buffer::size)
        .chain(images.iter().map(|&(_, w, h)| (w * h * 4) as u64))
    {
        offsets.push(total);
        total = (total + size).next_multiple_of(16);
    }
    let mut staging = allocator
        .create_buffer(&BufferDesc {
            size: total,
            usage: VkBufferUsageFlags2::TRANSFER_SRC,
            location: MemoryLocation::CpuToGpu,
        })
        .unwrap();
    let readback = allocator
        .create_buffer(&BufferDesc {
            size: total,
            usage: VkBufferUsageFlags2::TRANSFER_DST,
            location: MemoryLocation::GpuToCpu,
        })
        .unwrap();
    let sizes: Vec<u64> = offsets
        .windows(2)
        .map(|w| w[1] - w[0])
        .chain([total - offsets.last().unwrap()])
        .collect();
    let bytes = staging.mapped_slice_mut().unwrap();
    for (id, (&offset, &size)) in offsets.iter().zip(&sizes).enumerate() {
        for i in 0..size as usize {
            bytes[offset as usize + i] = pattern(id, i);
        }
    }
    allocator
        .flush(staging.allocation(), 0, VK_WHOLE_SIZE)
        .unwrap();

    let d = &ctx.device;
    ctx.submit(|cmd| unsafe {
        let to_dst: Vec<_> = images
            .iter()
            .map(|(image, ..)| {
                layout_barrier(
                    image.handle(),
                    VkImageLayout::Undefined,
                    VkImageLayout::TransferDstOptimal,
                )
            })
            .collect();
        let transfer = VkPipelineStageFlags::TRANSFER;
        let none = VkDependencyFlags::empty();
        d.vkCmdPipelineBarrier(
            cmd,
            VkPipelineStageFlags::TOP_OF_PIPE,
            transfer,
            none,
            0,
            null(),
            0,
            null(),
            to_dst.len() as u32,
            to_dst.as_ptr(),
        );
        for (buffer, &offset) in buffers.iter().zip(&offsets) {
            let region = VkBufferCopy {
                srcOffset: offset,
                dstOffset: 0,
                size: buffer.size(),
            };
            d.vkCmdCopyBuffer(cmd, staging.handle(), buffer.handle(), 1, &region);
        }
        for ((image, w, h), &offset) in images.iter().zip(&offsets[buffers.len()..]) {
            let region = VkBufferImageCopy {
                bufferOffset: offset,
                bufferRowLength: 0,
                bufferImageHeight: 0,
                imageSubresource: color_layers(),
                imageOffset: VkOffset3D { x: 0, y: 0, z: 0 },
                imageExtent: VkExtent3D {
                    width: *w,
                    height: *h,
                    depth: 1,
                },
            };
            d.vkCmdCopyBufferToImage(
                cmd,
                staging.handle(),
                image.handle(),
                VkImageLayout::TransferDstOptimal,
                1,
                &region,
            );
        }
        let memory = VkMemoryBarrier {
            sType: VkStructureType::MemoryBarrier,
            pNext: null(),
            srcAccessMask: VkAccessFlags::TRANSFER_WRITE,
            dstAccessMask: VkAccessFlags::TRANSFER_READ,
        };
        let to_src: Vec<_> = images
            .iter()
            .map(|(image, ..)| {
                layout_barrier(
                    image.handle(),
                    VkImageLayout::TransferDstOptimal,
                    VkImageLayout::TransferSrcOptimal,
                )
            })
            .collect();
        d.vkCmdPipelineBarrier(
            cmd,
            transfer,
            transfer,
            none,
            1,
            &memory,
            0,
            null(),
            to_src.len() as u32,
            to_src.as_ptr(),
        );
        for (buffer, &offset) in buffers.iter().zip(&offsets) {
            let region = VkBufferCopy {
                srcOffset: 0,
                dstOffset: offset,
                size: buffer.size(),
            };
            d.vkCmdCopyBuffer(cmd, buffer.handle(), readback.handle(), 1, &region);
        }
        for ((image, w, h), &offset) in images.iter().zip(&offsets[buffers.len()..]) {
            let region = VkBufferImageCopy {
                bufferOffset: offset,
                bufferRowLength: 0,
                bufferImageHeight: 0,
                imageSubresource: color_layers(),
                imageOffset: VkOffset3D { x: 0, y: 0, z: 0 },
                imageExtent: VkExtent3D {
                    width: *w,
                    height: *h,
                    depth: 1,
                },
            };
            d.vkCmdCopyImageToBuffer(
                cmd,
                image.handle(),
                VkImageLayout::TransferSrcOptimal,
                readback.handle(),
                1,
                &region,
            );
        }
        let host = VkMemoryBarrier {
            sType: VkStructureType::MemoryBarrier,
            pNext: null(),
            srcAccessMask: VkAccessFlags::TRANSFER_WRITE,
            dstAccessMask: VkAccessFlags::HOST_READ,
        };
        d.vkCmdPipelineBarrier(
            cmd,
            transfer,
            VkPipelineStageFlags::HOST,
            none,
            1,
            &host,
            0,
            null(),
            0,
            null(),
        );
    });
    allocator
        .invalidate(readback.allocation(), 0, VK_WHOLE_SIZE)
        .unwrap();

    let sent = staging.mapped_slice().unwrap();
    let received = readback.mapped_slice().unwrap();
    let ids = buffers.len() + images.len();
    for id in 0..ids {
        let (offset, size) = (
            offsets[id] as usize,
            if id < buffers.len() {
                buffers[id].size()
            } else {
                let (_, w, h) = images[id - buffers.len()];
                (w * h * 4) as u64
            } as usize,
        );
        assert!(
            sent[offset..offset + size] == received[offset..offset + size],
            "resource {id} corrupted"
        );
    }

    let stats = allocator.statistics();
    println!(
        "{} buffers, {} images, {stats:?}",
        buffers.len(),
        images.len()
    );
    for buffer in buffers.into_iter().chain([staging, readback]) {
        allocator.destroy_buffer(buffer);
    }
    for (image, ..) in images {
        allocator.destroy_image(image);
    }
    let stats = allocator.statistics();
    assert_eq!((stats.allocation_count, stats.allocation_bytes), (0, 0));
    drop(allocator);
    assert_no_validation_errors();
}

#[test]
fn descriptor_heaps_bind_and_accept_descriptors() {
    let ctx = Context::new();
    let allocator = ctx.allocator();
    let mut resources = allocator
        .create_descriptor_heap(DescriptorHeapKind::Resource, 64 << 10)
        .unwrap();
    let samplers = allocator
        .create_descriptor_heap(DescriptorHeapKind::Sampler, 4 << 10)
        .unwrap();
    assert!(resources.capacity() >= 64 << 10);
    assert!(resources.buffer().allocation().mapped_ptr().is_some());

    let info = resources.bind_info();
    assert!(info.reservedRangeOffset + info.reservedRangeSize <= info.heapRange.size);
    assert_eq!(info.reservedRangeOffset, resources.capacity());

    let storage = allocator
        .create_buffer(&BufferDesc {
            size: 4096,
            usage: VkBufferUsageFlags2::STORAGE_BUFFER,
            location: MemoryLocation::GpuOnly,
        })
        .unwrap();
    assert_ne!(storage.address(), 0);
    unsafe {
        let descriptor_size = ctx.instance.vkGetPhysicalDeviceDescriptorSizeEXT(
            ctx.physical_device,
            VkDescriptorType::StorageBuffer,
        );
        let range = storage.range();
        let resource = VkResourceDescriptorInfoEXT {
            sType: VkStructureType::ResourceDescriptorInfoExt,
            pNext: null(),
            r#type: VkDescriptorType::StorageBuffer,
            data: VkResourceDescriptorDataEXT {
                pAddressRange: &range,
            },
        };
        let target = resources.host_range(0, descriptor_size);
        check(
            ctx.device
                .vkWriteResourceDescriptorsEXT(ctx.device_handle, 1, &resource, &target),
        );
        assert!(
            resources.descriptors_mut()[..descriptor_size as usize]
                .iter()
                .any(|&b| b != 0)
        );

        let (resource_info, sampler_info) = (resources.bind_info(), samplers.bind_info());
        ctx.submit(|cmd| {
            ctx.device.vkCmdBindResourceHeapEXT(cmd, &resource_info);
            ctx.device.vkCmdBindSamplerHeapEXT(cmd, &sampler_info);
        });
    }

    allocator.destroy_buffer(storage);
    allocator.destroy_descriptor_heap(resources);
    allocator.destroy_descriptor_heap(samplers);
    drop(allocator);
    assert_no_validation_errors();
}

#[test]
fn large_buffers_get_dedicated_memory_and_stats_balance() {
    let ctx = Context::new();
    let allocator = ctx.allocator();
    let usage = VkBufferUsageFlags2::STORAGE_BUFFER;
    let small = allocator
        .create_buffer(&BufferDesc {
            size: 1 << 20,
            usage,
            location: MemoryLocation::GpuOnly,
        })
        .unwrap();
    let large = allocator
        .create_buffer(&BufferDesc {
            size: 200 << 20,
            usage,
            location: MemoryLocation::GpuOnly,
        })
        .unwrap();
    assert!(!small.allocation().is_dedicated());
    assert!(large.allocation().is_dedicated());

    let mut upload = allocator
        .create_buffer(&BufferDesc {
            size: 1 << 16,
            usage,
            location: MemoryLocation::CpuToGpu,
        })
        .unwrap();
    upload.mapped_slice_mut().unwrap().fill(0xab);
    allocator
        .flush(upload.allocation(), 0, VK_WHOLE_SIZE)
        .unwrap();

    let stats = allocator.statistics();
    assert_eq!(stats.allocation_count, 3);
    assert!(stats.block_bytes >= stats.allocation_bytes);
    let budgets = allocator.budgets();
    assert!(
        budgets
            .iter()
            .any(|b| b.statistics.block_bytes >= 200 << 20)
    );
    for budget in &budgets {
        println!("{budget:?}");
    }

    allocator.destroy_buffer(small);
    allocator.destroy_buffer(large);
    allocator.destroy_buffer(upload);
    let stats = allocator.statistics();
    assert_eq!((stats.allocation_count, stats.allocation_bytes), (0, 0));
    assert!(
        stats.block_count <= 2,
        "at most one empty block kept per memory type"
    );
    drop(allocator);
    assert_no_validation_errors();
}

/// Hammers the internally synchronized pools from many threads and checks that no two live
/// allocations overlap in the same `VkDeviceMemory`.
#[test]
fn concurrent_allocations_never_overlap() {
    let ctx = Context::new();
    let allocator = ctx.allocator();
    let probe = allocator
        .create_buffer(&BufferDesc {
            size: 256,
            usage: VkBufferUsageFlags2::STORAGE_BUFFER,
            location: MemoryLocation::GpuOnly,
        })
        .unwrap();
    let type_bits = 1 << probe.allocation().memory_type();
    allocator.destroy_buffer(probe);

    let live: Vec<Allocation> = std::thread::scope(|scope| {
        let handles: Vec<_> = (0..8)
            .map(|t| {
                let allocator = &allocator;
                scope.spawn(move || {
                    let mut rng = Rng(0x1234_5678 + t);
                    let mut mine = Vec::new();
                    for _ in 0..40_000 {
                        if mine.is_empty() || rng.below(100) < 52 {
                            let desc = AllocationDesc {
                                requirements: VkMemoryRequirements {
                                    size: rng.log_uniform(16, 1 << 20),
                                    alignment: 1 << rng.below(12),
                                    memoryTypeBits: type_bits,
                                },
                                location: MemoryLocation::GpuOnly,
                                kind: if rng.below(2) == 0 {
                                    ResourceKind::Linear
                                } else {
                                    ResourceKind::Optimal
                                },
                                dedicated: false,
                            };
                            mine.push(allocator.allocate(&desc).unwrap());
                        } else {
                            allocator.free(mine.swap_remove(rng.below(mine.len() as u64) as usize));
                        }
                    }
                    mine
                })
            })
            .collect();
        handles
            .into_iter()
            .flat_map(|h| h.join().unwrap())
            .collect()
    });

    let mut by_memory: HashMap<VkDeviceMemory, Vec<(u64, u64)>> = HashMap::new();
    for a in &live {
        by_memory
            .entry(a.memory())
            .or_default()
            .push((a.offset(), a.offset() + a.size()));
    }
    for ranges in by_memory.values_mut() {
        ranges.sort_unstable();
        for pair in ranges.windows(2) {
            assert!(pair[0].1 <= pair[1].0, "overlap {:?}", pair);
        }
    }
    let stats = allocator.statistics();
    assert_eq!(stats.allocation_count as usize, live.len());
    for allocation in live {
        allocator.free(allocation);
    }
    assert_eq!(allocator.statistics().allocation_count, 0);
    drop(allocator);
    assert_no_validation_errors();
}