Ensure scratch buffers have enough memory to meet alignment requirements

.gitignore

@@ -1 +1,2 @@
 /target
+.vscode

kea_gpu/src/ray_tracing/acceleration_structures/build.rs

@@ -1,16 +1,15 @@
-use std::sync::Arc;
-
-use ash::vk;
-use gpu_allocator::MemoryLocation;
-
 use crate::{
     core::buffer::{AllocatedBuffer, Buffer},
     device::Device,
     storage::memory,
 };
+use ash::vk;
+use gpu_allocator::MemoryLocation;
+use std::sync::Arc;
 
 pub struct ScratchBuffer {
     buffer: AllocatedBuffer,
+    size: u64,
     alignment: u32,
 }
 
@@ -21,10 +20,14 @@ impl ScratchBuffer {
             ..
         } = device.physical_device().acceleration_structure_properties();
 
-        let aligned_size = memory::align(size, alignment as u64);
+        // The buffer that is created may not meet the alignment requirements
+        // of the scratch, so we need the extra space to allow for passing an
+        // appropriately aligned address.
+        let max_size = size + alignment as u64;
+
         let buffer = Buffer::new(
             device,
-            aligned_size,
+            max_size,
             vk::BufferUsageFlags::STORAGE_BUFFER | vk::BufferUsageFlags::SHADER_DEVICE_ADDRESS,
         )
         .allocate(
@@ -32,10 +35,21 @@ impl ScratchBuffer {
             MemoryLocation::GpuOnly,
         );
 
-        ScratchBuffer { buffer, alignment }
+        ScratchBuffer {
+            buffer,
+            size,
+            alignment,
+        }
     }
 
     pub fn device_address(&self) -> vk::DeviceAddress {
-        memory::align(self.buffer.device_address(), self.alignment as _)
+        let address = self.buffer.device_address();
+        let aligned_address = memory::align(address, self.alignment as _);
+
+        assert!(
+            (aligned_address + self.size as u64) <= (address + self.buffer.allocated_size() as u64)
+        );
+
+        aligned_address
     }
 }