use std::num::NonZeroU64;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
pub(crate) const UNIFORM_SIZE: u64 = 512;
pub(crate) const KERNEL_SIZE: u64 =
(crate::planner::gaussian::MAX_KERNEL_SAMPLES * std::mem::size_of::<[f32; 4]>()) as u64;
const FRAMES: usize = 3;
const IDLE_PASSES: u8 = 3;
const SLOTS_PER_BLOCK: u64 = 1024;
const VERTEX_BLOCK_SIZE: u64 = 256 * 1024;
pub struct HostBuffer {
factory: BlockFactory,
frames: [FrameArena; FRAMES],
frame: usize,
stride: u64,
kernel_stride: u64,
uniform_block_size: u64,
blocks_created: u32,
}
struct BlockFactory {
device: wgpu::Device,
record_layout: wgpu::BindGroupLayout,
kernel_layout: wgpu::BindGroupLayout,
writer: BlockWriter,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum BlockWriter {
Queue,
Mapped,
}
#[derive(Default)]
struct FrameArena {
uniforms: Vec<Block<UniformBindGroups>>,
cursor: Cursor,
vertices: Vec<Block<()>>,
vertex_cursor: Cursor,
}
#[derive(Default, Clone, Copy)]
struct Cursor {
block: usize,
offset: u64,
}
struct Block<Views> {
buffer: wgpu::Buffer,
views: Views,
scratch: Vec<u8>,
used: u64,
idle: u8,
writable: Arc<AtomicBool>,
}
struct UniformBindGroups {
record: wgpu::BindGroup,
kernel: wgpu::BindGroup,
}
#[derive(Clone, Copy, Debug)]
pub(crate) struct UniformSlot {
pub block: usize,
pub offset: u32,
}
#[derive(Clone, Copy, Debug)]
pub(crate) struct VertexSlot {
pub block: usize,
pub offset: u64,
pub bytes: u64,
}
#[must_use = "a flushed frame's blocks are remapped by `after_submit`"]
pub struct Flushed {
pub uniform_bytes: u64,
pub vertex_bytes: u64,
pub blocks_created: u32,
remaps: Vec<Remap>,
}
struct Remap {
buffer: wgpu::Buffer,
writable: Arc<AtomicBool>,
}
impl HostBuffer {
pub fn new(device: &wgpu::Device) -> Self {
let record_layout = dynamic_uniform_layout(
device,
"valo.host_buffer",
wgpu::ShaderStages::VERTEX_FRAGMENT,
UNIFORM_SIZE,
);
let kernel_layout = dynamic_uniform_layout(
device,
"valo.host_buffer.kernel",
wgpu::ShaderStages::FRAGMENT,
KERNEL_SIZE,
);
let stride = (device.limits().min_uniform_buffer_offset_alignment as u64).max(UNIFORM_SIZE);
Self {
factory: BlockFactory {
device: device.clone(),
record_layout,
kernel_layout,
writer: BlockWriter::for_device(device),
},
frames: Default::default(),
frame: 0,
stride,
kernel_stride: KERNEL_SIZE.next_multiple_of(stride),
uniform_block_size: stride * SLOTS_PER_BLOCK,
blocks_created: 0,
}
}
pub fn bind_group_layout(&self) -> &wgpu::BindGroupLayout {
&self.factory.record_layout
}
pub fn kernel_bind_group_layout(&self) -> &wgpu::BindGroupLayout {
&self.factory.kernel_layout
}
pub fn begin_frame(&mut self) {
if self.factory.writer == BlockWriter::Mapped {
let _ = self.factory.device.poll(wgpu::PollType::Poll);
}
self.blocks_created = 0;
self.frame = (self.frame + 1) % FRAMES;
let arena = &mut self.frames[self.frame];
arena.cursor = Cursor::default();
arena.vertex_cursor = Cursor::default();
recycle(&mut arena.uniforms);
recycle(&mut arena.vertices);
}
pub(crate) fn alloc_uniform(&mut self, bytes: &[u8]) -> UniformSlot {
debug_assert!(bytes.len() as u64 <= self.stride);
self.alloc_uniform_footprint(bytes, self.stride)
}
pub(crate) fn alloc_kernel(&mut self, bytes: &[u8]) -> UniformSlot {
debug_assert!(bytes.len() as u64 == KERNEL_SIZE);
self.alloc_uniform_footprint(bytes, self.kernel_stride)
}
fn alloc_uniform_footprint(&mut self, bytes: &[u8], footprint: u64) -> UniformSlot {
let arena = &mut self.frames[self.frame];
let writer = self.factory.writer;
advance_cursor(&mut arena.cursor, &arena.uniforms, footprint, writer);
if arena.cursor.block >= arena.uniforms.len() {
let block = self.factory.uniform_block(self.uniform_block_size);
arena.uniforms.push(block);
self.blocks_created += 1;
}
let cursor = arena.cursor;
write_scratch(&mut arena.uniforms[cursor.block], cursor.offset, bytes);
arena.cursor.offset += footprint;
UniformSlot {
block: cursor.block,
offset: cursor.offset as u32,
}
}
pub(crate) fn alloc_vertices(&mut self, bytes: &[u8]) -> VertexSlot {
let len = bytes.len() as u64;
let arena = &mut self.frames[self.frame];
let writer = self.factory.writer;
advance_cursor(&mut arena.vertex_cursor, &arena.vertices, len, writer);
if arena.vertex_cursor.block >= arena.vertices.len() {
let block = self.factory.vertex_block(VERTEX_BLOCK_SIZE.max(len));
arena.vertices.push(block);
self.blocks_created += 1;
}
let cursor = arena.vertex_cursor;
write_scratch(&mut arena.vertices[cursor.block], cursor.offset, bytes);
arena.vertex_cursor.offset += len.next_multiple_of(4);
VertexSlot {
block: cursor.block,
offset: cursor.offset,
bytes: len,
}
}
pub fn flush(&mut self, queue: &wgpu::Queue) -> Flushed {
let arena = &mut self.frames[self.frame];
let mut flushed = Flushed {
uniform_bytes: arena.uniforms.iter().map(|b| b.used).sum(),
vertex_bytes: arena.vertices.iter().map(|b| b.used).sum(),
blocks_created: self.blocks_created,
remaps: Vec::new(),
};
let writer = self.factory.writer;
for block in arena.uniforms.iter_mut().filter(|block| block.used > 0) {
writer.write(block, queue, &mut flushed.remaps);
}
for block in arena.vertices.iter_mut().filter(|block| block.used > 0) {
writer.write(block, queue, &mut flushed.remaps);
}
flushed
}
pub fn after_submit(&mut self, flushed: Flushed) {
for remap in flushed.remaps {
let writable = remap.writable;
remap
.buffer
.slice(..)
.map_async(wgpu::MapMode::Write, move |result| {
if result.is_ok() {
writable.store(true, Ordering::Release);
}
});
}
}
pub(crate) fn report(&self) -> crate::PoolReport {
let mut count = 0u32;
let mut bytes = 0u64;
for arena in &self.frames {
let scratch = arena.uniforms.iter().map(|b| b.scratch.len());
for len in scratch.chain(arena.vertices.iter().map(|b| b.scratch.len())) {
count += 1;
bytes += len as u64;
}
}
crate::PoolReport { count, bytes }
}
pub(crate) fn bind_group(&self, block: usize) -> &wgpu::BindGroup {
&self.frames[self.frame].uniforms[block].views.record
}
pub(crate) fn kernel_bind_group(&self, block: usize) -> &wgpu::BindGroup {
&self.frames[self.frame].uniforms[block].views.kernel
}
pub(crate) fn vertex_buffer(&self, block: usize) -> &wgpu::Buffer {
&self.frames[self.frame].vertices[block].buffer
}
}
impl BlockWriter {
fn for_device(device: &wgpu::Device) -> Self {
if device
.features()
.contains(wgpu::Features::MAPPABLE_PRIMARY_BUFFERS)
{
BlockWriter::Mapped
} else {
BlockWriter::Queue
}
}
fn usage(self, base: wgpu::BufferUsages) -> wgpu::BufferUsages {
match self {
BlockWriter::Queue => base | wgpu::BufferUsages::COPY_DST,
BlockWriter::Mapped => base | wgpu::BufferUsages::MAP_WRITE,
}
}
fn can_write<Views>(self, block: &Block<Views>) -> bool {
match self {
BlockWriter::Queue => true,
BlockWriter::Mapped => block.writable.load(Ordering::Acquire),
}
}
fn write<Views>(self, block: &mut Block<Views>, queue: &wgpu::Queue, remaps: &mut Vec<Remap>) {
let used = &block.scratch[..block.used as usize];
match self {
BlockWriter::Queue => queue.write_buffer(&block.buffer, 0, used),
BlockWriter::Mapped => {
let mut view = block
.buffer
.slice(..block.used)
.get_mapped_range_mut()
.expect("a block allocated from is mapped for writing");
view.copy_from_slice(used);
drop(view);
block.buffer.unmap();
block.writable.store(false, Ordering::Release);
remaps.push(Remap {
buffer: block.buffer.clone(),
writable: Arc::clone(&block.writable),
});
}
}
}
}
impl BlockFactory {
fn uniform_block(&self, size: u64) -> Block<UniformBindGroups> {
let buffer = self.buffer(
"valo.host_buffer.uniforms",
wgpu::BufferUsages::UNIFORM,
size,
);
let bind_groups = UniformBindGroups {
record: uniform_view(&self.device, &self.record_layout, &buffer, UNIFORM_SIZE),
kernel: uniform_view(&self.device, &self.kernel_layout, &buffer, KERNEL_SIZE),
};
self.block(buffer, bind_groups, size)
}
fn vertex_block(&self, size: u64) -> Block<()> {
let buffer = self.buffer(
"valo.host_buffer.vertices",
wgpu::BufferUsages::VERTEX,
size,
);
self.block(buffer, (), size)
}
fn buffer(&self, label: &str, usage: wgpu::BufferUsages, size: u64) -> wgpu::Buffer {
let mapped = self.writer == BlockWriter::Mapped;
self.device.create_buffer(&wgpu::BufferDescriptor {
label: Some(label),
size,
usage: self.writer.usage(usage),
mapped_at_creation: mapped,
})
}
fn block<Views>(&self, buffer: wgpu::Buffer, views: Views, size: u64) -> Block<Views> {
Block {
buffer,
views,
scratch: vec![0; size as usize],
used: 0,
idle: 0,
writable: Arc::new(AtomicBool::new(self.writer == BlockWriter::Mapped)),
}
}
}
fn advance_cursor<Views>(
cursor: &mut Cursor,
blocks: &[Block<Views>],
needed: u64,
writer: BlockWriter,
) {
while let Some(block) = blocks.get(cursor.block) {
if writer.can_write(block) && cursor.offset + needed <= block.scratch.len() as u64 {
return;
}
cursor.block += 1;
cursor.offset = 0;
}
}
fn dynamic_uniform_layout(
device: &wgpu::Device,
label: &str,
visibility: wgpu::ShaderStages,
size: u64,
) -> wgpu::BindGroupLayout {
device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some(label),
entries: &[wgpu::BindGroupLayoutEntry {
binding: 0,
visibility,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: true,
min_binding_size: NonZeroU64::new(size),
},
count: None,
}],
})
}
fn uniform_view(
device: &wgpu::Device,
layout: &wgpu::BindGroupLayout,
buffer: &wgpu::Buffer,
size: u64,
) -> wgpu::BindGroup {
device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("valo.host_buffer.uniforms"),
layout,
entries: &[wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
buffer,
offset: 0,
size: NonZeroU64::new(size),
}),
}],
})
}
fn recycle<Views>(blocks: &mut Vec<Block<Views>>) {
for b in blocks.iter_mut() {
b.idle = if b.used > 0 {
0
} else {
b.idle.saturating_add(1)
};
b.used = 0;
}
while blocks.len() > 1 && blocks.last().is_some_and(|b| b.idle >= IDLE_PASSES) {
blocks.pop();
}
}
fn write_scratch<Views>(block: &mut Block<Views>, offset: u64, bytes: &[u8]) {
block.scratch[offset as usize..offset as usize + bytes.len()].copy_from_slice(bytes);
block.used = block.used.max(offset + bytes.len() as u64);
}
#[cfg(test)]
mod tests {
use super::*;
fn mappable_device() -> Option<(wgpu::Device, wgpu::Queue)> {
let instance = wgpu::Instance::default();
let adapter = pollster::block_on(instance.request_adapter(&Default::default())).ok()?;
if !adapter
.features()
.contains(wgpu::Features::MAPPABLE_PRIMARY_BUFFERS)
{
return None;
}
pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
required_features: wgpu::Features::MAPPABLE_PRIMARY_BUFFERS,
..Default::default()
}))
.ok()
}
#[test]
fn a_mapped_block_is_written_in_place_and_writable_again_by_its_next_turn() {
let Some((device, queue)) = mappable_device() else {
eprintln!("SKIP a_mapped_block_is_written_in_place: no mappable adapter");
return;
};
let mut host = HostBuffer::new(&device);
assert_eq!(host.factory.writer, BlockWriter::Mapped);
host.begin_frame();
let slot = host.alloc_uniform(&[7u8; 64]);
let flushed = host.flush(&queue);
assert_eq!(
(flushed.uniform_bytes, flushed.vertex_bytes),
(64, 0),
"the bytes used, not the stride"
);
let block = &host.frames[host.frame].uniforms[slot.block];
assert!(
!block.writable.load(Ordering::Acquire),
"unmapped for the GPU"
);
host.after_submit(flushed);
for _ in 0..FRAMES {
host.begin_frame();
}
let block = &host.frames[host.frame].uniforms[slot.block];
assert!(block.writable.load(Ordering::Acquire));
let again = host.alloc_uniform(&[8u8; 64]);
assert_eq!(again.block, slot.block, "the block is written again");
assert_eq!(host.blocks_created, 0, "and no block is created for it");
}
fn headless() -> Option<wgpu::Device> {
let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
let adapter =
pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default()))
.ok()?;
let (device, _queue) =
pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor::default())).ok()?;
Some(device)
}
#[test]
fn kernels_share_the_uniform_blocks() {
let Some(device) = headless() else {
eprintln!("SKIP kernels_share_the_uniform_blocks: no GPU adapter");
return;
};
let mut host = HostBuffer::new(&device);
host.begin_frame();
let before = host.alloc_uniform(&[1u8; UNIFORM_SIZE as usize]);
let kernel = host.alloc_kernel(&[2u8; KERNEL_SIZE as usize]);
let after = host.alloc_uniform(&[3u8; UNIFORM_SIZE as usize]);
assert_eq!((before.block, kernel.block, after.block), (0, 0, 0));
assert_eq!(u64::from(kernel.offset) % host.stride, 0);
assert!(u64::from(after.offset) >= u64::from(kernel.offset) + KERNEL_SIZE);
assert_eq!(host.blocks_created, 1);
}
#[test]
fn retained_mixed_size_blocks_never_overrun() {
let Some(device) = headless() else {
eprintln!("SKIP retained_mixed_size_blocks_never_overrun: no GPU adapter");
return;
};
let mut host = HostBuffer::new(&device);
host.begin_frame();
host.alloc_vertices(&vec![1u8; 1024 * 1024]);
host.alloc_vertices(&[2u8; 64]);
for _ in 0..FRAMES {
host.begin_frame();
}
let a = host.alloc_vertices(&vec![3u8; 200 * 1024]);
let b = host.alloc_vertices(&vec![4u8; 800 * 1024]); assert_eq!(a.block, 0);
assert_eq!(b.block, 0, "800 KB still fits the 1 MB block");
let c = host.alloc_vertices(&vec![5u8; 900 * 1024]);
assert_eq!(c.bytes, 900 * 1024);
assert!(c.block >= 2, "small retained block is skipped, not overrun");
}
}