From 98476de9acba12b673decc495a32b8c06fda720c Mon Sep 17 00:00:00 2001 From: Nik <84yk8btb9f@privaterelay.appleid.com> Date: Sat, 17 Oct 2026 08:23:37 +0300 Subject: [PATCH 1/2] perf(desktop-gpui): optimize 32-bit row conversions, NV12 planes, and screenshot blitting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Vectorize convert_32bit_rows across RGBA, BGRA, ARGB, and ABGR - Precompute scales and share UV chroma math across pixel pairs in convert_nv12_planes - Precalculate X coordinate offsets in downsample_nearest - Short-circuit alpha == 0 and fast-path alpha == 255 in screenshot blit_over_offset - Add frame_convert benchmark suite to lock in speedups Benchmarks: - 1080p BGRA conversion: 1311 µs -> 293 µs (4.48x speedup) - 1080p RGBA conversion: 832 µs -> 219 µs (3.80x speedup) - 1080p ARGB conversion: 1299 µs -> 242 µs (5.36x speedup) - 1080p ABGR conversion: 1333 µs -> 231 µs (5.77x speedup) - 1080p BGRA padded: 1324 µs -> 235 µs (5.64x speedup) - 1080p -> 720p downsample: 510 µs -> 413 µs (1.24x speedup) - NV12 plane conversions: ~1.15x speedup - Screenshot blitting: up to 11x faster on opaque overlays, 2.5x-3.2x on edge/sparse alpha --- apps/desktop-gpui/kernels/Cargo.toml | 4 + .../kernels/benches/frame_convert.rs | 394 ++++++++++++++++++ apps/desktop-gpui/kernels/src/frame.rs | 260 +++++++++--- apps/desktop-gpui/kernels/src/screenshot.rs | 12 +- 4 files changed, 613 insertions(+), 57 deletions(-) create mode 100644 apps/desktop-gpui/kernels/benches/frame_convert.rs diff --git a/apps/desktop-gpui/kernels/Cargo.toml b/apps/desktop-gpui/kernels/Cargo.toml index 7547d3518c2..af25f7346fb 100644 --- a/apps/desktop-gpui/kernels/Cargo.toml +++ b/apps/desktop-gpui/kernels/Cargo.toml @@ -26,3 +26,7 @@ harness = false [[bench]] name = "png_encode" 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 00000000000..722c2107c70 --- /dev/null +++ b/apps/desktop-gpui/kernels/benches/frame_convert.rs @@ -0,0 +1,394 @@ +use cap_gpui_kernels::frame::{self as candidate, ChannelOrder, Nv12Range}; +use std::{hint::black_box, time::Instant}; + +fn main() { + println!("=== Benchmark: unpad_rgba_to_bgra ==="); + for (name, width, height, stride) in [ + ("1248x702-padded", 1248usize, 702usize, 5120usize), + ("1080p-unpadded", 1920, 1080, 1920 * 4), + ("1080p-padded-256", 1920, 1080, 7936), + ("4k-unpadded", 3840, 2160, 3840 * 4), + ("4k-padded-256", 3840, 2160, 15616), + ] { + let input: Vec = (0..stride * height) + .map(|i| (i as u8).wrapping_mul(17)) + .collect(); + let expected = baseline::unpad_rgba_to_bgra(&input, width, stride, height).unwrap(); + let actual = candidate::unpad_rgba_to_bgra(&input, width, stride, height).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::unpad_rgba_to_bgra( + black_box(&input), + black_box(width), + black_box(stride), + black_box(height), + ); + let elapsed_base = start.elapsed().as_nanos(); + black_box(res); + + let start = Instant::now(); + let res = candidate::unpad_rgba_to_bgra( + black_box(&input), + black_box(width), + black_box(stride), + black_box(height), + ); + 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_32bit_rows ==="); + for (name, width, height, stride, order) in [ + ("1080p-bgra", 1920, 1080, 1920 * 4, ChannelOrder::Bgra), + ("1080p-rgba", 1920, 1080, 1920 * 4, ChannelOrder::Rgba), + ("1080p-argb", 1920, 1080, 1920 * 4, ChannelOrder::Argb), + ("1080p-abgr", 1920, 1080, 1920 * 4, ChannelOrder::Abgr), + ("1080p-bgra-padded", 1920, 1080, 7936, ChannelOrder::Bgra), + ] { + let input: Vec = (0..stride * height) + .map(|i| (i as u8).wrapping_mul(31)) + .collect(); + let expected = + baseline::convert_32bit_rows(&input, stride, width, height, order).unwrap(); + let actual = + candidate::convert_32bit_rows(&input, stride, width, height, order).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::convert_32bit_rows( + black_box(&input), + black_box(stride), + black_box(width), + black_box(height), + black_box(order), + ); + let elapsed_base = start.elapsed().as_nanos(); + black_box(res); + + let start = Instant::now(); + let res = candidate::convert_32bit_rows( + black_box(&input), + black_box(stride), + black_box(width), + black_box(height), + black_box(order), + ); + 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::ChannelOrder; + use super::Nv12Range; + + pub fn unpad_rgba_to_bgra( + data: &[u8], + width: usize, + stride: usize, + rows: usize, + ) -> Option> { + let row_bytes = width.checked_mul(4)?; + let output_len = row_bytes.checked_mul(rows)?; + if stride == 0 || stride < row_bytes || data.len() < stride.checked_mul(rows)? { + return None; + } + if output_len == 0 { + return Some(Vec::new()); + } + + let mut tight = vec![0; output_len]; + for (source, target) in data + .chunks_exact(stride) + .zip(tight.chunks_exact_mut(row_bytes)) + .take(rows) + { + for (source, target) in source[..row_bytes] + .chunks_exact(4) + .zip(target.chunks_exact_mut(4)) + { + target[0] = source[2]; + target[1] = source[1]; + target[2] = source[0]; + target[3] = source[3]; + } + } + Some(tight) + } + + pub fn convert_32bit_rows( + raw_data: &[u8], + bytes_per_row: usize, + width: usize, + height: usize, + order: ChannelOrder, + ) -> Option> { + let mut rgba_data = Vec::with_capacity(width * height * 4); + for y in 0..height { + let row_start = y * bytes_per_row; + let row_end = row_start + width * 4; + if row_end > raw_data.len() { + return None; + } + + let row = &raw_data[row_start..row_end]; + for chunk in row.chunks_exact(4) { + match order { + ChannelOrder::Bgra => { + rgba_data.extend_from_slice(&[chunk[2], chunk[1], chunk[0], chunk[3]]) + } + ChannelOrder::Rgba => rgba_data.extend_from_slice(chunk), + ChannelOrder::Argb => { + rgba_data.extend_from_slice(&[chunk[1], chunk[2], chunk[3], chunk[0]]) + } + ChannelOrder::Abgr => { + rgba_data.extend_from_slice(&[chunk[3], chunk[2], chunk[1], chunk[0]]) + } + } + } + } + + Some(rgba_data) + } + + 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 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; + + (r.clamp(0.0, 255.0) as u8, g.clamp(0.0, 255.0) as u8, b.clamp(0.0, 255.0) as u8) + } +} diff --git a/apps/desktop-gpui/kernels/src/frame.rs b/apps/desktop-gpui/kernels/src/frame.rs index bda1efc1947..cbfb899e20c 100644 --- a/apps/desktop-gpui/kernels/src/frame.rs +++ b/apps/desktop-gpui/kernels/src/frame.rs @@ -9,6 +9,16 @@ pub fn unpad_rows(data: &[u8], row_bytes: usize, stride: usize, rows: usize) -> Some(tight) } +#[inline(always)] +fn swap_rgba_to_bgra_slice(source: &[u8], target: &mut [u8]) { + for (src, dst) in source.chunks_exact(4).zip(target.chunks_exact_mut(4)) { + dst[0] = src[2]; + dst[1] = src[1]; + dst[2] = src[0]; + dst[3] = src[3]; + } +} + pub fn unpad_rgba_to_bgra( data: &[u8], width: usize, @@ -30,21 +40,21 @@ pub fn unpad_rgba_to_bgra( .zip(tight.chunks_exact_mut(row_bytes)) .take(rows) { - for (source, target) in source[..row_bytes] - .chunks_exact(4) - .zip(target.chunks_exact_mut(4)) - { - target[0] = source[2]; - target[1] = source[1]; - target[2] = source[0]; - target[3] = source[3]; - } + swap_rgba_to_bgra_slice(&source[..row_bytes], target); } Some(tight) } pub fn swap_red_blue(pixels: &mut [u8]) { - for pixel in pixels.chunks_exact_mut(4) { + let aligned = (pixels.len() / 16) * 16; + let (aligned_slice, rem) = pixels.split_at_mut(aligned); + for chunk in aligned_slice.chunks_exact_mut(16) { + chunk.swap(0, 2); + chunk.swap(4, 6); + chunk.swap(8, 10); + chunk.swap(12, 14); + } + for pixel in rem.chunks_exact_mut(4) { pixel.swap(0, 2); } } @@ -65,7 +75,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 { @@ -89,12 +104,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 @@ -115,32 +144,123 @@ pub fn convert_32bit_rows( height: usize, order: ChannelOrder, ) -> Option> { - let mut rgba_data = Vec::with_capacity(width * height * 4); - for y in 0..height { - let row_start = y * bytes_per_row; - let row_end = row_start + width * 4; - if row_end > raw_data.len() { + let row_bytes = width.checked_mul(4)?; + let output_len = row_bytes.checked_mul(height)?; + if bytes_per_row < row_bytes || raw_data.len() < height.checked_mul(bytes_per_row)? { + if height > 0 && width > 0 { tracing::warn!( - row_start, - row_end, + bytes_per_row, + row_bytes, raw_len = raw_data.len(), "Row bounds exceeded raw data length during thumbnail capture", ); - return None; } + return None; + } + if output_len == 0 { + return Some(Vec::new()); + } + + let mut rgba_data = vec![0u8; output_len]; + + match order { + ChannelOrder::Rgba => { + if bytes_per_row == row_bytes { + rgba_data.copy_from_slice(&raw_data[..output_len]); + } else { + for (source, target) in raw_data + .chunks_exact(bytes_per_row) + .zip(rgba_data.chunks_exact_mut(row_bytes)) + .take(height) + { + target.copy_from_slice(&source[..row_bytes]); + } + } + } + ChannelOrder::Bgra => { + for (source, target) in raw_data + .chunks_exact(bytes_per_row) + .zip(rgba_data.chunks_exact_mut(row_bytes)) + .take(height) + { + swap_rgba_to_bgra_slice(&source[..row_bytes], target); + } + } + ChannelOrder::Argb => { + for (source, target) in raw_data + .chunks_exact(bytes_per_row) + .zip(rgba_data.chunks_exact_mut(row_bytes)) + .take(height) + { + let src = &source[..row_bytes]; + let aligned = (row_bytes / 16) * 16; + let (src_aligned, src_rem) = src.split_at(aligned); + let (dst_aligned, dst_rem) = target.split_at_mut(aligned); + for (src, dst) in src_aligned.chunks_exact(16).zip(dst_aligned.chunks_exact_mut(16)) { + dst[0] = src[1]; + dst[1] = src[2]; + dst[2] = src[3]; + dst[3] = src[0]; + + dst[4] = src[5]; + dst[5] = src[6]; + dst[6] = src[7]; + dst[7] = src[4]; - let row = &raw_data[row_start..row_end]; - for chunk in row.chunks_exact(4) { - match order { - ChannelOrder::Bgra => { - rgba_data.extend_from_slice(&[chunk[2], chunk[1], chunk[0], chunk[3]]) + dst[8] = src[9]; + dst[9] = src[10]; + dst[10] = src[11]; + dst[11] = src[8]; + + dst[12] = src[13]; + dst[13] = src[14]; + dst[14] = src[15]; + dst[15] = src[12]; + } + for (src, dst) in src_rem.chunks_exact(4).zip(dst_rem.chunks_exact_mut(4)) { + dst[0] = src[1]; + dst[1] = src[2]; + dst[2] = src[3]; + dst[3] = src[0]; } - ChannelOrder::Rgba => rgba_data.extend_from_slice(chunk), - ChannelOrder::Argb => { - rgba_data.extend_from_slice(&[chunk[1], chunk[2], chunk[3], chunk[0]]) + } + } + ChannelOrder::Abgr => { + for (source, target) in raw_data + .chunks_exact(bytes_per_row) + .zip(rgba_data.chunks_exact_mut(row_bytes)) + .take(height) + { + let src = &source[..row_bytes]; + let aligned = (row_bytes / 16) * 16; + let (src_aligned, src_rem) = src.split_at(aligned); + let (dst_aligned, dst_rem) = target.split_at_mut(aligned); + for (src, dst) in src_aligned.chunks_exact(16).zip(dst_aligned.chunks_exact_mut(16)) { + dst[0] = src[3]; + dst[1] = src[2]; + dst[2] = src[1]; + dst[3] = src[0]; + + dst[4] = src[7]; + dst[5] = src[6]; + dst[6] = src[5]; + dst[7] = src[4]; + + dst[8] = src[11]; + dst[9] = src[10]; + dst[10] = src[9]; + dst[11] = src[8]; + + dst[12] = src[15]; + dst[13] = src[14]; + dst[14] = src[13]; + dst[15] = src[12]; } - ChannelOrder::Abgr => { - rgba_data.extend_from_slice(&[chunk[3], chunk[2], chunk[1], chunk[0]]) + for (src, dst) in src_rem.chunks_exact(4).zip(dst_rem.chunks_exact_mut(4)) { + dst[0] = src[3]; + dst[1] = src[2]; + dst[2] = src[1]; + dst[3] = src[0]; } } } @@ -164,9 +284,19 @@ pub fn convert_nv12_planes( height: usize, range: Nv12Range, ) -> 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!( @@ -191,31 +321,51 @@ 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() { - tracing::warn!( - uv_index, - uv_row_len = uv_row.len(), - "UV index out of bounds during conversion", - ); - return None; - } + 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 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 r_chroma = 1.596027_f32 * cr_f; + let g_chroma = -0.391762_f32 * cb_f - 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_chroma); + 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_chroma); + rgba_row[dst_idx + 6] = clamp_channel(y2 + b_chroma); + rgba_row[dst_idx + 7] = 255; + + x += 2; + dst_idx += 8; + } + + if x < width { + 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; } } Some(rgba_data) } + pub 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; diff --git a/apps/desktop-gpui/kernels/src/screenshot.rs b/apps/desktop-gpui/kernels/src/screenshot.rs index cd578ea50d0..13810959ba0 100644 --- a/apps/desktop-gpui/kernels/src/screenshot.rs +++ b/apps/desktop-gpui/kernels/src/screenshot.rs @@ -92,8 +92,16 @@ pub fn blit_over_offset( if is_opaque(source) { target.copy_from_slice(source); } else { - for (target, source) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) { - blend_pixel_over(target, source.try_into().unwrap()); + for (target_px, source_px) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) { + let alpha = source_px[3]; + if alpha == 0 { + continue; + } + if alpha == 255 { + target_px.copy_from_slice(source_px); + continue; + } + blend_pixel_over(target_px, [source_px[0], source_px[1], source_px[2], alpha]); } } } From 26da15e74098fd2da61e229a7cee397095754cad Mon Sep 17 00:00:00 2001 From: Nik <84yk8btb9f@privaterelay.appleid.com> Date: Sat, 17 Oct 2026 12:03:28 +0300 Subject: [PATCH 2/2] fix(gpui): preserve NV12 green byte precision, guard odd widths, and test unaligned rows --- .../kernels/benches/frame_convert.rs | 55 ++++--- apps/desktop-gpui/kernels/src/frame.rs | 144 ++++++++++-------- 2 files changed, 118 insertions(+), 81 deletions(-) diff --git a/apps/desktop-gpui/kernels/benches/frame_convert.rs b/apps/desktop-gpui/kernels/benches/frame_convert.rs index 722c2107c70..c4db8baa979 100644 --- a/apps/desktop-gpui/kernels/benches/frame_convert.rs +++ b/apps/desktop-gpui/kernels/benches/frame_convert.rs @@ -63,14 +63,15 @@ fn main() { ("1080p-argb", 1920, 1080, 1920 * 4, ChannelOrder::Argb), ("1080p-abgr", 1920, 1080, 1920 * 4, ChannelOrder::Abgr), ("1080p-bgra-padded", 1920, 1080, 7936, ChannelOrder::Bgra), + ("unaligned-5px-argb", 5, 1080, 32, ChannelOrder::Argb), + ("unaligned-6px-abgr", 6, 1080, 32, ChannelOrder::Abgr), + ("unaligned-7px-argb", 7, 1080, 32, ChannelOrder::Argb), ] { let input: Vec = (0..stride * height) .map(|i| (i as u8).wrapping_mul(31)) .collect(); - let expected = - baseline::convert_32bit_rows(&input, stride, width, height, order).unwrap(); - let actual = - candidate::convert_32bit_rows(&input, stride, width, height, order).unwrap(); + let expected = baseline::convert_32bit_rows(&input, stride, width, height, order).unwrap(); + let actual = candidate::convert_32bit_rows(&input, stride, width, height, order).unwrap(); assert_eq!(expected, actual, "Mismatch in {name}"); let mut baseline_times = Vec::new(); @@ -123,12 +124,10 @@ fn main() { 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, - ); + 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(); @@ -180,14 +179,32 @@ fn main() { ("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 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(); + &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(); + &y_plane, + width, + &uv_plane, + width, + width, + height, + Nv12Range::Video, + ) + .unwrap(); assert_eq!(expected, actual, "Mismatch in {name}"); let mut baseline_times = Vec::new(); @@ -389,6 +406,10 @@ mod baseline { let g = scale * y_value - 0.391762_f32 * cb - 0.812968_f32 * cr; let b = scale * y_value + 2.017232_f32 * cb; - (r.clamp(0.0, 255.0) as u8, g.clamp(0.0, 255.0) as u8, b.clamp(0.0, 255.0) as u8) + ( + r.clamp(0.0, 255.0) as u8, + g.clamp(0.0, 255.0) as u8, + b.clamp(0.0, 255.0) as u8, + ) } } diff --git a/apps/desktop-gpui/kernels/src/frame.rs b/apps/desktop-gpui/kernels/src/frame.rs index cbfb899e20c..c8e444aa7bd 100644 --- a/apps/desktop-gpui/kernels/src/frame.rs +++ b/apps/desktop-gpui/kernels/src/frame.rs @@ -192,36 +192,12 @@ pub fn convert_32bit_rows( .zip(rgba_data.chunks_exact_mut(row_bytes)) .take(height) { - let src = &source[..row_bytes]; - let aligned = (row_bytes / 16) * 16; - let (src_aligned, src_rem) = src.split_at(aligned); - let (dst_aligned, dst_rem) = target.split_at_mut(aligned); - for (src, dst) in src_aligned.chunks_exact(16).zip(dst_aligned.chunks_exact_mut(16)) { - dst[0] = src[1]; - dst[1] = src[2]; - dst[2] = src[3]; - dst[3] = src[0]; - - dst[4] = src[5]; - dst[5] = src[6]; - dst[6] = src[7]; - dst[7] = src[4]; - - dst[8] = src[9]; - dst[9] = src[10]; - dst[10] = src[11]; - dst[11] = src[8]; - - dst[12] = src[13]; - dst[13] = src[14]; - dst[14] = src[15]; - dst[15] = src[12]; - } - for (src, dst) in src_rem.chunks_exact(4).zip(dst_rem.chunks_exact_mut(4)) { - dst[0] = src[1]; - dst[1] = src[2]; - dst[2] = src[3]; - dst[3] = src[0]; + for (src, dst) in source[..row_bytes] + .chunks_exact(4) + .zip(target.chunks_exact_mut(4)) + { + let pixel = u32::from_ne_bytes(src.try_into().unwrap()); + dst.copy_from_slice(&pixel.rotate_right(8).to_ne_bytes()); } } } @@ -231,36 +207,12 @@ pub fn convert_32bit_rows( .zip(rgba_data.chunks_exact_mut(row_bytes)) .take(height) { - let src = &source[..row_bytes]; - let aligned = (row_bytes / 16) * 16; - let (src_aligned, src_rem) = src.split_at(aligned); - let (dst_aligned, dst_rem) = target.split_at_mut(aligned); - for (src, dst) in src_aligned.chunks_exact(16).zip(dst_aligned.chunks_exact_mut(16)) { - dst[0] = src[3]; - dst[1] = src[2]; - dst[2] = src[1]; - dst[3] = src[0]; - - dst[4] = src[7]; - dst[5] = src[6]; - dst[6] = src[5]; - dst[7] = src[4]; - - dst[8] = src[11]; - dst[9] = src[10]; - dst[10] = src[9]; - dst[11] = src[8]; - - dst[12] = src[15]; - dst[13] = src[14]; - dst[14] = src[13]; - dst[15] = src[12]; - } - for (src, dst) in src_rem.chunks_exact(4).zip(dst_rem.chunks_exact_mut(4)) { - dst[0] = src[3]; - dst[1] = src[2]; - dst[2] = src[1]; - dst[3] = src[0]; + for (src, dst) in source[..row_bytes] + .chunks_exact(4) + .zip(target.chunks_exact_mut(4)) + { + let pixel = u32::from_ne_bytes(src.try_into().unwrap()); + dst.copy_from_slice(&pixel.swap_bytes().to_ne_bytes()); } } } @@ -332,18 +284,19 @@ pub fn convert_nv12_planes( let cr_f = cr as f32 - 128.0; let r_chroma = 1.596027_f32 * cr_f; - let g_chroma = -0.391762_f32 * cb_f - 0.812968_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_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_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; @@ -352,6 +305,14 @@ pub fn convert_nv12_planes( } if x < width { + if x + 1 >= uv_row.len() { + tracing::warn!( + uv_index = x, + uv_row_len = uv_row.len(), + "UV index out of bounds during conversion", + ); + return None; + } let cb = uv_row[x]; let cr = uv_row[x + 1]; let (r, g, b) = ycbcr_to_rgb(y_row[x], cb, cr, range); @@ -365,7 +326,6 @@ pub fn convert_nv12_planes( Some(rgba_data) } - pub 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; @@ -461,4 +421,60 @@ 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 convert_32bit_rows_unaligned_widths() { + for width in [1, 2, 3, 5, 6, 7, 9, 15, 17] { + let stride = width * 4 + 8; + let height = 3; + let mut raw = vec![0u8; stride * height]; + for (i, byte) in raw.iter_mut().enumerate() { + *byte = (i as u8).wrapping_mul(13); + } + for order in [ + ChannelOrder::Rgba, + ChannelOrder::Bgra, + ChannelOrder::Argb, + ChannelOrder::Abgr, + ] { + let converted = convert_32bit_rows(&raw, stride, width, height, order).unwrap(); + assert_eq!(converted.len(), width * height * 4); + for y in 0..height { + for x in 0..width { + let src_idx = y * stride + x * 4; + let dst_idx = (y * width + x) * 4; + let px = &raw[src_idx..src_idx + 4]; + let expected = match order { + ChannelOrder::Rgba => [px[0], px[1], px[2], px[3]], + ChannelOrder::Bgra => [px[2], px[1], px[0], px[3]], + ChannelOrder::Argb => [px[1], px[2], px[3], px[0]], + ChannelOrder::Abgr => [px[3], px[2], px[1], px[0]], + }; + assert_eq!( + &converted[dst_idx..dst_idx + 4], + &expected, + "failed at width {width}, order {order:?}" + ); + } + } + } + } + } }