alpha2_render/

uniform.rs

use std::collections::HashMap;
use std::num::NonZeroU64;
use std::sync::Mutex;

use crate::Shared;
use crate::retire::{Retired, Retirement};

/// Bytes per effect parameter in `EffectParams`; every type pads to one 16-byte slot.
pub(crate) const PARAM_SLOT_BYTES: usize = 16;

/// Bytes of an effect's parameter uniform: one slot per parameter, minimum one since a uniform
/// cannot be empty. Shared by the graph reservation and the pass write.
pub(crate) const fn param_uniform_bytes(count: usize) -> usize {
    let slots = if count == 0 { 1 } else { count };
    slots.saturating_mul(PARAM_SLOT_BYTES)
}

/// 32-bit words of the runtime uniform: two `vec4f` of values, then a `vec4u` with the ticks.
pub(crate) const RUNTIME_UNIFORM_WORDS: usize = 12;

/// Floats in composite.wesl's `BatchParams`: three `vec4f` of controls, then four per layer.
pub(crate) const BATCH_UNIFORM_FLOATS: usize = 12 + 16 * crate::composite::COMPOSITE_BATCH_LAYERS;

pub(crate) const fn float_uniform_bytes(values: usize) -> u64 {
    (values * size_of::<f32>()) as u64
}

// Bounds on what an unusual frame can leave idle; normal frames stay within one capacity's peak.
const MAX_IDLE_UNIFORM_BUFFERS: usize = 64;
const MAX_IDLE_UNIFORM_BYTES: u64 = 4 * 1024 * 1024;

pub(crate) fn f32_bytes<const N: usize>(values: [f32; N]) -> Vec<u8> {
    values
        .into_iter()
        .flat_map(f32::to_ne_bytes)
        .collect::<Vec<_>>()
}

/// Allocation and retention counters for submission-scoped uniform arenas.
#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
pub struct UniformBufferPoolStats {
    /// Physical buffers created since renderer construction.
    pub created: u64,
    /// Graph executions served by an idle buffer.
    pub reused: u64,
    /// Buffers dropped instead of retained, including failed graph encodes.
    pub discarded: u64,
    /// Arenas currently being encoded or awaiting submission.
    pub active: usize,
    /// Bytes currently owned by active arenas.
    pub active_bytes: u64,
    /// Maximum active bytes observed since renderer construction.
    pub high_water_bytes: u64,
    /// Unmapped buffers currently available for reuse.
    pub available: usize,
    /// Bytes retained by idle reusable buffers.
    pub available_bytes: u64,
}

#[derive(Debug)]
pub(crate) struct UniformBufferPool {
    alignment: u64,
    state: Mutex<UniformBufferPoolState>,
    retirement: Shared<Retirement>,
}

#[derive(Debug, Default)]
struct UniformBufferPoolState {
    /// Oldest release first: capacity, buffer, bytes, creation generation.
    available: Vec<(u64, wgpu::Buffer, Vec<u8>, u64)>,
    /// Generation whose untrusted idle buffers were last dropped.
    checked_generation: u64,
    active_by_capacity: HashMap<u64, usize>,
    peak_active_by_capacity: HashMap<u64, usize>,
    created: u64,
    reused: u64,
    discarded: u64,
    active: usize,
    active_bytes: u64,
    high_water_bytes: u64,
    available_bytes: u64,
}

#[derive(Debug)]
pub(crate) struct UniformArenaLease {
    pool: Shared<UniformBufferPool>,
    buffer: Option<wgpu::Buffer>,
    capacity: u64,
    alignment: u64,
    cursor: u64,
    bytes: Vec<u8>,
    generation: u64,
}

#[derive(Copy, Clone, Debug)]
pub(crate) struct UniformSlice {
    offset: u64,
    size: NonZeroU64,
}

impl UniformSlice {
    /// Offset in the arena, a multiple of the device's uniform offset alignment.
    pub(crate) const fn offset(self) -> u64 {
        self.offset
    }
}

impl UniformBufferPool {
    pub(crate) fn new(alignment: u64, retirement: Shared<Retirement>) -> Self {
        Self {
            alignment: alignment.max(1),
            state: Mutex::default(),
            retirement,
        }
    }

    /// Leases an arena of at least `capacity` bytes, rounded up to a power of two so similar graphs
    /// share arenas.
    pub(crate) fn acquire(
        self: &Shared<Self>,
        device: &wgpu::Device,
        capacity: u64,
    ) -> UniformArenaLease {
        let alignment = self.alignment;
        let capacity = align_up(capacity.max(1), alignment)
            .checked_next_power_of_two()
            .unwrap_or(u64::MAX);
        let reused = {
            let mut state = self
                .state
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            let untrusted = self.drop_untrusted(&mut state);
            let index = state
                .available
                .iter()
                .position(|(available_capacity, _, _, _)| *available_capacity == capacity);
            let reused = index.map(|index| {
                state.reused = state.reused.saturating_add(1);
                state.available_bytes = state.available_bytes.saturating_sub(capacity);
                let (_, buffer, bytes, generation) = state.available.remove(index);
                (buffer, bytes, generation)
            });
            drop(state);
            self.retirement.retire(untrusted);
            reused
        };
        let (buffer, bytes, generation) = reused.unwrap_or_else(|| {
            let mut state = self
                .state
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            state.created = state.created.saturating_add(1);
            drop(state);
            let generation = self.retirement.generation();
            (
                device.create_buffer(&wgpu::BufferDescriptor {
                    label: Some("alpha2 uniform arena"),
                    size: capacity,
                    usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                }),
                vec![0; usize::try_from(capacity).expect("uniform arena fits host indexing")],
                generation,
            )
        });
        {
            let mut state = self
                .state
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            state.active = state.active.saturating_add(1);
            state.active_bytes = state.active_bytes.saturating_add(capacity);
            state.high_water_bytes = state.high_water_bytes.max(state.active_bytes);
            let active = state.active_by_capacity.entry(capacity).or_default();
            *active = active.saturating_add(1);
            let active = *active;
            let peak = state.peak_active_by_capacity.entry(capacity).or_default();
            *peak = (*peak).max(active);
        }
        UniformArenaLease {
            pool: Shared::clone(self),
            buffer: Some(buffer),
            capacity,
            alignment,
            cursor: 0,
            bytes,
            generation,
        }
    }

    /// Removes idle arenas created before the last reported GPU error, once per report.
    fn drop_untrusted(&self, state: &mut UniformBufferPoolState) -> Vec<(Retired, u64)> {
        let generation = self.retirement.generation();
        if state.checked_generation == generation {
            return Vec::new();
        }
        state.checked_generation = generation;
        let (kept, untrusted): (Vec<_>, Vec<_>) = std::mem::take(&mut state.available)
            .into_iter()
            .partition(|(_, _, _, created)| *created == generation);
        state.available = kept;
        state.available_bytes = state
            .available
            .iter()
            .map(|(capacity, _, _, _)| *capacity)
            .sum();
        state.discarded = state
            .discarded
            .saturating_add(u64::try_from(untrusted.len()).unwrap_or(u64::MAX));
        untrusted
            .into_iter()
            .map(|(capacity, buffer, _, _)| {
                prune_capacity(state, capacity);
                (Retired::Buffer(buffer), capacity)
            })
            .collect()
    }

    /// Destroys every idle uniform buffer.
    pub(crate) fn clear(&self) {
        let mut state = self
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let discarded = u64::try_from(state.available.len()).unwrap_or(u64::MAX);
        let idle = std::mem::take(&mut state.available);
        state.discarded = state.discarded.saturating_add(discarded);
        state.available_bytes = 0;
        drop(state);
        self.retirement.retire(
            idle.into_iter()
                .map(|(capacity, buffer, _, _)| (Retired::Buffer(buffer), capacity)),
        );
    }

    pub(crate) fn stats(&self) -> UniformBufferPoolStats {
        let state = self
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        UniformBufferPoolStats {
            created: state.created,
            reused: state.reused,
            discarded: state.discarded,
            active: state.active,
            active_bytes: state.active_bytes,
            high_water_bytes: state.high_water_bytes,
            available: state.available.len(),
            available_bytes: state.available_bytes,
        }
    }

    /// Keeps a finished arena while its capacity has fewer idle arenas than its peak in use, within
    /// the idle bounds (oldest evicted first); otherwise destroys it.
    fn release(&self, capacity: u64, buffer: wgpu::Buffer, bytes: Vec<u8>, generation: u64) {
        if !self.retirement.trusts(generation) {
            self.discard(capacity, buffer);
            return;
        }
        let mut retired = Vec::new();
        self.keep_or_retire(capacity, buffer, bytes, generation, &mut retired);
        self.retirement.retire(retired);
    }

    fn keep_or_retire(
        &self,
        capacity: u64,
        buffer: wgpu::Buffer,
        bytes: Vec<u8>,
        generation: u64,
        retired: &mut Vec<(Retired, u64)>,
    ) {
        let mut state = self
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        finish_active(&mut state, capacity);
        let retention_target = state
            .peak_active_by_capacity
            .get(&capacity)
            .copied()
            .unwrap_or(1);
        let idle_at_capacity = state
            .available
            .iter()
            .filter(|(available, _, _, _)| *available == capacity)
            .count();
        if idle_at_capacity >= retention_target || capacity > MAX_IDLE_UNIFORM_BYTES {
            state.discarded = state.discarded.saturating_add(1);
            retired.push((Retired::Buffer(buffer), capacity));
            return;
        }
        while state.available.len() >= MAX_IDLE_UNIFORM_BUFFERS
            || state.available_bytes.saturating_add(capacity) > MAX_IDLE_UNIFORM_BYTES
        {
            let (discarded, discarded_buffer, _bytes, _) = state.available.remove(0);
            state.available_bytes = state.available_bytes.saturating_sub(discarded);
            state.discarded = state.discarded.saturating_add(1);
            prune_capacity(&mut state, discarded);
            retired.push((Retired::Buffer(discarded_buffer), discarded));
        }
        state.available.push((capacity, buffer, bytes, generation));
        state.available_bytes = state.available_bytes.saturating_add(capacity);
    }

    /// Takes back an arena whose encode failed or was abandoned.
    fn discard(&self, capacity: u64, buffer: wgpu::Buffer) {
        let mut state = self
            .state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        finish_active(&mut state, capacity);
        state.discarded = state.discarded.saturating_add(1);
        prune_capacity(&mut state, capacity);
        drop(state);
        self.retirement
            .retire([(Retired::Buffer(buffer), capacity)]);
    }
}

impl UniformArenaLease {
    pub(crate) fn write(&mut self, contents: &[u8]) -> UniformSlice {
        let offset = align_up(self.cursor, self.alignment);
        let size = u64::try_from(contents.len()).expect("uniform contents fit u64");
        let end = offset.checked_add(size).expect("uniform arena offset");
        assert!(end <= self.capacity, "uniform arena capacity was reserved");
        let start = usize::try_from(offset).expect("uniform offset fits host indexing");
        let end = usize::try_from(end).expect("uniform end fits host indexing");
        self.bytes[start..end].copy_from_slice(contents);
        self.cursor = u64::try_from(end).expect("uniform cursor fits u64");
        UniformSlice {
            offset,
            size: NonZeroU64::new(size).expect("uniform binding is non-empty"),
        }
    }

    /// The arena's first `size` bytes as a dynamic-offset binding.
    pub(crate) fn dynamic_binding(&self, size: NonZeroU64) -> wgpu::BindingResource<'_> {
        wgpu::BindingResource::Buffer(wgpu::BufferBinding {
            buffer: self.buffer.as_ref().expect("live uniform arena"),
            offset: 0,
            size: Some(size),
        })
    }

    pub(crate) fn binding(&self, slice: UniformSlice) -> wgpu::BindingResource<'_> {
        wgpu::BindingResource::Buffer(wgpu::BufferBinding {
            buffer: self.buffer.as_ref().expect("live uniform arena"),
            offset: slice.offset,
            size: Some(slice.size),
        })
    }

    #[cfg(test)]
    pub(crate) fn buffer(&self) -> &wgpu::Buffer {
        self.buffer.as_ref().expect("live uniform arena")
    }

    pub(crate) fn flush(&self, queue: &wgpu::Queue) {
        if self.cursor > 0 {
            let used = usize::try_from(self.cursor).expect("uniform cursor fits host indexing");
            queue.write_buffer(
                self.buffer.as_ref().expect("live uniform arena"),
                0,
                &self.bytes[..used],
            );
        }
    }

    pub(crate) fn recycle(mut self) {
        let buffer = self.buffer.take().expect("live uniform arena");
        let bytes = std::mem::take(&mut self.bytes);
        self.pool
            .release(self.capacity, buffer, bytes, self.generation);
    }
}

impl Drop for UniformArenaLease {
    fn drop(&mut self) {
        if let Some(buffer) = self.buffer.take() {
            self.pool.discard(self.capacity, buffer);
        }
    }
}

fn finish_active(state: &mut UniformBufferPoolState, capacity: u64) {
    state.active = state.active.saturating_sub(1);
    state.active_bytes = state.active_bytes.saturating_sub(capacity);
    if let Some(active) = state.active_by_capacity.get_mut(&capacity) {
        *active = active.saturating_sub(1);
        if *active == 0 {
            state.active_by_capacity.remove(&capacity);
        }
    }
}

/// Forgets a capacity's peak once nothing of it is active or idle.
fn prune_capacity(state: &mut UniformBufferPoolState, capacity: u64) {
    if !state.active_by_capacity.contains_key(&capacity)
        && !state
            .available
            .iter()
            .any(|(available, _, _, _)| *available == capacity)
    {
        state.peak_active_by_capacity.remove(&capacity);
    }
}

pub(crate) const fn aligned_uniform_end(cursor: u64, size: u64, alignment: u64) -> u64 {
    align_up(cursor, alignment).saturating_add(size)
}

const fn align_up(value: u64, alignment: u64) -> u64 {
    value
        .saturating_add(alignment.saturating_sub(1))
        .div_euclid(alignment)
        .saturating_mul(alignment)
}