mirror of
https://github.com/storytold/effectcraft.git
synced 2026-11-12 00:42:36 +00:00
404 lines
18 KiB
Rust
404 lines
18 KiB
Rust
//! Export frames converted on the device to the H.264 encoder's input (8-bit BT.709 limited
|
||
//! range 4:2:0, see `shaders/yuv.wgsl`): a 4K frame reads back 12 MB of planes instead of
|
||
//! 133 MB of RGBA f32, and the CPU skips `to_rgba8_over` + `rgba_to_yuv420`.
|
||
|
||
use wgpu::util::DeviceExt;
|
||
|
||
use crate::context::{Enc, GpuImage};
|
||
|
||
pub(crate) struct YuvPipe {
|
||
bgl: wgpu::BindGroupLayout,
|
||
pipe: wgpu::ComputePipeline,
|
||
}
|
||
|
||
impl YuvPipe {
|
||
pub(crate) fn new(device: &wgpu::Device) -> Result<YuvPipe, String> {
|
||
crate::context::init_resource(device, "yuv420", || {
|
||
let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
|
||
label: Some("yuv420"),
|
||
source: wgpu::ShaderSource::Wgsl(include_str!("shaders/yuv.wgsl").into()),
|
||
});
|
||
let storage =
|
||
wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only: false }, has_dynamic_offset: false, min_binding_size: None };
|
||
let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
|
||
label: Some("yuv420"),
|
||
entries: &[
|
||
wgpu::BindGroupLayoutEntry {
|
||
binding: 0,
|
||
visibility: wgpu::ShaderStages::COMPUTE,
|
||
ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Uniform, has_dynamic_offset: false, min_binding_size: None },
|
||
count: None,
|
||
},
|
||
wgpu::BindGroupLayoutEntry {
|
||
binding: 1,
|
||
visibility: wgpu::ShaderStages::COMPUTE,
|
||
ty: wgpu::BindingType::Texture {
|
||
sample_type: wgpu::TextureSampleType::Float { filterable: false },
|
||
view_dimension: wgpu::TextureViewDimension::D2,
|
||
multisampled: false,
|
||
},
|
||
count: None,
|
||
},
|
||
wgpu::BindGroupLayoutEntry { binding: 2, visibility: wgpu::ShaderStages::COMPUTE, ty: storage, count: None },
|
||
],
|
||
});
|
||
let layout =
|
||
device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("yuv420"), bind_group_layouts: &[Some(&bgl)], immediate_size: 0 });
|
||
let pipe = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
|
||
label: Some("yuv420"),
|
||
layout: Some(&layout),
|
||
module: &module,
|
||
entry_point: Some("yuv420"),
|
||
compilation_options: Default::default(),
|
||
cache: None,
|
||
});
|
||
YuvPipe { bgl, pipe }
|
||
})
|
||
}
|
||
}
|
||
|
||
/// The frame sizes the kernel handles: width a multiple of 8 (4 chroma samples per
|
||
/// invocation), height even.
|
||
pub(crate) fn supported(w: u32, h: u32) -> bool {
|
||
w > 0 && h > 0 && w.is_multiple_of(8) && h.is_multiple_of(2)
|
||
}
|
||
|
||
/// Bytes of the planes: Y `w×h`, then U and V `w/2 × h/2`.
|
||
pub(crate) fn planes_len(w: u32, h: u32) -> u64 {
|
||
w as u64 * h as u64 + 2 * ((w / 2) as u64 * (h / 2) as u64)
|
||
}
|
||
|
||
/// Whether the [`planes_len`] storage and staging buffers for `w×h` fit within `limits`.
|
||
pub(crate) fn fits_limits(w: u32, h: u32, limits: &wgpu::Limits) -> bool {
|
||
let size = planes_len(w, h);
|
||
size <= limits.max_storage_buffer_binding_size && size <= limits.max_buffer_size
|
||
}
|
||
|
||
/// Record the conversion of `img` over `bg` into a new `MAP_READ` buffer of
|
||
/// [`planes_len`] bytes and submit it. `None`: the size is not [`supported`], exceeds the
|
||
/// device's buffer limits, or buffer allocation failed.
|
||
pub(crate) fn record(e: &mut Enc, yp: &YuvPipe, img: &GpuImage, bg: [f32; 3]) -> Option<wgpu::Buffer> {
|
||
let (w, h) = (img.width, img.height);
|
||
let g = e.g;
|
||
if !supported(w, h) || !fits_limits(w, h, &g.device.limits()) {
|
||
return None;
|
||
}
|
||
let cw = w / 2;
|
||
let size = planes_len(w, h);
|
||
let mut params = [0u8; 32];
|
||
// w * h fits in u32: both are at most the texture limit (16384).
|
||
for (i, v) in [w, h, cw, w * h / 4].into_iter().enumerate() {
|
||
params[i * 4..i * 4 + 4].copy_from_slice(&v.to_le_bytes());
|
||
}
|
||
for (i, v) in [bg[0], bg[1], bg[2], 1.0].into_iter().enumerate() {
|
||
params[16 + i * 4..16 + i * 4 + 4].copy_from_slice(&v.to_le_bytes());
|
||
}
|
||
let (out, stg, ub) = match crate::context::init_resource(&g.device, "yuv420 buffers", || {
|
||
let out = g.device.create_buffer(&wgpu::BufferDescriptor {
|
||
label: Some("yuv420 out"),
|
||
size,
|
||
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
|
||
mapped_at_creation: false,
|
||
});
|
||
let stg = g.device.create_buffer(&wgpu::BufferDescriptor {
|
||
label: Some("yuv420 staging"),
|
||
size,
|
||
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
|
||
mapped_at_creation: false,
|
||
});
|
||
let ub = g.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||
label: Some("yuv420 params"),
|
||
contents: ¶ms,
|
||
usage: wgpu::BufferUsages::UNIFORM,
|
||
});
|
||
(out, stg, ub)
|
||
}) {
|
||
Ok(buffers) => buffers,
|
||
Err(err) => {
|
||
log::warn!(target: "effectcraft_gpu", "yuv420 buffers: {err}");
|
||
return None;
|
||
}
|
||
};
|
||
let view = img.texture.create_view(&Default::default());
|
||
let group = g.device.create_bind_group(&wgpu::BindGroupDescriptor {
|
||
label: Some("yuv420"),
|
||
layout: &yp.bgl,
|
||
entries: &[
|
||
wgpu::BindGroupEntry { binding: 0, resource: ub.as_entire_binding() },
|
||
wgpu::BindGroupEntry { binding: 1, resource: wgpu::BindingResource::TextureView(&view) },
|
||
wgpu::BindGroupEntry { binding: 2, resource: out.as_entire_binding() },
|
||
],
|
||
});
|
||
{
|
||
let enc = e.record(&[img]);
|
||
{
|
||
let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("yuv420"), timestamp_writes: None });
|
||
pass.set_pipeline(&yp.pipe);
|
||
pass.set_bind_group(0, &group, &[]);
|
||
pass.dispatch_workgroups((w / 8).div_ceil(8), (h / 2).div_ceil(8), 1);
|
||
}
|
||
enc.copy_buffer_to_buffer(&out, 0, &stg, 0, size);
|
||
}
|
||
e.submit();
|
||
Some(stg)
|
||
}
|
||
|
||
/// Convert `img` over `bg` and read the planes back, waiting for the device.
|
||
pub(crate) fn convert(e: &mut Enc, yp: &YuvPipe, img: &GpuImage, bg: [f32; 3]) -> Option<Vec<u8>> {
|
||
let stg = record(e, yp, img, bg)?;
|
||
let n = planes_len(img.width, img.height) as usize;
|
||
e.g.read_buffer_after(&stg, e.last_submission(), n, n, 1).map_err(|err| log::error!(target: "effectcraft_gpu", "yuv420 readback: {err}")).ok()
|
||
}
|
||
|
||
#[cfg(test)]
|
||
pub(crate) mod tests {
|
||
use effectcraft_raster::Image;
|
||
|
||
/// The export's conversion (`to_rgba8_over` + `rgba_to_yuv420`), the reference the kernel
|
||
/// is held to.
|
||
pub(crate) fn cpu(img: &Image, bg: [f32; 3]) -> Vec<u8> {
|
||
effectcraft_raster::yuv::yuv420_over(img, bg)
|
||
}
|
||
|
||
/// Premultiplied noise with some out-of-range values, opaque pixels, exact 8-bit levels and
|
||
/// a few NaN and infinite components (the CPU writes 0 for NaN).
|
||
pub(crate) fn noise(w: u32, h: u32, seed: u32) -> Image {
|
||
let mut img = Image::new(w, h);
|
||
let mut s = seed | 1;
|
||
let mut rnd = || {
|
||
s ^= s << 13;
|
||
s ^= s >> 17;
|
||
s ^= s << 5;
|
||
s as f32 / u32::MAX as f32
|
||
};
|
||
for (i, p) in img.data.iter_mut().enumerate() {
|
||
let a = if i % 7 == 0 { 1.0 } else { rnd() };
|
||
*p = if i % 97 == 3 {
|
||
[f32::NAN, 0.5, f32::INFINITY, a]
|
||
} else if i % 101 == 5 {
|
||
[0.2, 0.3, 0.4, f32::NAN]
|
||
} else if i % 103 == 7 {
|
||
[f32::NEG_INFINITY, 0.1, 0.9, f32::INFINITY]
|
||
} else if i % 5 == 0 {
|
||
let q = (i % 256) as f32 / 255.0;
|
||
[q, 1.0 - q, q * 0.5, 1.0]
|
||
} else {
|
||
[rnd() * a * 1.1 - 0.02, rnd() * a, rnd() * a, a]
|
||
};
|
||
}
|
||
img
|
||
}
|
||
|
||
/// Check `got` against the CPU's `want`. Software adapters (llvmpipe, WARP) must match
|
||
/// exactly. A hardware driver may contract multiply-adds or divide less exactly, which
|
||
/// moves a value near a rounding boundary by one level (seen on an RTX 3090 under Vulkan):
|
||
/// there every byte must be within one level and at most `budget` of them off (about 2–3×
|
||
/// the rate measured on the 3090). A wrong formula, rounding mode or NaN rule fails either
|
||
/// way on software adapters.
|
||
pub(crate) fn assert_parity(g: &crate::Gpu, got: &[u8], want: &[u8], budget: f64, what: &str) {
|
||
assert_eq!(got.len(), want.len(), "{what}: length");
|
||
let diff: Vec<usize> = (0..got.len()).filter(|&i| got[i] != want[i]).collect();
|
||
let max = diff.iter().map(|&i| (got[i] as i32 - want[i] as i32).abs()).max().unwrap_or(0);
|
||
eprintln!("{what} on {}: {} of {} bytes differ (max {max})", g.ctx.name, diff.len(), got.len());
|
||
if g.ctx.software {
|
||
assert!(diff.is_empty(), "{what}: a software adapter must match exactly, first diffs at {:?}", &diff[..diff.len().min(10)]);
|
||
return;
|
||
}
|
||
assert!(max <= 1, "{what}: a byte is {max} levels off, first diffs at {:?}", &diff[..diff.len().min(10)]);
|
||
assert!(diff.len() as f64 <= budget * got.len() as f64, "{what}: {} of {} bytes off by one", diff.len(), got.len());
|
||
}
|
||
|
||
#[test]
|
||
fn yuv420_matches_cpu() {
|
||
let Some(g) = crate::tests::gpu() else { return };
|
||
let yp = super::YuvPipe::new(&g.ctx.device).expect("yuv420 pipe");
|
||
for (w, h, bg) in [(264u32, 150u32, [0.05, 0.1, 0.2]), (8, 2, [1.0, 1.0, 1.0]), (1928, 6, [0.0, 0.0, 0.0])] {
|
||
let img = noise(w, h, w ^ (h << 16));
|
||
let src = g.ctx.upload_image(&img).expect("upload");
|
||
let got = super::convert(&mut crate::context::Enc::new(&g.ctx), &yp, &src, bg).expect("readback");
|
||
assert_parity(g, &got, &cpu(&img, bg), 0.00005, &format!("noise {w}x{h}"));
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn unsupported_sizes_are_declined() {
|
||
assert!(super::supported(8, 2));
|
||
for (w, h) in [(0, 2), (8, 0), (12, 2), (8, 3), (7, 4)] {
|
||
assert!(!super::supported(w, h), "{w}x{h}");
|
||
}
|
||
assert_eq!(super::planes_len(8, 2), 16 + 2 * 4);
|
||
}
|
||
|
||
#[test]
|
||
fn sizes_exceeding_buffer_limits_are_declined() {
|
||
let def = wgpu::Limits::default();
|
||
assert!(super::fits_limits(3840, 2160, &def));
|
||
// 16384×16384 planes are 384 MiB, over default max_storage_buffer_binding_size (128 MiB)
|
||
// and max_buffer_size (256 MiB).
|
||
assert!(!super::fits_limits(16384, 16384, &def));
|
||
let size = super::planes_len(64, 64);
|
||
assert!(super::fits_limits(64, 64, &wgpu::Limits { max_storage_buffer_binding_size: size, max_buffer_size: size, ..def.clone() }));
|
||
assert!(!super::fits_limits(64, 64, &wgpu::Limits { max_storage_buffer_binding_size: size - 1, max_buffer_size: size, ..def.clone() }));
|
||
assert!(!super::fits_limits(64, 64, &wgpu::Limits { max_storage_buffer_binding_size: size, max_buffer_size: size - 1, ..def.clone() }));
|
||
|
||
let Some(_) = crate::tests::gpu() else { return };
|
||
let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
|
||
let adapter = pollster::block_on(instance.request_adapter(&Default::default())).expect("adapter");
|
||
let required_limits = wgpu::Limits { max_storage_buffer_binding_size: 4096, max_buffer_size: 4096, ..def };
|
||
let (device, queue) = pollster::block_on(adapter.request_device(&crate::context::device_descriptor(required_limits))).expect("limited device");
|
||
let g = crate::Gpu::new(&adapter, device, queue).expect("limited Gpu");
|
||
let yp = super::YuvPipe::new(&g.ctx.device).expect("yuv420 pipe");
|
||
let src = g.ctx.upload_image(&noise(64, 64, 1)).expect("upload");
|
||
assert!(super::convert(&mut crate::context::Enc::new(&g.ctx), &yp, &src, [0.0; 3]).is_none());
|
||
assert!(g.ctx.check_health().is_ok());
|
||
}
|
||
|
||
#[test]
|
||
fn unsupported_sizes_are_not_converted() {
|
||
let Some(g) = crate::tests::gpu() else { return };
|
||
let yp = super::YuvPipe::new(&g.ctx.device).expect("yuv420 pipe");
|
||
let src = g.ctx.upload_image(&noise(12, 4, 3)).expect("upload");
|
||
assert!(super::convert(&mut crate::context::Enc::new(&g.ctx), &yp, &src, [0.0; 3]).is_none());
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod accel_tests {
|
||
use effectcraft_color::Label;
|
||
use effectcraft_project::{BitDepth, Comp, ItemKind};
|
||
use effectcraft_render::{Accelerator, Backend, RenderOpts, Renderer};
|
||
use effectcraft_time::{FrameRate, Tick};
|
||
|
||
use crate::tests::{Pattern, Scene, opts};
|
||
|
||
fn scene(w: u32, h: u32) -> Scene {
|
||
let mut s = Scene::new(BitDepth::Bpc8);
|
||
let comp = Comp::new(w, h, FrameRate::FPS_30, Tick::from_seconds_f64(2.0));
|
||
s.cid = s.p.add_item("Comp yuv", Label::Sandstone, None, ItemKind::Comp(comp.clone().into()));
|
||
s.comp = comp;
|
||
let l = s.footage(w + 7, h + 3);
|
||
s.push(l);
|
||
let sol = s.solid([0.9, 0.2, 0.4], w / 2, h / 3);
|
||
s.push(sol);
|
||
s
|
||
}
|
||
|
||
/// The accelerator's YUV frame is the export's CPU conversion of the very frame
|
||
/// `comp_frame` returns.
|
||
#[test]
|
||
fn comp_frame_yuv420_is_the_cpu_conversion_of_comp_frame() {
|
||
let Some(g) = crate::tests::gpu() else { return };
|
||
let s = scene(96, 60);
|
||
let mut r = Renderer::new(&s.p, &Pattern, RenderOpts { backend: Backend::Gpu, ..opts() });
|
||
r.accel = Some(g);
|
||
let bg = [0.2, 0.3, 0.4];
|
||
let rgba = g.comp_frame(&r, s.cid, Tick::ZERO).expect("rgba frame");
|
||
let yuv = g.comp_frame_yuv420(&r, s.cid, Tick::ZERO, bg).expect("yuv frame");
|
||
super::tests::assert_parity(g, &yuv, &super::tests::cpu(&rgba, bg), 0.006, "comp frame");
|
||
}
|
||
|
||
#[test]
|
||
fn odd_sized_comps_fall_back_to_the_cpu() {
|
||
let Some(g) = crate::tests::gpu() else { return };
|
||
let s = scene(97, 61);
|
||
let mut r = Renderer::new(&s.p, &Pattern, RenderOpts { backend: Backend::Gpu, ..opts() });
|
||
r.accel = Some(g);
|
||
assert!(g.comp_frame(&r, s.cid, Tick::ZERO).is_some());
|
||
assert!(g.comp_frame_yuv420(&r, s.cid, Tick::ZERO, [0.0; 3]).is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn comp_frame_yuv420_falls_back_when_pipeline_init_failed() {
|
||
let Some(_) = crate::tests::gpu() else { return };
|
||
let g = crate::Gpu::try_headless().expect("headless Gpu");
|
||
let _ = g.ctx.yuv.set(Err("simulated yuv420 init failure".into()));
|
||
let s = scene(96, 60);
|
||
let mut r = Renderer::new(&s.p, &Pattern, RenderOpts { backend: Backend::Gpu, ..opts() });
|
||
r.accel = Some(&g);
|
||
assert!(g.comp_frame_yuv420(&r, s.cid, Tick::ZERO, [0.0; 3]).is_none());
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tie_tests {
|
||
use effectcraft_raster::Image;
|
||
|
||
fn lum(r: u8, g: u8, b: u8) -> f32 {
|
||
let (r, g, b) = (r as f32 / 255.0, g as f32 / 255.0, b as f32 / 255.0);
|
||
0.2126 * r + 0.7152 * g + 0.0722 * b
|
||
}
|
||
|
||
fn px(p: [u8; 3]) -> [f32; 4] {
|
||
[p[0] as f32 / 255.0, p[1] as f32 / 255.0, p[2] as f32 / 255.0, 1.0]
|
||
}
|
||
|
||
/// Opaque pixels whose Y (or whose 2×2 block's U or V) lands exactly on .5 in the CPU's f32
|
||
/// math: a kernel that rounds ties the other way moves most of them.
|
||
#[test]
|
||
fn rounding_ties_match_cpu() {
|
||
let mut ties: Vec<[u8; 3]> = vec![];
|
||
for r in 0..=255u8 {
|
||
for g in 0..=255u8 {
|
||
for b in 0..=255u8 {
|
||
let x = 16.0 + 219.0 * lum(r, g, b);
|
||
if x - x.floor() == 0.5 {
|
||
ties.push([r, g, b]);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
assert!(!ties.is_empty());
|
||
let mut blocks: Vec<[[u8; 3]; 4]> = vec![];
|
||
let mut s = 12345u32;
|
||
let mut rnd = || {
|
||
s ^= s << 13;
|
||
s ^= s >> 17;
|
||
s ^= s << 5;
|
||
(s >> 24) as u8
|
||
};
|
||
// Bounded so the search stays quick in debug builds (a few hundred blocks).
|
||
for _ in 0..8_000_000 {
|
||
let blk = [[rnd(), rnd(), rnd()], [rnd(), rnd(), rnd()], [rnd(), rnd(), rnd()], [rnd(), rnd(), rnd()]];
|
||
let (mut a0, mut a1) = (0f32, 0f32);
|
||
for p in blk {
|
||
let l = lum(p[0], p[1], p[2]);
|
||
a0 += (p[2] as f32 / 255.0 - l) / 1.8556;
|
||
a1 += (p[0] as f32 / 255.0 - l) / 1.5748;
|
||
}
|
||
let u = 128.0 + 224.0 * a0 / 4.0;
|
||
let v = 128.0 + 224.0 * a1 / 4.0;
|
||
if u - u.floor() == 0.5 || v - v.floor() == 0.5 {
|
||
blocks.push(blk);
|
||
}
|
||
}
|
||
eprintln!("Y tie pixels: {}, chroma tie blocks: {}", ties.len(), blocks.len());
|
||
assert!(!blocks.is_empty());
|
||
let Some(g) = crate::tests::gpu() else { return };
|
||
let yp = super::YuvPipe::new(&g.ctx.device).expect("yuv420 pipe");
|
||
let w = 1024usize;
|
||
let h = (ties.len().div_ceil(w) + 1) & !1;
|
||
let mut img = Image::new(w as u32, h as u32);
|
||
for (i, p) in img.data.iter_mut().enumerate() {
|
||
*p = px(ties[i % ties.len()]);
|
||
}
|
||
let src = g.ctx.upload_image(&img).expect("upload");
|
||
let got = super::convert(&mut crate::context::Enc::new(&g.ctx), &yp, &src, [0.0; 3]).expect("readback");
|
||
super::tests::assert_parity(g, &got, &super::tests::cpu(&img, [0.0; 3]), 0.02, "Y ties");
|
||
let cw = 512usize;
|
||
let rows = blocks.len().div_ceil(cw);
|
||
let mut img = Image::new((cw * 2) as u32, (rows * 2) as u32);
|
||
for cy in 0..rows {
|
||
for cx in 0..cw {
|
||
let b = blocks[(cy * cw + cx) % blocks.len()];
|
||
for (k, p) in b.iter().enumerate() {
|
||
let (x, y) = (cx * 2 + k % 2, cy * 2 + k / 2);
|
||
img.data[y * cw * 2 + x] = px(*p);
|
||
}
|
||
}
|
||
}
|
||
let src = g.ctx.upload_image(&img).expect("upload");
|
||
let got = super::convert(&mut crate::context::Enc::new(&g.ctx), &yp, &src, [0.0; 3]).expect("readback");
|
||
super::tests::assert_parity(g, &got, &super::tests::cpu(&img, [0.0; 3]), 0.06, "chroma ties");
|
||
}
|
||
}
|