Convert export frames RGBA->YUV420p on the GPU
The export read back 8MB RGBA per frame and ran swscale RGBA->YUV420p on the UI thread (~6ms/frame). Add a tight GPU compute converter (gpu_yuv, BT.709 full-range matching the encoder tags) and wire it into the triple-buffered ReadbackPipeline: render to RGBA, convert on the GPU, read back ~3MB of planar YUV, and skip the CPU pass. Gated on a runtime check that the encoder's YUV420P plane strides are tight (no linesize padding), with the swscale path as fallback for other dimensions; LB_DISABLE_GPU_YUV forces the CPU path. Includes a CPU reference + unit tests for the packing. Also guard render_next_video_frame against re-initializing/re-emitting "Complete" every frame after the render finishes while the encoder/mux drains (the completion nulled gpu_resources but left video_state set).
This commit is contained in:
Skyler Lehmkuhl
parent
70ac0cde00
commit
2564d807a0
|
|
@ -0,0 +1,292 @@
|
|||
//! Tight GPU RGBA→YUV420p converter for video export.
|
||||
//!
|
||||
//! Unlike [`lightningbeam_core::gpu::YuvConverter`] (which writes one byte per
|
||||
//! `Rgba8Unorm` texel — a 4× readback), this writes **packed planar YUV420p** into a
|
||||
//! storage buffer, so the readback is exactly `W*H*3/2` bytes (~3.1 MB at 1080p vs
|
||||
//! 8.3 MB RGBA) and — more importantly — the per-frame CPU `rgba_to_yuv420p` (swscale)
|
||||
//! is eliminated.
|
||||
//!
|
||||
//! Color math is BT.709 **full-range** (JPEG range), matching the encoder color tags
|
||||
//! set in `setup_video_encoder` (`Space::BT709` + `Range::JPEG`).
|
||||
//!
|
||||
//! Output buffer layout (tight, little-endian byte packing into `array<u32>`):
|
||||
//! - `[0, W*H)` Y plane, row stride `W`
|
||||
//! - `[W*H, W*H + CW*CH)` U plane, row stride `CW` (`CW=W/2`, `CH=H/2`)
|
||||
//! - `[W*H+CW*CH, end)` V plane, row stride `CW`
|
||||
//!
|
||||
//! Dimension requirement: `W % 8 == 0 && H % 2 == 0` (so `W/4` and `CW/4` are whole —
|
||||
//! the shader packs 4 bytes per `u32`). [`GpuYuv::supports`] reports this; callers
|
||||
//! fall back to the CPU converter otherwise.
|
||||
|
||||
/// `true` when [`GpuYuv`] can convert these dimensions (else use the CPU path).
|
||||
pub fn supports(width: u32, height: u32) -> bool {
|
||||
width % 8 == 0 && height % 2 == 0 && width > 0 && height > 0
|
||||
}
|
||||
|
||||
/// Tight planar YUV420p byte length for `width`×`height`.
|
||||
pub fn yuv420p_len(width: u32, height: u32) -> usize {
|
||||
let y = (width * height) as usize;
|
||||
let c = ((width / 2) * (height / 2)) as usize;
|
||||
y + 2 * c
|
||||
}
|
||||
|
||||
/// GPU compute pipeline: `Rgba8Unorm` texture → tight planar YUV420p storage buffer.
|
||||
pub struct GpuYuv {
|
||||
y_pipeline: wgpu::ComputePipeline,
|
||||
uv_pipeline: wgpu::ComputePipeline,
|
||||
bind_group_layout: wgpu::BindGroupLayout,
|
||||
}
|
||||
|
||||
impl GpuYuv {
|
||||
pub fn new(device: &wgpu::Device) -> Self {
|
||||
let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
|
||||
label: Some("gpu_yuv_bgl"),
|
||||
entries: &[
|
||||
// 0: input RGBA (non-filterable, read via textureLoad)
|
||||
wgpu::BindGroupLayoutEntry {
|
||||
binding: 0,
|
||||
visibility: wgpu::ShaderStages::COMPUTE,
|
||||
ty: wgpu::BindingType::Texture {
|
||||
sample_type: wgpu::TextureSampleType::Float { filterable: false },
|
||||
view_dimension: wgpu::TextureViewDimension::D2,
|
||||
multisampled: false,
|
||||
},
|
||||
count: None,
|
||||
},
|
||||
// 1: output packed YUV (read_write so 4-byte packing writes whole u32s)
|
||||
wgpu::BindGroupLayoutEntry {
|
||||
binding: 1,
|
||||
visibility: wgpu::ShaderStages::COMPUTE,
|
||||
ty: wgpu::BindingType::Buffer {
|
||||
ty: wgpu::BufferBindingType::Storage { read_only: false },
|
||||
has_dynamic_offset: false,
|
||||
min_binding_size: None,
|
||||
},
|
||||
count: None,
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
|
||||
label: Some("gpu_yuv_pl"),
|
||||
bind_group_layouts: &[&bind_group_layout],
|
||||
push_constant_ranges: &[],
|
||||
});
|
||||
|
||||
let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
|
||||
label: Some("gpu_yuv_shader"),
|
||||
source: wgpu::ShaderSource::Wgsl(SHADER.into()),
|
||||
});
|
||||
|
||||
let mk = |entry: &str| {
|
||||
device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
|
||||
label: Some("gpu_yuv_pipeline"),
|
||||
layout: Some(&pipeline_layout),
|
||||
module: &shader,
|
||||
entry_point: Some(entry),
|
||||
compilation_options: wgpu::PipelineCompilationOptions::default(),
|
||||
cache: None,
|
||||
})
|
||||
};
|
||||
|
||||
Self {
|
||||
y_pipeline: mk("y_main"),
|
||||
uv_pipeline: mk("uv_main"),
|
||||
bind_group_layout,
|
||||
}
|
||||
}
|
||||
|
||||
/// Record the RGBA→YUV420p conversion into `encoder`.
|
||||
///
|
||||
/// `rgba_view` is the rendered frame (`Rgba8Unorm`, `width`×`height`, must have
|
||||
/// `TEXTURE_BINDING` usage). `yuv_buffer` must be a `STORAGE | COPY_SRC` buffer of
|
||||
/// at least [`yuv420p_len`] bytes (rounded up to 4). Caller must ensure
|
||||
/// [`supports`]`(width, height)`.
|
||||
pub fn convert(
|
||||
&self,
|
||||
device: &wgpu::Device,
|
||||
encoder: &mut wgpu::CommandEncoder,
|
||||
rgba_view: &wgpu::TextureView,
|
||||
yuv_buffer: &wgpu::Buffer,
|
||||
width: u32,
|
||||
height: u32,
|
||||
) {
|
||||
debug_assert!(supports(width, height));
|
||||
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
|
||||
label: Some("gpu_yuv_bg"),
|
||||
layout: &self.bind_group_layout,
|
||||
entries: &[
|
||||
wgpu::BindGroupEntry { binding: 0, resource: wgpu::BindingResource::TextureView(rgba_view) },
|
||||
wgpu::BindGroupEntry { binding: 1, resource: yuv_buffer.as_entire_binding() },
|
||||
],
|
||||
});
|
||||
|
||||
let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
|
||||
label: Some("gpu_yuv_pass"),
|
||||
timestamp_writes: None,
|
||||
});
|
||||
pass.set_bind_group(0, &bind_group, &[]);
|
||||
|
||||
// Y: one thread per 4 horizontal luma samples → (W/4)×H threads.
|
||||
pass.set_pipeline(&self.y_pipeline);
|
||||
let wg = 8u32;
|
||||
pass.dispatch_workgroups(((width / 4) + wg - 1) / wg, (height + wg - 1) / wg, 1);
|
||||
|
||||
// UV: one thread per 4 horizontal chroma samples → (CW/4)×CH = (W/8)×(H/2) threads.
|
||||
pass.set_pipeline(&self.uv_pipeline);
|
||||
let cw = width / 2;
|
||||
let ch = height / 2;
|
||||
pass.dispatch_workgroups(((cw / 4) + wg - 1) / wg, (ch + wg - 1) / wg, 1);
|
||||
}
|
||||
}
|
||||
|
||||
/// CPU reference for the exact math/layout the shader produces — used by unit tests so
|
||||
/// the packing and BT.709 coefficients stay verifiable without a GPU.
|
||||
#[cfg(test)]
|
||||
fn cpu_reference(rgba: &[u8], width: u32, height: u32) -> Vec<u8> {
|
||||
let w = width as usize;
|
||||
let h = height as usize;
|
||||
let cw = w / 2;
|
||||
let ch = h / 2;
|
||||
let mut out = vec![0u8; yuv420p_len(width, height)];
|
||||
let to_byte = |v: f32| (v.clamp(0.0, 1.0) * 255.0 + 0.5) as u8;
|
||||
let px = |x: usize, y: usize| {
|
||||
let i = (y * w + x) * 4;
|
||||
[rgba[i] as f32 / 255.0, rgba[i + 1] as f32 / 255.0, rgba[i + 2] as f32 / 255.0]
|
||||
};
|
||||
// Y
|
||||
for y in 0..h {
|
||||
for x in 0..w {
|
||||
let p = px(x, y);
|
||||
out[y * w + x] = to_byte(0.2126 * p[0] + 0.7152 * p[1] + 0.0722 * p[2]);
|
||||
}
|
||||
}
|
||||
// U/V (2x2 average)
|
||||
let y_size = w * h;
|
||||
let uv_size = cw * ch;
|
||||
for cy in 0..ch {
|
||||
for cx in 0..cw {
|
||||
let mut acc = [0.0f32; 3];
|
||||
for (dx, dy) in [(0, 0), (1, 0), (0, 1), (1, 1)] {
|
||||
let p = px(2 * cx + dx, 2 * cy + dy);
|
||||
acc[0] += p[0]; acc[1] += p[1]; acc[2] += p[2];
|
||||
}
|
||||
let a = [acc[0] / 4.0, acc[1] / 4.0, acc[2] / 4.0];
|
||||
let u = -0.1146 * a[0] - 0.3854 * a[1] + 0.5000 * a[2] + 0.5;
|
||||
let v = 0.5000 * a[0] - 0.4542 * a[1] - 0.0458 * a[2] + 0.5;
|
||||
out[y_size + cy * cw + cx] = to_byte(u);
|
||||
out[y_size + uv_size + cy * cw + cx] = to_byte(v);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
const SHADER: &str = r#"
|
||||
// RGBA -> tight planar YUV420p (BT.709 full-range), packed 4 bytes/u32.
|
||||
@group(0) @binding(0) var input_rgba: texture_2d<f32>;
|
||||
@group(0) @binding(1) var<storage, read_write> out_buf: array<u32>;
|
||||
|
||||
fn to_byte(v: f32) -> u32 { return u32(clamp(v, 0.0, 1.0) * 255.0 + 0.5); }
|
||||
|
||||
// Y plane: each thread packs 4 horizontal luma bytes.
|
||||
@compute @workgroup_size(8, 8, 1)
|
||||
fn y_main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||
let dims = textureDimensions(input_rgba);
|
||||
let w = dims.x;
|
||||
let h = dims.y;
|
||||
let x4 = gid.x * 4u;
|
||||
let y = gid.y;
|
||||
if (x4 >= w || y >= h) { return; }
|
||||
var packed: u32 = 0u;
|
||||
for (var i = 0u; i < 4u; i = i + 1u) {
|
||||
let c = textureLoad(input_rgba, vec2<u32>(x4 + i, y), 0).rgb;
|
||||
let yy = 0.2126 * c.r + 0.7152 * c.g + 0.0722 * c.b;
|
||||
packed = packed | (to_byte(yy) << (8u * i));
|
||||
}
|
||||
out_buf[(y * w + x4) / 4u] = packed;
|
||||
}
|
||||
|
||||
// U/V planes: each thread packs 4 horizontal chroma bytes (2x2 box-averaged).
|
||||
@compute @workgroup_size(8, 8, 1)
|
||||
fn uv_main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||
let dims = textureDimensions(input_rgba);
|
||||
let w = dims.x;
|
||||
let h = dims.y;
|
||||
let cw = w / 2u;
|
||||
let ch = h / 2u;
|
||||
let cx4 = gid.x * 4u;
|
||||
let cy = gid.y;
|
||||
if (cx4 >= cw || cy >= ch) { return; }
|
||||
let y_size = w * h;
|
||||
let uv_size = cw * ch;
|
||||
var up: u32 = 0u;
|
||||
var vp: u32 = 0u;
|
||||
for (var i = 0u; i < 4u; i = i + 1u) {
|
||||
let cx = cx4 + i;
|
||||
let sx = 2u * cx;
|
||||
let sy = 2u * cy;
|
||||
let p00 = textureLoad(input_rgba, vec2<u32>(sx, sy), 0).rgb;
|
||||
let p10 = textureLoad(input_rgba, vec2<u32>(sx + 1u, sy), 0).rgb;
|
||||
let p01 = textureLoad(input_rgba, vec2<u32>(sx, sy + 1u), 0).rgb;
|
||||
let p11 = textureLoad(input_rgba, vec2<u32>(sx + 1u, sy + 1u), 0).rgb;
|
||||
let a = (p00 + p10 + p01 + p11) * 0.25;
|
||||
let u = -0.1146 * a.r - 0.3854 * a.g + 0.5000 * a.b + 0.5;
|
||||
let v = 0.5000 * a.r - 0.4542 * a.g - 0.0458 * a.b + 0.5;
|
||||
up = up | (to_byte(u) << (8u * i));
|
||||
vp = vp | (to_byte(v) << (8u * i));
|
||||
}
|
||||
out_buf[(y_size + cy * cw + cx4) / 4u] = up;
|
||||
out_buf[(y_size + uv_size + cy * cw + cx4) / 4u] = vp;
|
||||
}
|
||||
"#;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn supports_dims() {
|
||||
assert!(supports(1920, 1080));
|
||||
assert!(supports(1280, 720));
|
||||
assert!(supports(8, 2));
|
||||
assert!(!supports(6, 2)); // width not %8
|
||||
assert!(!supports(8, 3)); // height odd
|
||||
assert!(!supports(0, 0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn len_matches() {
|
||||
assert_eq!(yuv420p_len(1920, 1080), 1920 * 1080 * 3 / 2);
|
||||
assert_eq!(yuv420p_len(8, 2), 8 * 2 + 2 * (4 * 1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reference_known_colors() {
|
||||
// 8x2 solid white → Y≈255, U≈V≈128. Solid black → Y=0, U=V≈128.
|
||||
let white = vec![255u8; 8 * 2 * 4];
|
||||
let out = cpu_reference(&white, 8, 2);
|
||||
let (cw, ch) = (4usize, 1usize);
|
||||
let y_size = 8 * 2;
|
||||
for &y in &out[..y_size] { assert!(y >= 254, "white Y={y}"); }
|
||||
for &u in &out[y_size..y_size + cw * ch] { assert!((u as i32 - 128).abs() <= 1, "white U={u}"); }
|
||||
|
||||
let black = vec![0u8; 8 * 2 * 4];
|
||||
let out = cpu_reference(&black, 8, 2);
|
||||
for &y in &out[..y_size] { assert_eq!(y, 0); }
|
||||
for &v in &out[y_size + cw * ch..] { assert!((v as i32 - 128).abs() <= 1, "black V={v}"); }
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reference_red_bt709() {
|
||||
// Solid red (255,0,0): Y=0.2126*255≈54; V high, U low (full range).
|
||||
let red: Vec<u8> = (0..8 * 2).flat_map(|_| [255u8, 0, 0, 255]).collect();
|
||||
let out = cpu_reference(&red, 8, 2);
|
||||
assert!((out[0] as i32 - 54).abs() <= 1, "red Y={}", out[0]);
|
||||
let y_size = 8 * 2;
|
||||
let u = out[y_size];
|
||||
let v = out[y_size + 4];
|
||||
// U = -0.1146*1*255+128 ≈ 99 ; V = 0.5*255+128 → clamps to 255
|
||||
assert!((u as i32 - 99).abs() <= 2, "red U={u}");
|
||||
assert_eq!(v, 255, "red V={v}");
|
||||
}
|
||||
}
|
||||
|
|
@ -10,6 +10,7 @@ pub mod video_exporter;
|
|||
pub mod readback_pipeline;
|
||||
pub mod perf_metrics;
|
||||
pub mod cpu_yuv_converter;
|
||||
pub mod gpu_yuv;
|
||||
|
||||
use lightningbeam_core::export::{AudioExportSettings, ImageExportSettings, VideoExportSettings, ExportProgress};
|
||||
use lightningbeam_core::document::Document;
|
||||
|
|
@ -1038,6 +1039,18 @@ impl ExportOrchestrator {
|
|||
let state = self.video_state.as_mut()
|
||||
.ok_or("No video export in progress")?;
|
||||
|
||||
// Already completed (Done sent, all frames done): don't re-initialize and
|
||||
// re-run. The completion path nulls gpu_resources but leaves video_state set
|
||||
// (cleared only when the parallel export finishes); without this guard the
|
||||
// function would re-create the GPU pipeline and re-emit "Complete" every frame
|
||||
// while the encoder/mux drains.
|
||||
if state.frame_tx.is_none()
|
||||
&& state.current_frame >= state.total_frames
|
||||
&& state.frames_in_flight == 0
|
||||
{
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let width = state.width;
|
||||
let height = state.height;
|
||||
|
||||
|
|
@ -1045,7 +1058,19 @@ impl ExportOrchestrator {
|
|||
if state.gpu_resources.is_none() {
|
||||
println!("🎬 [VIDEO EXPORT] Initializing HDR GPU + async pipeline {}x{}", width, height);
|
||||
state.gpu_resources = Some(video_exporter::ExportGpuResources::new(device, width, height));
|
||||
state.readback_pipeline = Some(readback_pipeline::ReadbackPipeline::new(device, queue, width, height));
|
||||
// Enable GPU YUV only when the encoder's YUV420P planes are tight (no linesize
|
||||
// padding) — then the packed GPU planes copy in without row misalignment.
|
||||
// Otherwise fall back to RGBA readback + CPU swscale.
|
||||
let gpu_yuv_tight = std::env::var("LB_DISABLE_GPU_YUV").is_err() && {
|
||||
let probe = ffmpeg_next::frame::Video::new(
|
||||
ffmpeg_next::format::Pixel::YUV420P, width, height,
|
||||
);
|
||||
probe.stride(0) == width as usize && probe.stride(1) == (width / 2) as usize
|
||||
};
|
||||
if !gpu_yuv_tight {
|
||||
println!("🎬 [VIDEO EXPORT] YUV planes are padded at {width}x{height}; using CPU YUV path");
|
||||
}
|
||||
state.readback_pipeline = Some(readback_pipeline::ReadbackPipeline::new(device, queue, width, height, gpu_yuv_tight));
|
||||
state.cpu_yuv_converter = Some(cpu_yuv_converter::CpuYuvConverter::new(width, height)?);
|
||||
println!("🚀 [ASYNC PIPELINE] Triple-buffered pipeline initialized");
|
||||
println!("🚀 [CPU YUV] swscale converter initialized");
|
||||
|
|
@ -1075,14 +1100,18 @@ impl ExportOrchestrator {
|
|||
}
|
||||
}
|
||||
|
||||
// Extract RGBA data (timed)
|
||||
// Extract readback data (timed)
|
||||
let extraction_start = Instant::now();
|
||||
let rgba_data = pipeline.extract_rgba_data(result.buffer_id);
|
||||
let data = pipeline.extract_rgba_data(result.buffer_id);
|
||||
let extraction_end = Instant::now();
|
||||
|
||||
// CPU YUV conversion (timed)
|
||||
// YUV planes: GPU-converted (just slice) or CPU swscale fallback (timed).
|
||||
let conversion_start = Instant::now();
|
||||
let (y, u, v) = cpu_converter.convert(&rgba_data)?;
|
||||
let (y, u, v) = if pipeline.is_yuv_mode() {
|
||||
pipeline.split_yuv(&data)
|
||||
} else {
|
||||
cpu_converter.convert(&data)?
|
||||
};
|
||||
let conversion_end = Instant::now();
|
||||
|
||||
if let Some(m) = metrics.as_mut() {
|
||||
|
|
|
|||
|
|
@ -41,7 +41,10 @@ struct PipelineBuffer {
|
|||
/// RGBA texture for GPU rendering output (Rgba8Unorm)
|
||||
rgba_texture: wgpu::Texture,
|
||||
rgba_texture_view: wgpu::TextureView,
|
||||
/// Staging buffer for GPU→CPU transfer (MAP_READ)
|
||||
/// In YUV mode: packed planar YUV420p the compute shader writes (STORAGE | COPY_SRC).
|
||||
/// `None` in RGBA fallback mode.
|
||||
yuv_buffer: Option<wgpu::Buffer>,
|
||||
/// Staging buffer for GPU→CPU transfer (MAP_READ). Holds YUV in YUV mode, RGBA otherwise.
|
||||
staging_buffer: wgpu::Buffer,
|
||||
/// Current state in the pipeline
|
||||
state: BufferState,
|
||||
|
|
@ -71,6 +74,10 @@ pub struct ReadbackPipeline {
|
|||
/// Buffer dimensions
|
||||
width: u32,
|
||||
height: u32,
|
||||
/// `Some` when converting RGBA→YUV420p on the GPU (skips the CPU swscale pass and
|
||||
/// reads back ~3 MB of planar YUV instead of 8 MB RGBA). `None` falls back to RGBA
|
||||
/// readback + CPU conversion for dimensions the packed shader can't handle.
|
||||
gpu_yuv: Option<super::gpu_yuv::GpuYuv>,
|
||||
}
|
||||
|
||||
impl ReadbackPipeline {
|
||||
|
|
@ -81,13 +88,31 @@ impl ReadbackPipeline {
|
|||
/// * `queue` - GPU queue (will be cloned for async operations)
|
||||
/// * `width` - Frame width in pixels
|
||||
/// * `height` - Frame height in pixels
|
||||
pub fn new(device: &wgpu::Device, queue: &wgpu::Queue, width: u32, height: u32) -> Self {
|
||||
/// `enable_gpu_yuv` should be `true` only when the caller has verified the encoder's
|
||||
/// `YUV420P` plane strides are tight (== width / width-2), so the packed GPU planes
|
||||
/// drop straight into the `AVFrame` without row misalignment.
|
||||
pub fn new(device: &wgpu::Device, queue: &wgpu::Queue, width: u32, height: u32, enable_gpu_yuv: bool) -> Self {
|
||||
let (readback_tx, readback_rx) = channel();
|
||||
|
||||
// GPU YUV conversion when enabled AND the dimensions fit the packed shader; else RGBA + CPU.
|
||||
let gpu_yuv = if enable_gpu_yuv && super::gpu_yuv::supports(width, height) {
|
||||
Some(super::gpu_yuv::GpuYuv::new(device))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let yuv_mode = gpu_yuv.is_some();
|
||||
|
||||
// Staging size: planar YUV420p (W*H*3/2) in YUV mode, else RGBA (W*H*4).
|
||||
let staging_size = if yuv_mode {
|
||||
super::gpu_yuv::yuv420p_len(width, height) as u64
|
||||
} else {
|
||||
(width * height * 4) as u64
|
||||
};
|
||||
|
||||
// Create 3 buffers for triple buffering
|
||||
let mut buffers = Vec::new();
|
||||
for id in 0..3 {
|
||||
// RGBA texture (Rgba8Unorm)
|
||||
// RGBA texture (Rgba8Unorm). TEXTURE_BINDING lets the YUV compute shader read it.
|
||||
let rgba_texture = device.create_texture(&wgpu::TextureDescriptor {
|
||||
label: Some(&format!("readback_rgba_texture_{}", id)),
|
||||
size: wgpu::Extent3d {
|
||||
|
|
@ -99,17 +124,29 @@ impl ReadbackPipeline {
|
|||
sample_count: 1,
|
||||
dimension: wgpu::TextureDimension::D2,
|
||||
format: wgpu::TextureFormat::Rgba8Unorm,
|
||||
usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
|
||||
usage: wgpu::TextureUsages::RENDER_ATTACHMENT
|
||||
| wgpu::TextureUsages::COPY_SRC
|
||||
| wgpu::TextureUsages::TEXTURE_BINDING,
|
||||
view_formats: &[],
|
||||
});
|
||||
|
||||
let rgba_texture_view = rgba_texture.create_view(&wgpu::TextureViewDescriptor::default());
|
||||
|
||||
let yuv_buffer = if yuv_mode {
|
||||
Some(device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some(&format!("readback_yuv_buffer_{}", id)),
|
||||
size: staging_size,
|
||||
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
|
||||
mapped_at_creation: false,
|
||||
}))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Staging buffer for GPU→CPU readback
|
||||
let rgba_buffer_size = (width * height * 4) as u64; // Rgba8Unorm = 4 bytes/pixel
|
||||
let staging_buffer = device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some(&format!("readback_staging_buffer_{}", id)),
|
||||
size: rgba_buffer_size,
|
||||
size: staging_size,
|
||||
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
|
||||
mapped_at_creation: false,
|
||||
});
|
||||
|
|
@ -118,6 +155,7 @@ impl ReadbackPipeline {
|
|||
id,
|
||||
rgba_texture,
|
||||
rgba_texture_view,
|
||||
yuv_buffer,
|
||||
staging_buffer,
|
||||
state: BufferState::Free,
|
||||
frame_num: None,
|
||||
|
|
@ -133,9 +171,23 @@ impl ReadbackPipeline {
|
|||
queue: queue.clone(),
|
||||
width,
|
||||
height,
|
||||
gpu_yuv,
|
||||
}
|
||||
}
|
||||
|
||||
/// `true` when frames are read back as planar YUV420p (GPU-converted) — the caller
|
||||
/// should slice planes with [`Self::split_yuv`] instead of running the CPU converter.
|
||||
pub fn is_yuv_mode(&self) -> bool {
|
||||
self.gpu_yuv.is_some()
|
||||
}
|
||||
|
||||
/// Split a YUV-mode readback buffer into tight (Y, U, V) planes.
|
||||
pub fn split_yuv(&self, data: &[u8]) -> (Vec<u8>, Vec<u8>, Vec<u8>) {
|
||||
let y = (self.width * self.height) as usize;
|
||||
let c = ((self.width / 2) * (self.height / 2)) as usize;
|
||||
(data[..y].to_vec(), data[y..y + c].to_vec(), data[y + c..y + 2 * c].to_vec())
|
||||
}
|
||||
|
||||
/// Acquire a free buffer for rendering (non-blocking)
|
||||
///
|
||||
/// Returns None if all buffers are in use (caller should poll and retry)
|
||||
|
|
@ -166,7 +218,12 @@ impl ReadbackPipeline {
|
|||
let buffer = &mut self.buffers[buffer_id];
|
||||
assert_eq!(buffer.state, BufferState::Rendering, "Buffer not in Rendering state");
|
||||
|
||||
// Copy RGBA texture to staging buffer
|
||||
if let (Some(gpu_yuv), Some(yuv_buffer)) = (self.gpu_yuv.as_ref(), buffer.yuv_buffer.as_ref()) {
|
||||
// GPU RGBA→YUV420p, then copy the packed YUV buffer to staging (~3 MB).
|
||||
gpu_yuv.convert(&self.device, &mut encoder, &buffer.rgba_texture_view, yuv_buffer, self.width, self.height);
|
||||
encoder.copy_buffer_to_buffer(yuv_buffer, 0, &buffer.staging_buffer, 0, buffer.staging_buffer.size());
|
||||
} else {
|
||||
// Fallback: copy the RGBA texture to staging (8 MB), CPU converts later.
|
||||
encoder.copy_texture_to_buffer(
|
||||
wgpu::TexelCopyTextureInfo {
|
||||
texture: &buffer.rgba_texture,
|
||||
|
|
@ -188,6 +245,7 @@ impl ReadbackPipeline {
|
|||
depth_or_array_layers: 1,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
// Submit GPU commands (non-blocking)
|
||||
self.queue.submit(Some(encoder.finish()));
|
||||
|
|
|
|||
Loading…
Reference in New Issue