perf(desktop-gpui): AVX2/SSSE3 rgba to bgra unpad

.github/workflows/instant-reliability.yml

@@ -316,8 +316,11 @@ jobs:
         run: |
           set -euo pipefail
           sudo apt-get update -qq
-          # v4l2loopback needs videodev, which is in the extra kernel modules.
-          sudo apt-get install -y -qq "linux-modules-extra-$(uname -r)" v4l2loopback-dkms xvfb x11-utils xdotool \
+          if ! sudo modprobe videodev 2>/dev/null; then
+            sudo apt-get install -y -qq "linux-modules-extra-$(uname -r)"
+            sudo modprobe videodev
+          fi
+          sudo apt-get install -y -qq "linux-headers-$(uname -r)" v4l2loopback-dkms xvfb x11-utils xdotool \
             dbus ffmpeg pulseaudio pulseaudio-utils libasound2-plugins >/dev/null
           chmod +x "$RUNNER_TEMP/build/cap-gpui" "$RUNNER_TEMP/build/real-app-tests"
           sudo modprobe v4l2loopback video_nr=42 card_label="Cap Test Camera" exclusive_caps=1

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

@@ -4,18 +4,28 @@ use std::{hint::black_box, time::Instant};
 type Unpad = fn(&[u8], usize, usize, usize) -> Option<Vec<u8>>;
 
 fn main() {
+    #[cfg(target_arch = "x86_64")]
+    println!(
+        "arch={} avx2={} ssse3={}",
+        std::env::consts::ARCH,
+        std::arch::is_x86_feature_detected!("avx2"),
+        std::arch::is_x86_feature_detected!("ssse3"),
+    );
+    #[cfg(not(target_arch = "x86_64"))]
+    println!("arch={} avx2=false ssse3=false", std::env::consts::ARCH);
     println!("workload,sample,before_ns,after_ns");
     for (name, width, height) in [
         ("preview-1248x702", 1248usize, 702usize),
         ("1080p", 1920, 1080),
         ("4k", 3840, 2160),
     ] {
         let stride = (width * 4).next_multiple_of(256);
-        let src: Vec<u8> = (0..stride * height).map(|i| i as u8).collect();
+        let src: Vec<u8> = (0..stride * height).map(|i| (i * 31 + 9) as u8).collect();
         let before = baseline::unpad_rgba_to_bgra(&src, width, stride, height).unwrap();
         let after = candidate::unpad_rgba_to_bgra(&src, width, stride, height).unwrap();
         assert_eq!(before, after, "{name}");
-        for sample in 0..35 {
+        let mut samples = Vec::with_capacity(30);
+        for sample in 0..40 {
             let run = |f: Unpad| {
                 let start = Instant::now();
                 let _ = black_box(f(
@@ -33,14 +43,39 @@ fn main() {
                 let new = run(candidate::unpad_rgba_to_bgra);
                 (run(baseline::unpad_rgba_to_bgra), new)
             };
-            if sample >= 5 {
-                println!("{name},{},{old},{new}", sample - 5);
+            if sample >= 10 {
+                println!("{name},{},{old},{new}", sample - 10);
+                samples.push((old, new));
             }
         }
+        let (before_ns, after_ns) = median(&samples);
+        println!(
+            "{name},median,{before_ns},{after_ns},stride={stride},samples={}",
+            samples.len()
+        );
     }
 }
 
-// Inlined from main so before/after share one build and the same inputs.
+fn median(samples: &[(u128, u128)]) -> (u128, u128) {
+    fn mid(values: &mut [u128]) -> u128 {
+        values.sort_unstable();
+        let center = values.len() / 2;
+        if values.len().is_multiple_of(2) {
+            values[center - 1].midpoint(values[center])
+        } else {
+            values[center]
+        }
+    }
+    let before = mid(&mut samples
+        .iter()
+        .map(|(before, _)| *before)
+        .collect::<Vec<_>>());
+    let after = mid(&mut samples.iter().map(|(_, after)| *after).collect::<Vec<_>>());
+    (before, after)
+}
+
+// Scalar pointer loop this kernel used before the AVX2/SSSE3 dispatch.
+// Inlined so before/after share one build and the same inputs.
 mod baseline {
     pub fn unpad_rgba_to_bgra(
         data: &[u8],
@@ -57,21 +92,22 @@ mod baseline {
             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];
+        let mut tight = Vec::with_capacity(output_len);
+        let dest: *mut u8 = tight.as_mut_ptr();
+        unsafe {
+            for row in 0..rows {
+                let src = data.as_ptr().add(row * stride);
+                let dst = dest.add(row * row_bytes);
+                for col in 0..width {
+                    let s = src.add(col * 4);
+                    let d = dst.add(col * 4);
+                    *d = *s.add(2);
+                    *d.add(1) = *s.add(1);
+                    *d.add(2) = *s;
+                    *d.add(3) = *s.add(3);
+                }
             }
+            tight.set_len(output_len);
         }
         Some(tight)
     }

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

@@ -26,25 +26,145 @@ pub fn unpad_rgba_to_bgra(
         return Some(Vec::new());
     }
 
+    #[cfg(target_arch = "x86_64")]
+    {
+        // AVX2 implies SSSE3 on hardware; check the wider kernel first.
+        if is_x86_feature_detected!("avx2") {
+            // SAFETY: feature detected. Each source row has `stride` bytes
+            // (`stride * rows <= data.len()`), and `stride >= width * 4`.
+            return Some(unsafe { unpad_rgba_to_bgra_avx2(data, width, stride, rows, row_bytes) });
+        }
+        if is_x86_feature_detected!("ssse3") {
+            // SAFETY: same row bounds as the AVX2 path.
+            return Some(unsafe { unpad_rgba_to_bgra_ssse3(data, width, stride, rows, row_bytes) });
+        }
+    }
+
+    Some(unpad_rgba_to_bgra_scalar(
+        data, width, stride, rows, row_bytes,
+    ))
+}
+
+fn unpad_rgba_to_bgra_scalar(
+    data: &[u8],
+    width: usize,
+    stride: usize,
+    rows: usize,
+    row_bytes: usize,
+) -> Vec<u8> {
+    let output_len = row_bytes * rows;
     let mut tight = Vec::with_capacity(output_len);
     let dest: *mut u8 = tight.as_mut_ptr();
     // SAFETY: dest is unused spare capacity; every output byte is written before set_len.
     unsafe {
         for row in 0..rows {
+            swap_rgba_row_scalar(
+                data.as_ptr().add(row * stride),
+                dest.add(row * row_bytes),
+                0,
+                width,
+            );
+        }
+        tight.set_len(output_len);
+    }
+    tight
+}
+
+#[inline(always)]
+fn swap_rgba_row_scalar(src: *const u8, dst: *mut u8, start: usize, width: usize) {
+    for col in start..width {
+        // SAFETY: caller passes a pixel range inside both the source row and the dest row.
+        unsafe {
+            let s = src.add(col * 4);
+            let d = dst.add(col * 4);
+            *d = *s.add(2);
+            *d.add(1) = *s.add(1);
+            *d.add(2) = *s;
+            *d.add(3) = *s.add(3);
+        }
+    }
+}
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2")]
+unsafe fn unpad_rgba_to_bgra_avx2(
+    data: &[u8],
+    width: usize,
+    stride: usize,
+    rows: usize,
+    row_bytes: usize,
+) -> Vec<u8> {
+    use std::arch::x86_64::*;
+
+    let output_len = row_bytes * rows;
+    let mut tight = Vec::with_capacity(output_len);
+    let dest: *mut u8 = tight.as_mut_ptr();
+    // vpshufb indexes each 16-byte lane separately, so the mask repeats.
+    let shuffle = _mm256_setr_epi8(
+        2, 1, 0, 3, 6, 5, 4, 7, 10, 9, 8, 11, 14, 13, 12, 15, 2, 1, 0, 3, 6, 5, 4, 7, 10, 9, 8, 11,
+        14, 13, 12, 15,
+    );
+    let simd_pixels = width & !7;
+    let simd_bytes = simd_pixels * 4;
+    for row in 0..rows {
+        // SAFETY: row * stride + simd_bytes <= row * stride + row_bytes <= stride * rows,
+        // which the caller checked against data.len(). Stores cover exactly row_bytes
+        // once the scalar tail finishes the leftover pixels. set_len is below.
+        unsafe {
             let src = data.as_ptr().add(row * stride);
             let dst = dest.add(row * row_bytes);
-            for col in 0..width {
-                let s = src.add(col * 4);
-                let d = dst.add(col * 4);
-                *d = *s.add(2);
-                *d.add(1) = *s.add(1);
-                *d.add(2) = *s;
-                *d.add(3) = *s.add(3);
+            let mut offset = 0;
+            while offset < simd_bytes {
+                debug_assert!(offset + 32 <= row_bytes);
+                let loaded = _mm256_loadu_si256(src.add(offset).cast());
+                let shuffled = _mm256_shuffle_epi8(loaded, shuffle);
+                _mm256_storeu_si256(dst.add(offset).cast(), shuffled);
+                offset += 32;
             }
+            swap_rgba_row_scalar(src, dst, simd_pixels, width);
         }
-        tight.set_len(output_len);
     }
-    Some(tight)
+    // SAFETY: every output byte was written by the vector stores or the scalar tail.
+    unsafe { tight.set_len(output_len) };
+    tight
+}
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "ssse3")]
+unsafe fn unpad_rgba_to_bgra_ssse3(
+    data: &[u8],
+    width: usize,
+    stride: usize,
+    rows: usize,
+    row_bytes: usize,
+) -> Vec<u8> {
+    use std::arch::x86_64::*;
+
+    let output_len = row_bytes * rows;
+    let mut tight = Vec::with_capacity(output_len);
+    let dest: *mut u8 = tight.as_mut_ptr();
+    let shuffle = _mm_setr_epi8(2, 1, 0, 3, 6, 5, 4, 7, 10, 9, 8, 11, 14, 13, 12, 15);
+    let simd_pixels = width & !3;
+    let simd_bytes = simd_pixels * 4;
+    for row in 0..rows {
+        // SAFETY: same bound as the AVX2 kernel, with 16-byte groups instead of 32.
+        unsafe {
+            let src = data.as_ptr().add(row * stride);
+            let dst = dest.add(row * row_bytes);
+            let mut offset = 0;
+            while offset < simd_bytes {
+                debug_assert!(offset + 16 <= row_bytes);
+                let loaded = _mm_loadu_si128(src.add(offset).cast());
+                let shuffled = _mm_shuffle_epi8(loaded, shuffle);
+                _mm_storeu_si128(dst.add(offset).cast(), shuffled);
+                offset += 16;
+            }
+            swap_rgba_row_scalar(src, dst, simd_pixels, width);
+        }
+    }
+    // SAFETY: every output byte was written by the vector stores or the scalar tail.
+    unsafe { tight.set_len(output_len) };
+    tight
 }
 
 pub fn swap_red_blue(pixels: &mut [u8]) {
@@ -337,6 +457,108 @@ mod tests {
         assert_eq!(unpad_rgba_to_bgra(&[0; 16], 4, 0, 1), None);
     }
 
+    fn swapped_copy(data: &[u8], width: usize, stride: usize, rows: usize) -> Vec<u8> {
+        let row_bytes = width * 4;
+        let mut out = vec![0u8; row_bytes * rows];
+        for row in 0..rows {
+            for col in 0..width {
+                let s = row * stride + col * 4;
+                let d = row * row_bytes + col * 4;
+                out[d] = data[s + 2];
+                out[d + 1] = data[s + 1];
+                out[d + 2] = data[s];
+                out[d + 3] = data[s + 3];
+            }
+        }
+        out
+    }
+
+    fn assert_unpad_helpers(
+        data: &[u8],
+        width: usize,
+        stride: usize,
+        rows: usize,
+        expected: &[u8],
+    ) {
+        let row_bytes = width * 4;
+        assert_eq!(
+            unpad_rgba_to_bgra(data, width, stride, rows).unwrap(),
+            expected,
+            "dispatch width {width} stride {stride} rows {rows}"
+        );
+        assert_eq!(
+            unpad_rgba_to_bgra_scalar(data, width, stride, rows, row_bytes),
+            expected,
+            "scalar width {width} stride {stride} rows {rows}"
+        );
+        #[cfg(target_arch = "x86_64")]
+        {
+            if is_x86_feature_detected!("ssse3") {
+                let ssse3 =
+                    unsafe { unpad_rgba_to_bgra_ssse3(data, width, stride, rows, row_bytes) };
+                assert_eq!(
+                    ssse3, expected,
+                    "ssse3 width {width} stride {stride} rows {rows}"
+                );
+            }
+            if is_x86_feature_detected!("avx2") {
+                let avx2 = unsafe { unpad_rgba_to_bgra_avx2(data, width, stride, rows, row_bytes) };
+                assert_eq!(
+                    avx2, expected,
+                    "avx2 width {width} stride {stride} rows {rows}"
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn unpad_matches_an_independent_swap_across_tails_and_padding() {
+        let widths = [1usize, 3, 4, 7, 8, 9, 15, 16, 17, 31, 32, 33, 1080, 1248];
+        let rows = [1usize, 2, 3];
+        for width in widths {
+            for rows in rows {
+                let row_bytes = width * 4;
+                for stride in [row_bytes, row_bytes + 4, row_bytes.next_multiple_of(256)] {
+                    let mut data = vec![0xA5u8; stride * rows];
+                    for row in 0..rows {
+                        for byte in 0..row_bytes {
+                            data[row * stride + byte] = (row * 17 + byte * 3 + 1) as u8;
+                        }
+                    }
+                    let expected = swapped_copy(&data, width, stride, rows);
+                    assert_unpad_helpers(&data, width, stride, rows, &expected);
+                }
+            }
+        }
+    }
+
+    #[cfg(target_arch = "x86_64")]
+    #[test]
+    fn ssse3_and_avx2_unpad_match_scalar_on_the_preview_frame() {
+        let width = 1248usize;
+        let rows = 702usize;
+        let stride = (width * 4).next_multiple_of(256);
+        let mut data = vec![0xA5u8; stride * rows];
+        for row in 0..rows {
+            for byte in 0..width * 4 {
+                data[row * stride + byte] = (row.wrapping_mul(13) + byte) as u8;
+            }
+        }
+        let scalar = unpad_rgba_to_bgra_scalar(&data, width, stride, rows, width * 4);
+        if is_x86_feature_detected!("ssse3") {
+            let ssse3 = unsafe { unpad_rgba_to_bgra_ssse3(&data, width, stride, rows, width * 4) };
+            assert_eq!(ssse3, scalar);
+        }
+        if is_x86_feature_detected!("avx2") {
+            let avx2 = unsafe { unpad_rgba_to_bgra_avx2(&data, width, stride, rows, width * 4) };
+            assert_eq!(avx2, scalar);
+        }
+        assert_eq!(
+            unpad_rgba_to_bgra(&data, width, stride, rows).unwrap(),
+            scalar
+        );
+    }
+
     #[test]
     fn copies_rows_and_mirrors_each_row() {
         assert_eq!(