diff --git a/apps/desktop-gpui/kernels/Cargo.toml b/apps/desktop-gpui/kernels/Cargo.toml index fab553c503..63f7bfa506 100644 --- a/apps/desktop-gpui/kernels/Cargo.toml +++ b/apps/desktop-gpui/kernels/Cargo.toml @@ -54,3 +54,7 @@ harness = false [[bench]] name = "screenshot_blur" harness = false + +[[bench]] +name = "frame_convert" +harness = false diff --git a/apps/desktop-gpui/kernels/benches/frame_convert.rs b/apps/desktop-gpui/kernels/benches/frame_convert.rs new file mode 100644 index 0000000000..18ed21fba9 --- /dev/null +++ b/apps/desktop-gpui/kernels/benches/frame_convert.rs @@ -0,0 +1,313 @@ +use cap_gpui_kernels::frame::{self as candidate, Nv12Range}; +use std::{hint::black_box, time::Instant}; + +fn main() { + println!("=== Benchmark: copy_rows ==="); + for (name, row_bytes, rows, stride, mirrored) in [ + ("1080p-unpadded", 1920 * 4, 1080, 1920 * 4, false), + ("1080p-padded-256", 1920 * 4, 1080, 7936, false), + ("4k-unpadded", 3840 * 4, 2160, 3840 * 4, false), + ("4k-padded-256-mirrored", 3840 * 4, 2160, 15616, true), + ("1248x702-padded-mirrored", 1248 * 4, 702, 5120, true), + ] { + let input: Vec = (0..stride * rows) + .map(|i| (i as u8).wrapping_mul(17)) + .collect(); + let expected = baseline::copy_rows(&input, row_bytes, stride, rows, mirrored).unwrap(); + let actual = candidate::copy_rows(&input, row_bytes, stride, rows, mirrored).unwrap(); + assert_eq!(expected, actual, "Mismatch in {name}"); + + let mut baseline_times = Vec::new(); + let mut candidate_times = Vec::new(); + + for i in 0..25 { + let start = Instant::now(); + let res = baseline::copy_rows( + black_box(&input), + black_box(row_bytes), + black_box(stride), + black_box(rows), + black_box(mirrored), + ); + let elapsed_base = start.elapsed().as_nanos(); + black_box(res); + + let start = Instant::now(); + let res = candidate::copy_rows( + black_box(&input), + black_box(row_bytes), + black_box(stride), + black_box(rows), + black_box(mirrored), + ); + let elapsed_cand = start.elapsed().as_nanos(); + black_box(res); + + if i >= 5 { + baseline_times.push(elapsed_base); + candidate_times.push(elapsed_cand); + } + } + baseline_times.sort(); + candidate_times.sort(); + let base_med = baseline_times[baseline_times.len() / 2] as f64 / 1_000.0; + let cand_med = candidate_times[candidate_times.len() / 2] as f64 / 1_000.0; + let speedup = base_med / cand_med; + println!( + "{name}: baseline = {base_med:.1} us, candidate = {cand_med:.1} us, speedup = {speedup:.2}x" + ); + } + + println!("\n=== Benchmark: downsample_nearest ==="); + for (name, width, height, stride, target_w, target_h) in [ + ("1080p->720p", 1920, 1080, 1920 * 4, 1280, 720), + ("1080p->360p", 1920, 1080, 1920 * 4, 640, 360), + ("4k->1080p", 3840, 2160, 3840 * 4, 1920, 1080), + ] { + let input: Vec = (0..stride * height) + .map(|i| (i as u8).wrapping_mul(13)) + .collect(); + let expected = + baseline::downsample_nearest(&input, width, height, stride, target_w, target_h); + let actual = + candidate::downsample_nearest(&input, width, height, stride, target_w, target_h); + assert_eq!(expected, actual, "Mismatch in {name}"); + + let mut baseline_times = Vec::new(); + let mut candidate_times = Vec::new(); + + for i in 0..25 { + let start = Instant::now(); + let res = baseline::downsample_nearest( + black_box(&input), + black_box(width), + black_box(height), + black_box(stride), + black_box(target_w), + black_box(target_h), + ); + let elapsed_base = start.elapsed().as_nanos(); + black_box(res); + + let start = Instant::now(); + let res = candidate::downsample_nearest( + black_box(&input), + black_box(width), + black_box(height), + black_box(stride), + black_box(target_w), + black_box(target_h), + ); + let elapsed_cand = start.elapsed().as_nanos(); + black_box(res); + + if i >= 5 { + baseline_times.push(elapsed_base); + candidate_times.push(elapsed_cand); + } + } + baseline_times.sort(); + candidate_times.sort(); + let base_med = baseline_times[baseline_times.len() / 2] as f64 / 1_000.0; + let cand_med = candidate_times[candidate_times.len() / 2] as f64 / 1_000.0; + let speedup = base_med / cand_med; + println!( + "{name}: baseline = {base_med:.1} us, candidate = {cand_med:.1} us, speedup = {speedup:.2}x" + ); + } + + println!("\n=== Benchmark: convert_nv12_planes ==="); + for (name, width, height) in [ + ("720p-nv12", 1280, 720), + ("1080p-nv12", 1920, 1080), + ("4k-nv12", 3840, 2160), + ] { + let y_plane: Vec = (0..width * height) + .map(|i| (i as u8).wrapping_mul(7)) + .collect(); + let uv_plane: Vec = (0..width * (height / 2)) + .map(|i| 128u8.wrapping_add(i as u8)) + .collect(); + let expected = baseline::convert_nv12_planes( + &y_plane, + width, + &uv_plane, + width, + width, + height, + Nv12Range::Video, + ) + .unwrap(); + let actual = candidate::convert_nv12_planes( + &y_plane, + width, + &uv_plane, + width, + width, + height, + Nv12Range::Video, + ) + .unwrap(); + assert_eq!(expected, actual, "Mismatch in {name}"); + + let mut baseline_times = Vec::new(); + let mut candidate_times = Vec::new(); + + for i in 0..20 { + let start = Instant::now(); + let res = baseline::convert_nv12_planes( + black_box(&y_plane), + black_box(width), + black_box(&uv_plane), + black_box(width), + black_box(width), + black_box(height), + black_box(Nv12Range::Video), + ); + let elapsed_base = start.elapsed().as_nanos(); + black_box(res); + + let start = Instant::now(); + let res = candidate::convert_nv12_planes( + black_box(&y_plane), + black_box(width), + black_box(&uv_plane), + black_box(width), + black_box(width), + black_box(height), + black_box(Nv12Range::Video), + ); + let elapsed_cand = start.elapsed().as_nanos(); + black_box(res); + + if i >= 5 { + baseline_times.push(elapsed_base); + candidate_times.push(elapsed_cand); + } + } + baseline_times.sort(); + candidate_times.sort(); + let base_med = baseline_times[baseline_times.len() / 2] as f64 / 1_000.0; + let cand_med = candidate_times[candidate_times.len() / 2] as f64 / 1_000.0; + let speedup = base_med / cand_med; + println!( + "{name}: baseline = {base_med:.1} us, candidate = {cand_med:.1} us, speedup = {speedup:.2}x" + ); + } +} + +// Baseline implementations prior to this optimization +mod baseline { + use super::Nv12Range; + + pub fn copy_rows( + data: &[u8], + row_bytes: usize, + stride: usize, + rows: usize, + mirrored: bool, + ) -> Option> { + if stride < row_bytes || data.len() < rows.checked_mul(stride)? { + return None; + } + let mut pixels = vec![0; rows.checked_mul(row_bytes)?]; + for (row, output) in pixels.chunks_exact_mut(row_bytes.max(1)).enumerate() { + let input = &data[row * stride..row * stride + row_bytes]; + if mirrored { + for (target, source) in output.chunks_exact_mut(4).zip(input.chunks_exact(4).rev()) + { + target.copy_from_slice(source); + } + } else { + output.copy_from_slice(input); + } + } + Some(pixels) + } + + pub fn downsample_nearest( + data: &[u8], + width: usize, + height: usize, + stride: usize, + target_width: u32, + target_height: u32, + ) -> Vec { + let target_width = target_width as usize; + let target_height = target_height as usize; + let mut pixels = Vec::with_capacity(target_width * target_height * 4); + for y in 0..target_height { + let row = &data[y * height / target_height * stride..]; + for x in 0..target_width { + let offset = x * width / target_width * 4; + pixels.extend_from_slice(&row[offset..offset + 4]); + } + } + pixels + } + + pub fn convert_nv12_planes( + y_plane: &[u8], + y_stride: usize, + uv_plane: &[u8], + uv_stride: usize, + width: usize, + height: usize, + range: Nv12Range, + ) -> Option> { + let mut rgba_data = vec![0u8; width * height * 4]; + + for y_idx in 0..height { + let y_row_start = y_idx * y_stride; + if y_row_start + width > y_plane.len() { + return None; + } + let y_row = &y_plane[y_row_start..y_row_start + width]; + + let uv_row_start = (y_idx / 2) * uv_stride; + if uv_row_start + width > uv_plane.len() { + return None; + } + let uv_row = &uv_plane[uv_row_start..uv_row_start + width]; + + for (x, y_val) in y_row.iter().enumerate().take(width) { + let uv_index = (x / 2) * 2; + if uv_index + 1 >= uv_row.len() { + return None; + } + + let cb = uv_row[uv_index]; + let cr = uv_row[uv_index + 1]; + let (r, g, b) = ycbcr_to_rgb(*y_val, cb, cr, range); + let out = (y_idx * width + x) * 4; + rgba_data[out] = r; + rgba_data[out + 1] = g; + rgba_data[out + 2] = b; + rgba_data[out + 3] = 255; + } + } + + Some(rgba_data) + } + + fn clamp_channel(value: f32) -> u8 { + value.clamp(0.0, 255.0) as u8 + } + + fn ycbcr_to_rgb(y: u8, cb: u8, cr: u8, range: Nv12Range) -> (u8, u8, u8) { + let y = y as f32; + let cb = cb as f32 - 128.0; + let cr = cr as f32 - 128.0; + + let (y_value, scale) = match range { + Nv12Range::Video => ((y - 16.0).max(0.0), 1.164383_f32), + Nv12Range::_Full => (y, 1.0_f32), + }; + + let r = scale * y_value + 1.596027_f32 * cr; + let g = scale * y_value - 0.391762_f32 * cb - 0.812968_f32 * cr; + let b = scale * y_value + 2.017232_f32 * cb; + + (clamp_channel(r), clamp_channel(g), clamp_channel(b)) + } +} diff --git a/apps/desktop-gpui/kernels/src/frame.rs b/apps/desktop-gpui/kernels/src/frame.rs index a5f6cc667d..e89667f484 100644 --- a/apps/desktop-gpui/kernels/src/frame.rs +++ b/apps/desktop-gpui/kernels/src/frame.rs @@ -68,7 +68,12 @@ pub fn copy_rows( if stride < row_bytes || data.len() < rows.checked_mul(stride)? { return None; } - let mut pixels = vec![0; rows.checked_mul(row_bytes)?]; + let total_bytes = rows.checked_mul(row_bytes)?; + if !mirrored && stride == row_bytes { + return Some(data[..total_bytes].to_vec()); + } + + let mut pixels = vec![0; total_bytes]; for (row, output) in pixels.chunks_exact_mut(row_bytes.max(1)).enumerate() { let input = &data[row * stride..row * stride + row_bytes]; if mirrored { @@ -92,12 +97,26 @@ pub fn downsample_nearest( ) -> Vec { let target_width = target_width as usize; let target_height = target_height as usize; - let mut pixels = Vec::with_capacity(target_width * target_height * 4); - for y in 0..target_height { - let row = &data[y * height / target_height * stride..]; - for x in 0..target_width { - let offset = x * width / target_width * 4; - pixels.extend_from_slice(&row[offset..offset + 4]); + if target_width == 0 || target_height == 0 || width == 0 || height == 0 { + return Vec::new(); + } + let target_row_bytes = target_width * 4; + let mut pixels = vec![0u8; target_row_bytes * target_height]; + + let x_offsets: Vec = (0..target_width) + .map(|x| (x * width / target_width) * 4) + .collect(); + + for (y, target_row) in pixels.chunks_exact_mut(target_row_bytes).enumerate() { + let src_row_offset = (y * height / target_height) * stride; + if src_row_offset >= data.len() { + break; + } + let src_row = &data[src_row_offset..]; + for (&src_offset, dst_pixel) in x_offsets.iter().zip(target_row.chunks_exact_mut(4)) { + if src_offset + 4 <= src_row.len() { + dst_pixel.copy_from_slice(&src_row[src_offset..src_offset + 4]); + } } } pixels @@ -172,7 +191,7 @@ fn reorder_32bit_row Option> { - let mut rgba_data = vec![0u8; width * height * 4]; + let row_bytes = width.checked_mul(4)?; + let output_len = row_bytes.checked_mul(height)?; + if output_len == 0 { + return Some(Vec::new()); + } + let mut rgba_data = vec![0u8; output_len]; + + let (scale, y_sub) = match range { + Nv12Range::Video => (1.164383_f32, 16.0_f32), + Nv12Range::_Full => (1.0_f32, 0.0_f32), + }; - for y_idx in 0..height { + for (y_idx, rgba_row) in rgba_data.chunks_exact_mut(row_bytes).enumerate() { let y_row_start = y_idx * y_stride; if y_row_start + width > y_plane.len() { tracing::warn!( @@ -214,25 +243,53 @@ pub fn convert_nv12_planes( } let uv_row = &uv_plane[uv_row_start..uv_row_start + width]; - for (x, y_val) in y_row.iter().enumerate().take(width) { - let uv_index = (x / 2) * 2; - if uv_index + 1 >= uv_row.len() { + let pairs = width / 2; + let mut x = 0; + let mut dst_idx = 0; + + for _ in 0..pairs { + let cb = uv_row[x]; + let cr = uv_row[x + 1]; + let cb_f = cb as f32 - 128.0; + let cr_f = cr as f32 - 128.0; + + let r_chroma = 1.596027_f32 * cr_f; + let g_cb = 0.391762_f32 * cb_f; + let g_cr = 0.812968_f32 * cr_f; + let b_chroma = 2.017232_f32 * cb_f; + + let y1 = ((f32::from(y_row[x]) - y_sub).max(0.0)) * scale; + rgba_row[dst_idx] = clamp_channel(y1 + r_chroma); + rgba_row[dst_idx + 1] = clamp_channel(y1 - g_cb - g_cr); + rgba_row[dst_idx + 2] = clamp_channel(y1 + b_chroma); + rgba_row[dst_idx + 3] = 255; + + let y2 = ((f32::from(y_row[x + 1]) - y_sub).max(0.0)) * scale; + rgba_row[dst_idx + 4] = clamp_channel(y2 + r_chroma); + rgba_row[dst_idx + 5] = clamp_channel(y2 - g_cb - g_cr); + rgba_row[dst_idx + 6] = clamp_channel(y2 + b_chroma); + rgba_row[dst_idx + 7] = 255; + + x += 2; + dst_idx += 8; + } + + if x < width { + if x + 1 >= uv_row.len() { tracing::warn!( - uv_index, + uv_index = x, uv_row_len = uv_row.len(), "UV index out of bounds during conversion", ); return None; } - - let cb = uv_row[uv_index]; - let cr = uv_row[uv_index + 1]; - let (r, g, b) = ycbcr_to_rgb(*y_val, cb, cr, range); - let out = (y_idx * width + x) * 4; - rgba_data[out] = r; - rgba_data[out + 1] = g; - rgba_data[out + 2] = b; - rgba_data[out + 3] = 255; + let cb = uv_row[x]; + let cr = uv_row[x + 1]; + let (r, g, b) = ycbcr_to_rgb(y_row[x], cb, cr, range); + rgba_row[dst_idx] = r; + rgba_row[dst_idx + 1] = g; + rgba_row[dst_idx + 2] = b; + rgba_row[dst_idx + 3] = 255; } } @@ -392,4 +449,74 @@ mod tests { assert_eq!(&out[0..8], &[0, 0, 0, 255, 254, 254, 254, 255]); assert!(convert_nv12_planes(&[16], 4, &uv_plane, 4, 2, 1, Nv12Range::Video).is_none()); } + + #[test] + fn nv12_precision_and_odd_width_rejection() { + let (r, g, b) = ycbcr_to_rgb(36, 80, 58, Nv12Range::Video); + assert_eq!(g, 99); + let y_plane = vec![36, 36]; + let uv_plane = vec![80, 58]; + let out = convert_nv12_planes(&y_plane, 2, &uv_plane, 2, 2, 1, Nv12Range::Video).unwrap(); + assert_eq!(out, vec![r, 99, b, 255, r, 99, b, 255]); + + let y_plane_odd = vec![36, 36, 36]; + let uv_plane_short = vec![80, 58, 80]; + assert!( + convert_nv12_planes(&y_plane_odd, 3, &uv_plane_short, 3, 3, 1, Nv12Range::Video) + .is_none() + ); + } + + #[test] + fn nv12_sweeps_y_cb_cr_ranges_against_ycbcr_to_rgb() { + for range in [Nv12Range::Video, Nv12Range::_Full] { + for &y in &[0u8, 16, 36, 81, 128, 200, 235, 255] { + for cb in (0..=255).step_by(15) { + for cr in (0..=255).step_by(15) { + let y_plane = vec![y, y]; + let uv_plane = vec![cb as u8, cr as u8]; + let out = convert_nv12_planes(&y_plane, 2, &uv_plane, 2, 2, 1, range) + .expect("conversion should succeed"); + let (er, eg, eb) = ycbcr_to_rgb(y, cb as u8, cr as u8, range); + assert_eq!( + &out[0..4], + &[er, eg, eb, 255], + "range={range:?} y={y} cb={cb} cr={cr}" + ); + assert_eq!( + &out[4..8], + &[er, eg, eb, 255], + "range={range:?} y={y} cb={cb} cr={cr}" + ); + } + } + } + } + } + + #[test] + fn copy_rows_returns_a_direct_copy_for_unpadded_unmirrored_rows() { + let data: Vec = (0..32).collect(); + + assert_eq!( + copy_rows(&data, 8, 8, 4, false).unwrap(), + data[..32].to_vec() + ); + assert_eq!( + copy_rows(&data, 8, 16, 2, false).unwrap(), + [&data[..8], &data[16..24]].concat(), + "padded sources still strip the stride padding" + ); + + let expected_mirrored: Vec = data[..8] + .chunks_exact(4) + .rev() + .flat_map(|px| px.to_vec()) + .collect(); + assert_eq!( + copy_rows(&data, 8, 8, 1, true).unwrap(), + expected_mirrored, + "mirrored rows bypass the direct-copy fast path" + ); + } }