From ec63a2aef2d110e68cac964c29a644668a64b780 Mon Sep 17 00:00:00 2001 From: mihir Date: Sun, 18 Oct 2026 18:15:25 +0000 Subject: [PATCH] perf(desktop-gpui): speed up the screenshot-export flatten and JPEG path flatten_onto_white fast-paths fully-opaque pixels (skip) and fully-transparent pixels (fill white) instead of running three integer divisions per pixel, and hoists the per-pixel 255 * (255 - alpha) recompute. A new fused flatten_rgba_to_rgb_over_white does the Share-JPEG flatten plus the alpha strip in one pass with one allocation, replacing two full-frame passes and a 33 MB transient at 4K, and rgba_to_rgb now pre-sizes its output instead of growing it through flat_map + collect. Outputs are bit-identical: the exhaustive all-channel/all-alpha flatten test, the JPEG round-trip test and a new fused-vs-two-step test all pass unchanged. Benchmarked by kernels/benches/frame_export.rs (median of 25 samples after warmup, alternating order, same buffers): flatten opaque 1080p 4.56 ms -> 1.48 ms 3.09x flatten opaque 4k 18.41 ms -> 6.01 ms 3.06x flatten mixed 4k 18.30 ms -> 7.70 ms 2.38x rgba strip 4k 31.00 ms -> 19.44 ms 1.59x export flatten+strip 49.19 ms -> 15.20 ms 3.24x --- apps/desktop-gpui/kernels/Cargo.toml | 4 + .../kernels/benches/frame_export.rs | 133 ++++++++++++++++++ apps/desktop-gpui/kernels/src/codec.rs | 53 ++++++- apps/desktop-gpui/kernels/src/screenshot.rs | 16 ++- apps/desktop-gpui/src/screenshot_export.rs | 5 +- 5 files changed, 203 insertions(+), 8 deletions(-) create mode 100644 apps/desktop-gpui/kernels/benches/frame_export.rs diff --git a/apps/desktop-gpui/kernels/Cargo.toml b/apps/desktop-gpui/kernels/Cargo.toml index 7547d3518c2..fbfcf8da89a 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_export" +harness = false diff --git a/apps/desktop-gpui/kernels/benches/frame_export.rs b/apps/desktop-gpui/kernels/benches/frame_export.rs new file mode 100644 index 00000000000..57740639dd0 --- /dev/null +++ b/apps/desktop-gpui/kernels/benches/frame_export.rs @@ -0,0 +1,133 @@ +use cap_gpui_kernels::{codec as candidate_codec, screenshot as candidate}; +use std::{hint::black_box, time::Instant}; + +fn opaque_rgba(pixels: usize) -> Vec { + let mut rgba = Vec::with_capacity(pixels * 4); + for i in 0..pixels { + rgba.extend_from_slice(&[ + (i & 0xFF) as u8, + ((i >> 8) & 0xFF) as u8, + (i.wrapping_mul(17) & 0xFF) as u8, + 255, + ]); + } + rgba +} + +fn mixed_rgba(pixels: usize) -> Vec { + let mut rgba = Vec::with_capacity(pixels * 4); + for i in 0..pixels { + rgba.extend_from_slice(&[ + (i & 0xFF) as u8, + ((i >> 8) & 0xFF) as u8, + (i.wrapping_mul(17) & 0xFF) as u8, + if i % 7 == 0 { (i % 200) as u8 } else { 255 }, + ]); + } + rgba +} + +fn fused_two_step(rgba: &[u8]) -> Vec { + baseline::rgba_to_rgb(&baseline::flatten_onto_white(rgba)) +} + +fn run_flatten(name: &str, rgba: &[u8]) { + assert_eq!( + baseline::flatten_onto_white(rgba), + candidate::flatten_onto_white(rgba), + "{name}" + ); + emit( + name, + baseline::flatten_onto_white, + candidate::flatten_onto_white, + rgba, + ); +} + +fn run_strip(name: &str, rgba: &[u8]) { + assert_eq!( + baseline::rgba_to_rgb(rgba), + candidate_codec::rgba_to_rgb(rgba), + "{name}" + ); + emit( + name, + baseline::rgba_to_rgb, + candidate_codec::rgba_to_rgb, + rgba, + ); +} + +fn run_fused(name: &str, rgba: &[u8]) { + assert_eq!( + fused_two_step(rgba), + candidate_codec::flatten_rgba_to_rgb_over_white(rgba), + "{name}" + ); + emit( + name, + fused_two_step, + candidate_codec::flatten_rgba_to_rgb_over_white, + rgba, + ); +} + +fn emit(name: &str, before: fn(T) -> Vec, after: fn(T) -> Vec, input: T) +where + T: Copy, +{ + for sample in 0..30 { + let (old, new) = if sample % 2 == 0 { + (timed(before, input), timed(after, input)) + } else { + let new = timed(after, input); + let old = timed(before, input); + (old, new) + }; + if sample >= 5 { + println!("{name},{},{old},{new}", sample - 5); + } + } +} + +fn timed(f: fn(T) -> Vec, input: T) -> u128 +where + T: Copy, +{ + let start = Instant::now(); + black_box(f(black_box(input))); + start.elapsed().as_nanos() +} + +fn main() { + let hd = opaque_rgba(1920 * 1080); + let uhd_opaque = opaque_rgba(3840 * 2160); + let uhd_mixed = mixed_rgba(3840 * 2160); + println!("workload,sample,before_ns,after_ns"); + run_flatten("flatten-opaque-1080p", &hd); + run_flatten("flatten-opaque-4k", &uhd_opaque); + run_flatten("flatten-mixed-4k", &uhd_mixed); + run_strip("rgba-strip-4k", &uhd_opaque); + run_fused("export-fused-4k", &uhd_mixed); +} + +mod baseline { + pub fn rgba_to_rgb(rgba: &[u8]) -> Vec { + rgba.chunks_exact(4) + .flat_map(|px| [px[0], px[1], px[2]]) + .collect() + } + + pub fn flatten_onto_white(rgba: &[u8]) -> Vec { + let mut rgba = rgba.to_vec(); + for pixel in rgba.chunks_exact_mut(4) { + let alpha = u32::from(pixel[3]); + for channel in &mut pixel[..3] { + *channel = ((u32::from(*channel) * alpha + 255 * (255 - alpha) + 127) / 255) as u8; + } + pixel[3] = 255; + } + rgba + } +} diff --git a/apps/desktop-gpui/kernels/src/codec.rs b/apps/desktop-gpui/kernels/src/codec.rs index b184447c339..334373f0684 100644 --- a/apps/desktop-gpui/kernels/src/codec.rs +++ b/apps/desktop-gpui/kernels/src/codec.rs @@ -17,9 +17,35 @@ pub fn read_cursor_events( } pub fn rgba_to_rgb(rgba: &[u8]) -> Vec { - rgba.chunks_exact(4) - .flat_map(|px| [px[0], px[1], px[2]]) - .collect() + let len = rgba.len() / 4 * 3; + let mut rgb = vec![0u8; len]; + for (src, dst) in rgba.chunks_exact(4).zip(rgb.chunks_exact_mut(3)) { + dst[0] = src[0]; + dst[1] = src[1]; + dst[2] = src[2]; + } + rgb +} + +pub fn flatten_rgba_to_rgb_over_white(rgba: &[u8]) -> Vec { + let mut rgb = Vec::with_capacity(rgba.len() / 4 * 3); + for px in rgba.chunks_exact(4) { + let alpha = px[3]; + if alpha == 255 { + rgb.extend_from_slice(&[px[0], px[1], px[2]]); + } else if alpha == 0 { + rgb.extend_from_slice(&[255, 255, 255]); + } else { + let a = u32::from(alpha); + let inv = 255 - a; + rgb.extend_from_slice(&[ + ((u32::from(px[0]) * a + 255 * inv + 127) / 255) as u8, + ((u32::from(px[1]) * a + 255 * inv + 127) / 255) as u8, + ((u32::from(px[2]) * a + 255 * inv + 127) / 255) as u8, + ]); + } + } + rgb } pub fn encode_png_rgba(rgba: &[u8], width: u32, height: u32) -> image::ImageResult> { @@ -192,6 +218,27 @@ mod tests { ); } + #[test] + fn fused_flatten_to_rgb_matches_two_step() { + let rgba: Vec = (0..=255u8) + .flat_map(|alpha| { + (0..=255u8).flat_map(move |channel| { + [channel, 255 - channel, channel.wrapping_mul(17), alpha] + }) + }) + .collect(); + assert_eq!( + flatten_rgba_to_rgb_over_white(&rgba), + rgba_to_rgb(&crate::screenshot::flatten_onto_white(&rgba)) + ); + assert_eq!(flatten_rgba_to_rgb_over_white(&[]), Vec::::new()); + let trailing: [u8; 7] = [10, 20, 30, 40, 50, 60, 70]; + assert_eq!( + flatten_rgba_to_rgb_over_white(&trailing), + rgba_to_rgb(&crate::screenshot::flatten_onto_white(&trailing)) + ); + } + #[test] fn png_round_trips_rgba() { let rgba = [10, 20, 30, 40, 50, 60, 70, 255]; diff --git a/apps/desktop-gpui/kernels/src/screenshot.rs b/apps/desktop-gpui/kernels/src/screenshot.rs index 93f969b67f5..e6000dcd106 100644 --- a/apps/desktop-gpui/kernels/src/screenshot.rs +++ b/apps/desktop-gpui/kernels/src/screenshot.rs @@ -19,9 +19,21 @@ pub fn has_translucent_pixel(rgba: &[u8]) -> bool { pub fn flatten_onto_white(rgba: &[u8]) -> Vec { let mut rgba = rgba.to_vec(); for pixel in rgba.chunks_exact_mut(4) { - let alpha = u32::from(pixel[3]); + let alpha = pixel[3]; + if alpha == 255 { + continue; + } + if alpha == 0 { + pixel[0] = 255; + pixel[1] = 255; + pixel[2] = 255; + pixel[3] = 255; + continue; + } + let a = u32::from(alpha); + let inv = 255 - a; for channel in &mut pixel[..3] { - *channel = ((u32::from(*channel) * alpha + 255 * (255 - alpha) + 127) / 255) as u8; + *channel = ((u32::from(*channel) * a + 255 * inv + 127) / 255) as u8; } pixel[3] = 255; } diff --git a/apps/desktop-gpui/src/screenshot_export.rs b/apps/desktop-gpui/src/screenshot_export.rs index 8e98a379c59..5533f188441 100644 --- a/apps/desktop-gpui/src/screenshot_export.rs +++ b/apps/desktop-gpui/src/screenshot_export.rs @@ -238,9 +238,8 @@ fn encode_png(out: &Composited) -> Result, String> { /// share flow only ever picks it when the output is fully opaque, and the /// flatten here is the same belt the browser's own encode wears. fn encode_jpeg(out: &Composited) -> Result, String> { - let flat = flatten_onto_white(out); - let rgb = codec::rgba_to_rgb(&flat.rgba); - codec::encode_jpeg_rgb(&rgb, flat.width, flat.height, 90) + let rgb = codec::flatten_rgba_to_rgb_over_white(&out.rgba); + codec::encode_jpeg_rgb(&rgb, out.width, out.height, 90) .map_err(|e| format!("Failed to encode screenshot export: {e}")) }