#![cfg_attr(
not(any(
feature = "vulkan",
all(feature = "dx12", target_os = "windows"),
all(feature = "metal", target_os = "macos"),
)),
allow(dead_code)
)]
use crate::slang::OwnedLayoutCheck;
use crate::types::{DepthStencilState, OptimizationLevel, PrimitiveTopology, TextureFormat, VertexBufferLayout};
use anyhow::Result;
use super::DeviceHandle;
#[cfg(any(
all(feature = "metal", target_os = "macos"),
all(feature = "dx12", target_os = "windows"),
))]
use super::ShaderHandle;
pub const MAX_BINDLESS_SLOTS: usize = 16;
pub const MAX_USER_SLOTS: usize = 8;
pub const TOTAL_PUSH_BYTES: usize = 128;
pub const DISPATCH_BATCH_STRIDE: usize = TOTAL_PUSH_BYTES + 3 * 4;
#[repr(C)]
#[derive(Default, Clone, Copy, Debug)]
pub struct PushLayout {
pub bindless: [u16; MAX_BINDLESS_SLOTS],
pub user: [u32; MAX_USER_SLOTS],
pub _reserved: [u32; 16],
}
const _: () = assert!(std::mem::size_of::<PushLayout>() == TOTAL_PUSH_BYTES);
unsafe impl bytemuck::Pod for PushLayout {}
unsafe impl bytemuck::Zeroable for PushLayout {}
impl PushLayout {
#[inline]
pub fn as_bytes(&self) -> &[u8] {
bytemuck::bytes_of(self)
}
}
#[inline]
pub fn fill_frame_table_dispatch(layout: &mut PushLayout, dispatch_base: u32, user: &[u32]) {
layout._reserved[crate::frame_table::dispatch_table_base_word_index()] = dispatch_base;
for (i, &val) in user.iter().enumerate().take(MAX_USER_SLOTS) {
layout.user[i] = val;
}
}
pub const FRAME_TABLE_SELECTOR_SLOT_WORD: usize = 1;
pub const FRAME_TABLE_TABLE_SLOT_WORD: usize = 2;
#[inline]
pub fn set_frame_table_slots(layout: &mut PushLayout, selector_slot: u32, table_slot: u32) {
layout._reserved[FRAME_TABLE_SELECTOR_SLOT_WORD] = selector_slot;
layout._reserved[FRAME_TABLE_TABLE_SLOT_WORD] = table_slot;
}
pub fn patch_dispatch_batch_frame_table_slots(arg_data: &mut [u8], count: usize, selector_slot: u32, table_slot: u32) {
let sel_off = (MAX_BINDLESS_SLOTS * 2) + (MAX_USER_SLOTS * 4) + FRAME_TABLE_SELECTOR_SLOT_WORD * 4;
let tab_off = (MAX_BINDLESS_SLOTS * 2) + (MAX_USER_SLOTS * 4) + FRAME_TABLE_TABLE_SLOT_WORD * 4;
for i in 0..count {
let base = i * DISPATCH_BATCH_STRIDE;
if base + TOTAL_PUSH_BYTES > arg_data.len() {
break;
}
arg_data[base + sel_off..base + sel_off + 4].copy_from_slice(&selector_slot.to_ne_bytes());
arg_data[base + tab_off..base + tab_off + 4].copy_from_slice(&table_slot.to_ne_bytes());
}
}
#[derive(Debug, Clone)]
pub struct SlotAllocator {
next: u32,
free: Vec<u32>,
}
impl Default for SlotAllocator {
fn default() -> Self {
Self::new(0)
}
}
impl SlotAllocator {
pub fn new(start: u32) -> Self {
Self {
next: start,
free: Vec::new(),
}
}
#[inline]
pub fn alloc(&mut self) -> u32 {
self.free.pop().unwrap_or_else(|| {
let i = self.next;
self.next += 1;
i
})
}
#[inline]
pub fn free(&mut self, slot: u32) {
self.free.push(slot);
}
#[inline]
pub fn live_count(&self) -> u32 {
self.next - self.free.len() as u32
}
#[cfg(any(test, all(feature = "metal", target_os = "macos")))]
#[inline]
pub fn free_count(&self) -> usize {
self.free.len()
}
#[cfg(any(test, all(feature = "metal", target_os = "macos")))]
#[inline]
pub fn next_fresh(&self) -> u32 {
self.next
}
#[inline]
pub fn ensure_minimum_next(&mut self, min: u32) {
if self.next < min {
self.next = min;
}
}
}
pub struct DeferredQueue<K, V> {
pending: Vec<(K, V)>,
}
impl<K, V> Default for DeferredQueue<K, V> {
fn default() -> Self {
Self { pending: Vec::new() }
}
}
impl<K, V> DeferredQueue<K, V> {
pub fn new() -> Self {
Self { pending: Vec::new() }
}
#[inline]
pub fn push(&mut self, key: K, value: V) {
self.pending.push((key, value));
}
#[inline]
pub fn len(&self) -> usize {
self.pending.len()
}
}
impl<K: PartialOrd + Copy, V> DeferredQueue<K, V> {
pub fn drain_up_to(&mut self, threshold: K) -> Vec<V> {
self.drain_up_to_filtered(threshold, |_| true)
}
pub fn drain_up_to_filtered<F>(&mut self, threshold: K, can_take: F) -> Vec<V>
where
F: Fn(&V) -> bool,
{
if self.pending.is_empty() {
return Vec::new();
}
let mut i = 0;
let mut eligible: Vec<V> = Vec::new();
while i < self.pending.len() {
if self.pending[i].0 <= threshold && can_take(&self.pending[i].1) {
let (_, v) = self.pending.swap_remove(i);
eligible.push(v);
} else {
i += 1;
}
}
eligible
}
}
impl<K, V> DeferredQueue<K, V> {
pub fn flush_all(&mut self) -> impl Iterator<Item = V> + '_ {
self.pending.drain(..).map(|(_, v)| v)
}
pub fn drain_where<F: Fn(&K) -> bool>(&mut self, ready: F) -> Vec<V> {
self.drain_where_with_keys(ready).into_iter().map(|(_, v)| v).collect()
}
pub fn drain_where_with_keys<F: Fn(&K) -> bool>(&mut self, ready: F) -> Vec<(K, V)> {
if self.pending.is_empty() {
return Vec::new();
}
let mut i = 0;
let mut eligible: Vec<(K, V)> = Vec::new();
while i < self.pending.len() {
if ready(&self.pending[i].0) {
let (k, v) = self.pending.swap_remove(i);
eligible.push((k, v));
} else {
i += 1;
}
}
eligible
}
}
pub const STAGING_COPY_ALIGN: u64 = 256;
pub const DEFAULT_STAGING_CHUNK_SIZE: u64 = 256 * 1024;
pub trait BeltChunk: Sized {
fn capacity(&self) -> u64;
fn offset(&self) -> u64;
fn offset_mut(&mut self) -> &mut u64;
fn mapped_ptr(&self) -> *mut u8;
#[inline]
fn reset(&mut self) {
*self.offset_mut() = 0;
}
}
pub struct StagingBeltCore<C> {
pub free: Vec<C>,
pub active: Vec<C>,
pub in_flight: Vec<(u64, Vec<C>)>,
pub chunk_size: u64,
}
impl<C: BeltChunk> StagingBeltCore<C> {
pub fn new(chunk_size: u64) -> Self {
Self {
free: Vec::new(),
active: Vec::new(),
in_flight: Vec::new(),
chunk_size,
}
}
pub fn write(&mut self, data: &[u8], alloc: impl FnOnce(u64) -> Result<C>) -> Result<(usize, u64)> {
if data.is_empty() {
anyhow::bail!("StagingBeltCore::write: empty data");
}
let len = data.len() as u64;
if let Some(ch) = self.active.last_mut() {
let start = align_up(ch.offset(), STAGING_COPY_ALIGN);
if start + len <= ch.capacity() {
unsafe {
std::ptr::copy_nonoverlapping(data.as_ptr(), ch.mapped_ptr().add(start as usize), data.len());
}
*ch.offset_mut() = start + len;
return Ok((self.active.len() - 1, start));
}
}
let alloc_size = self.chunk_size.max(align_up(len, STAGING_COPY_ALIGN));
let mut chunk = if let Some(pos) = self.free.iter().rposition(|c| c.capacity() >= len) {
let mut c = self.free.swap_remove(pos);
c.reset();
c
} else {
alloc(alloc_size)?
};
debug_assert_eq!(chunk.offset(), 0);
let start = 0u64;
unsafe {
std::ptr::copy_nonoverlapping(data.as_ptr(), chunk.mapped_ptr().add(start as usize), data.len());
}
*chunk.offset_mut() = start + len;
self.active.push(chunk);
Ok((self.active.len() - 1, start))
}
pub fn finish(&mut self, token: u64) {
if self.active.is_empty() {
return;
}
self.in_flight.push((token, std::mem::take(&mut self.active)));
}
pub fn trim_free(&mut self, mut destroy: impl FnMut(C)) {
let limit = self.chunk_size;
let mut i = 0;
while i < self.free.len() {
if self.free[i].capacity() > limit {
destroy(self.free.swap_remove(i));
} else {
i += 1;
}
}
}
pub fn destroy_all(&mut self, mut destroy: impl FnMut(C)) {
for ch in self.free.drain(..) {
destroy(ch);
}
for ch in self.active.drain(..) {
destroy(ch);
}
for (_, chunks) in self.in_flight.drain(..) {
for ch in chunks {
destroy(ch);
}
}
}
}
#[inline]
pub fn align_up(x: u64, a: u64) -> u64 {
x.div_ceil(a) * a
}
#[cfg(test)]
mod tests {
use super::*;
struct MockChunk {
backing: Vec<u8>,
offset: u64,
}
impl MockChunk {
fn new(capacity: u64) -> Self {
Self {
backing: vec![0u8; capacity as usize],
offset: 0,
}
}
}
impl BeltChunk for MockChunk {
fn capacity(&self) -> u64 {
self.backing.len() as u64
}
fn offset(&self) -> u64 {
self.offset
}
fn offset_mut(&mut self) -> &mut u64 {
&mut self.offset
}
fn mapped_ptr(&self) -> *mut u8 {
self.backing.as_ptr() as *mut u8
}
}
fn new_core(chunk_size: u64) -> StagingBeltCore<MockChunk> {
StagingBeltCore::new(chunk_size)
}
#[test]
fn staging_belt_finish_isolates_in_flight_chunk() {
let mut core = new_core(256);
let payload_a = b"hello";
let (idx_a, start_a) = core.write(payload_a, |sz| Ok(MockChunk::new(sz))).unwrap();
assert_eq!(core.active.len(), 1, "one active chunk after first write");
assert_eq!(core.in_flight.len(), 0);
assert_eq!(
&core.active[idx_a].backing[start_a as usize..start_a as usize + payload_a.len()],
payload_a
);
core.finish(42);
assert_eq!(core.active.len(), 0, "active drained by finish");
assert_eq!(core.in_flight.len(), 1, "one in-flight batch");
assert_eq!(core.in_flight[0].0, 42, "token preserved");
assert_eq!(core.in_flight[0].1.len(), 1, "one chunk in batch");
let payload_b = b"world";
let (idx_b, _start_b) = core.write(payload_b, |sz| Ok(MockChunk::new(sz))).unwrap();
assert_eq!(core.active.len(), 1, "new active chunk opened");
assert_eq!(core.in_flight.len(), 1, "in-flight batch is untouched");
assert_eq!(idx_b, 0, "second write is in a fresh active chunk at index 0");
let inflight_chunk = &core.in_flight[0].1[0];
assert_eq!(
&inflight_chunk.backing[start_a as usize..start_a as usize + payload_a.len()],
payload_a
);
}
#[test]
fn staging_belt_reclaimed_chunk_is_reused() {
let mut core = new_core(256);
core.write(b"first", |sz| Ok(MockChunk::new(sz))).unwrap();
core.finish(1);
assert_eq!(core.in_flight.len(), 1);
let (token, mut chunks) = core.in_flight.remove(0);
assert_eq!(token, 1);
for ch in &mut chunks {
ch.reset();
}
core.free.extend(chunks);
assert_eq!(core.free.len(), 1, "chunk available for reuse");
let alloc_called = std::cell::Cell::new(false);
core.write(b"second", |sz| {
alloc_called.set(true);
Ok(MockChunk::new(sz))
})
.unwrap();
assert!(!alloc_called.get(), "free chunk must be reused, not a fresh allocation");
assert_eq!(core.free.len(), 0, "free list drained");
assert_eq!(core.active.len(), 1);
}
}
#[inline]
pub fn resolve_clear_size(buffer_size: u64, offset: u64, size: u64) -> u64 {
if size == 0 {
buffer_size.saturating_sub(offset)
} else {
size
}
}
#[derive(Debug)]
pub struct ShaderDesc<'a> {
pub device: DeviceHandle,
pub slang_source: &'a str,
pub search_paths: &'a [&'a str],
pub defines: &'a [(&'a str, &'a str)],
pub optimization_level: OptimizationLevel,
pub layout_checks: Vec<OwnedLayoutCheck>,
}
impl<'a> ShaderDesc<'a> {
#[inline]
pub fn new(
device: DeviceHandle,
slang_source: &'a str,
search_paths: &'a [&'a str],
defines: &'a [(&'a str, &'a str)],
optimization_level: OptimizationLevel,
) -> Self {
Self {
device,
slang_source,
search_paths,
defines,
optimization_level,
layout_checks: Vec::new(),
}
}
#[inline]
pub fn with_layout_checks(mut self, layout_checks: Vec<OwnedLayoutCheck>) -> Self {
self.layout_checks = layout_checks;
self
}
}
#[cfg(all(feature = "metal", target_os = "macos"))]
#[derive(Debug, Clone, Copy)]
pub struct ShaderStageCompileDesc<'a> {
pub slang_source: &'a str,
pub search_paths: &'a [&'a str],
pub entry_point: &'a str,
pub stage: crate::slang::SlangStage,
pub extra_defines: &'a [(&'a str, &'a str)],
pub layout_checks: &'a [OwnedLayoutCheck],
pub optimization_level: OptimizationLevel,
}
#[derive(Debug, Clone, Copy)]
pub struct PipelineDesc<'a> {
pub vertex_layout: &'a VertexBufferLayout,
pub topology: PrimitiveTopology,
pub target_format: TextureFormat,
pub depth_stencil: Option<&'a DepthStencilState>,
}
impl<'a> PipelineDesc<'a> {
#[inline]
pub fn new(
vertex_layout: &'a VertexBufferLayout,
topology: PrimitiveTopology,
target_format: TextureFormat,
) -> Self {
Self {
vertex_layout,
topology,
target_format,
depth_stencil: None,
}
}
#[inline]
pub fn with_depth_stencil(mut self, depth_stencil: Option<&'a DepthStencilState>) -> Self {
self.depth_stencil = depth_stencil;
self
}
}
#[cfg(any(
all(feature = "metal", target_os = "macos"),
all(feature = "dx12", target_os = "windows"),
))]
#[derive(Debug, Clone, Copy)]
pub struct GraphicsPipelineCreateDesc<'a> {
pub device_handle: DeviceHandle,
pub vertex_shader: ShaderHandle,
pub fragment_shader: ShaderHandle,
pub raster: &'a PipelineDesc<'a>,
}