perf(gpui-kernels): specialise the blit for uniform backdrops

apps/desktop-gpui/kernels/benches/screenshot_blit.rs

@@ -3,6 +3,15 @@ use std::{hint::black_box, time::Instant};
 
 type Blit = fn(&mut [u8], (u32, u32), &[u8], (u32, u32), (i64, i64));
 
+// The export canvas is cleared to zero for a transparent background and
+// filled with white otherwise, so those two are the shipping backdrops.
+#[derive(Clone, Copy)]
+enum Backdrop {
+    Varying,
+    Transparent,
+    OpaqueWhite,
+}
+
 #[derive(Clone, Copy)]
 enum Alpha {
     Solid(u8),
@@ -13,20 +22,57 @@ enum Alpha {
 
 fn main() {
     use Alpha::{Edge, Mixed, Solid, Sparse};
+    use Backdrop::{OpaqueWhite, Transparent, Varying};
 
-    println!("workload,sample,before_ns,after_ns");
-    for (name, size, offset, alpha) in [
-        ("opaque-1080p", (1920u32, 1080u32), (0, 0), Solid(255)),
-        ("opaque-4k", (3840, 2160), (0, 0), Solid(255)),
-        ("opaque-portrait", (2160, 3840), (0, 0), Solid(255)),
-        ("opaque-expanded", (3840, 2160), (100, 80), Solid(255)),
-        ("opaque-clipped", (3840, 2160), (-2000, -1200), Solid(255)),
-        ("translucent-4k", (3840, 2160), (0, 0), Solid(128)),
-        ("edge-alpha-4k", (3840, 2160), (0, 0), Edge),
-        ("sparse-alpha-4k", (3840, 2160), (0, 0), Sparse),
-        ("mixed-4k", (3840, 2160), (0, 0), Mixed),
-        ("transparent-4k", (3840, 2160), (0, 0), Solid(0)),
-        ("outside-4k", (3840, 2160), (4000, 0), Solid(255)),
+    println!("workload,backdrop,sample,before_ns,after_ns");
+    for (name, size, offset, alpha, backdrop) in [
+        (
+            "opaque-1080p",
+            (1920u32, 1080u32),
+            (0, 0),
+            Solid(255),
+            Varying,
+        ),
+        ("opaque-4k", (3840, 2160), (0, 0), Solid(255), Varying),
+        ("opaque-portrait", (2160, 3840), (0, 0), Solid(255), Varying),
+        (
+            "opaque-expanded",
+            (3840, 2160),
+            (100, 80),
+            Solid(255),
+            Varying,
+        ),
+        (
+            "opaque-clipped",
+            (3840, 2160),
+            (-2000, -1200),
+            Solid(255),
+            Varying,
+        ),
+        ("translucent-4k", (3840, 2160), (0, 0), Solid(128), Varying),
+        ("edge-alpha-4k", (3840, 2160), (0, 0), Edge, Varying),
+        ("sparse-alpha-4k", (3840, 2160), (0, 0), Sparse, Varying),
+        ("mixed-4k", (3840, 2160), (0, 0), Mixed, Varying),
+        ("transparent-4k", (3840, 2160), (0, 0), Solid(0), Varying),
+        ("outside-4k", (3840, 2160), (4000, 0), Solid(255), Varying),
+        (
+            "translucent-4k",
+            (3840, 2160),
+            (0, 0),
+            Solid(128),
+            Transparent,
+        ),
+        (
+            "translucent-4k",
+            (3840, 2160),
+            (0, 0),
+            Solid(128),
+            OpaqueWhite,
+        ),
+        ("mixed-4k", (3840, 2160), (0, 0), Mixed, Transparent),
+        ("mixed-4k", (3840, 2160), (0, 0), Mixed, OpaqueWhite),
+        ("sparse-alpha-4k", (3840, 2160), (0, 0), Sparse, OpaqueWhite),
+        ("edge-alpha-4k", (3840, 2160), (0, 0), Edge, OpaqueWhite),
     ] {
         let dst_size = if name == "opaque-expanded" {
             (size.0 + 200, size.1 + 160)
@@ -50,7 +96,11 @@ fn main() {
             })
             .collect();
         let initial: Vec<u8> = (0..dst_size.0 * dst_size.1)
-            .flat_map(|i| [53, 71, 89, (i % 256) as u8])
+            .flat_map(|i| match backdrop {
+                Varying => [53, 71, 89, (i % 256) as u8],
+                Transparent => [0, 0, 0, 0],
+                OpaqueWhite => [255, 255, 255, 255],
+            })
             .collect();
         let mut before = initial.clone();
         let mut after = initial.clone();
@@ -80,7 +130,12 @@ fn main() {
             };
             black_box((&before, &after));
             if sample >= 5 {
-                println!("{name},{},{old},{new}", sample - 5);
+                let backdrop = match backdrop {
+                    Varying => "varying",
+                    Transparent => "transparent",
+                    OpaqueWhite => "opaque-white",
+                };
+                println!("{name},{backdrop},{},{old},{new}", sample - 5);
             }
         }
     }

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

@@ -51,6 +51,20 @@ fn blend_pixel_over(dst: &mut [u8], src: [u8; 4]) {
     dst[3] = (oa * 255.).round() as u8;
 }
 
+// Source-over onto an opaque destination. `oa` is exactly 1.0 for every
+// 8-bit source alpha, so the divide drops out and the result is the integer
+// rounding of `(s*a + d*(255-a)) / 255`, which is bit-identical to the float
+// path for all 16.6M (src, dst, alpha) combinations.
+fn blend_pixel_over_opaque(dst: &mut [u8], src: [u8; 4]) {
+    let alpha = u32::from(src[3]);
+    let inverse = 255 - alpha;
+    for channel in 0..3 {
+        dst[channel] = ((u32::from(src[channel]) * alpha + u32::from(dst[channel]) * inverse + 127)
+            / 255) as u8;
+    }
+    dst[3] = 255;
+}
+
 pub fn blit_over(canvas: &mut [u8], size: (u32, u32), region: &image::RgbaImage, at: (u32, u32)) {
     blit_over_offset(
         canvas,
@@ -92,6 +106,78 @@ pub fn blit_over_offset(
         if is_opaque(source) {
             target.copy_from_slice(source);
         } else {
+            blend_row_over(target, source);
+        }
+    }
+}
+
+enum Backdrop {
+    Transparent,
+    Opaque,
+    Mixed,
+}
+
+// One pass over the destination alpha bytes instead of two predicate scans.
+fn classify_backdrop(rgba: &[u8]) -> Backdrop {
+    const ALPHA_MASK: u128 = 0xff000000ff000000ff000000ff000000;
+    let (mut transparent, mut opaque) = (true, true);
+    let mut blocks = rgba.chunks_exact(16);
+    for block in blocks.by_ref() {
+        let alpha = u128::from_le_bytes(block.try_into().unwrap()) & ALPHA_MASK;
+        transparent &= alpha == 0;
+        opaque &= alpha == ALPHA_MASK;
+        if !transparent && !opaque {
+            return Backdrop::Mixed;
+        }
+    }
+    for &alpha in blocks.remainder().iter().skip(3).step_by(4) {
+        transparent &= alpha == 0;
+        opaque &= alpha == 255;
+    }
+    match (transparent, opaque) {
+        (true, _) => Backdrop::Transparent,
+        (_, true) => Backdrop::Opaque,
+        _ => Backdrop::Mixed,
+    }
+}
+
+#[inline(never)]
+fn blend_row_over(target: &mut [u8], source: &[u8]) {
+    match classify_backdrop(target) {
+        Backdrop::Transparent => {
+            // Over a fully transparent row the destination contributes
+            // nothing, so each pixel is either the source or left alone.
+            for (target, source) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) {
+                if source[3] != 0 {
+                    target.copy_from_slice(source);
+                }
+            }
+        }
+        Backdrop::Opaque => {
+            // Blocks of four pixels: a fully opaque block is a straight copy,
+            // anything else takes the branchless integer blend so the common
+            // run of opaque pixels never pays for per-pixel branching.
+            const ALPHA_MASK: u128 = 0xff000000ff000000ff000000ff000000;
+            let mut targets = target.chunks_exact_mut(16);
+            let mut sources = source.chunks_exact(16);
+            for (target, source) in targets.by_ref().zip(sources.by_ref()) {
+                if u128::from_le_bytes(source.try_into().unwrap()) & ALPHA_MASK == ALPHA_MASK {
+                    target.copy_from_slice(source);
+                    continue;
+                }
+                for (target, source) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) {
+                    blend_pixel_over_opaque(target, source.try_into().unwrap());
+                }
+            }
+            for (target, source) in targets
+                .into_remainder()
+                .chunks_exact_mut(4)
+                .zip(sources.remainder().chunks_exact(4))
+            {
+                blend_pixel_over_opaque(target, source.try_into().unwrap());
+            }
+        }
+        Backdrop::Mixed => {
             for (target, source) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) {
                 blend_pixel_over(target, source.try_into().unwrap());
             }
@@ -240,6 +326,56 @@ mod tests {
         }
     }
 
+    #[test]
+    fn uniform_backdrops_match_scalar_source_over() {
+        // Widths either side of the four-pixel block stride the opaque
+        // backdrop path copies through.
+        for width in [1, 3, 4, 5, 7, 8, 15, 16, 17, 31, 32, 33] {
+            let size = (width, 3);
+            let src: Vec<u8> = (0..width * 3)
+                .flat_map(|i| {
+                    [
+                        i as u8,
+                        (i * 17) as u8,
+                        (i * 29) as u8,
+                        match i % 5 {
+                            0 => 0,
+                            1 => 255,
+                            2 => 128,
+                            3 => 1,
+                            _ => 254,
+                        },
+                    ]
+                })
+                .collect();
+            // The export canvas is either cleared to zero or filled with
+            // white, so both uniform backdrops are the real-world cases.
+            for backdrop in [
+                [0, 0, 0, 0],
+                [255, 255, 255, 255],
+                [11, 22, 33, 255],
+                [11, 22, 33, 0],
+                [11, 22, 33, 128],
+            ] {
+                for stray in [None, Some(0usize), Some(1), Some(4), Some(7)] {
+                    let mut initial: Vec<u8> = (0..width * 3).flat_map(|_| backdrop).collect();
+                    // A single differing alpha must drop the row out of the
+                    // uniform fast paths and back onto the general blend.
+                    if let Some(index) = stray
+                        && let Some(pixel) = initial.chunks_exact_mut(4).nth(index)
+                    {
+                        pixel[3] = 97;
+                    }
+                    let mut expected = initial.clone();
+                    let mut actual = initial;
+                    reference_blit(&mut expected, size, &src, size, (0, 0));
+                    blit_over_offset(&mut actual, size, &src, size, (0, 0));
+                    assert_eq!(actual, expected, "{width}, {backdrop:?}, {stray:?}");
+                }
+            }
+        }
+    }
+
     #[test]
     fn clipped_blits_match_scalar_source_over() {
         for width in [0, 1, 3, 4, 5, 15, 16, 17, 31, 32, 33] {