Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kyzer1023 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
10c0cc2
perf(desktop-gpui): AVX2 rgba to bgra unpad
kyzer1023 Oct 26, 2026
d59657d
test(desktop-gpui): cover ssse3 unpad tail cases
kyzer1023 Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 55 additions & 19 deletions apps/desktop-gpui/kernels/benches/frame_unpad.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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],
Expand All @@ -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)
}
Expand Down
240 changes: 231 additions & 9 deletions apps/desktop-gpui/kernels/src/frame.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]) {
Expand Down Expand Up @@ -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);
Comment on lines +548 to +550

@greptile-apps greptile-apps Bot Oct 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 SSSE3 edge cases go unchecked

On an AVX2 machine, the odd-width and padding test always selects AVX2. The only direct SSSE3 test uses width 1248, which is divisible by four, so it never checks where SSSE3 stops and the scalar tail starts. This leaves fallback boundary mistakes undetected on those machines.

Run the existing width-and-stride cases directly through each supported helper, keeping the feature checks.

Prompt To Fix With AI 
This is a comment left during a code review.
Path: apps/desktop-gpui/kernels/src/frame.rs
Line: 514-516

Comment:
**SSSE3 edge cases go unchecked**

On an AVX2 machine, the odd-width and padding test always selects AVX2. The only direct SSSE3 test uses width 1248, which is divisible by four, so it never checks where SSSE3 stops and the scalar tail starts. This leaves fallback boundary mistakes undetected on those machines.

Run the existing width-and-stride cases directly through each supported helper, keeping the feature checks.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

}
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!(
Expand Down
Loading