perf(desktop-gpui): accelerate screenshot-export and frame kernels

apps/desktop-gpui/kernels/Cargo.toml

@@ -26,3 +26,7 @@ harness = false
 [[bench]]
 name = "png_encode"
 harness = false
+
+[[bench]]
+name = "frame_export"
+harness = false

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<u8> {
+    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<u8> {
+    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<u8> {
+    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<T>(name: &str, before: fn(T) -> Vec<u8>, after: fn(T) -> Vec<u8>, 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<T>(f: fn(T) -> Vec<u8>, 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<u8> {
+        rgba.chunks_exact(4)
+            .flat_map(|px| [px[0], px[1], px[2]])
+            .collect()
+    }
+
+    pub fn flatten_onto_white(rgba: &[u8]) -> Vec<u8> {
+        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
+    }
+}

apps/desktop-gpui/kernels/src/codec.rs

@@ -17,9 +17,35 @@ pub fn read_cursor_events(
 }
 
 pub fn rgba_to_rgb(rgba: &[u8]) -> Vec<u8> {
-    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<u8> {
+    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<Vec<u8>> {
@@ -192,6 +218,27 @@ mod tests {
         );
     }
 
+    #[test]
+    fn fused_flatten_to_rgb_matches_two_step() {
+        let rgba: Vec<u8> = (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::<u8>::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];

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<u8> {
     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;
     }

apps/desktop-gpui/src/screenshot_export.rs

@@ -238,9 +238,8 @@ fn encode_png(out: &Composited) -> Result<Vec<u8>, 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<Vec<u8>, 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}"))
 }