use super::{OverlayImage, PlaneView, PooledRect, SampleU16Le, SampleU8};
use core::arch::x86_64::*;
use std::sync::OnceLock;
pub(super) fn enabled() -> bool {
static AVX2: OnceLock<bool> = OnceLock::new();
*AVX2.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
}
pub(super) unsafe fn blend_direct_u8(
plane: &mut PlaneView<'_>,
image: &OverlayImage<'_>,
src: u32,
alpha: u32,
(x0, y0): (usize, usize),
(xm0, ym0): (usize, usize),
(w, h): (usize, usize),
) {
debug_assert_eq!(plane.pixel_step, 1);
for row in 0..h {
let mask_start = (ym0 + row) * image.stride + xm0;
let mask_row = &image.bitmap[mask_start..mask_start + w];
let dst_start = (y0 + row) * plane.linesize + x0;
let dst_row = &mut plane.data[dst_start..dst_start + w];
blend_row_u8(dst_row, mask_row, src, alpha);
}
}
pub(super) unsafe fn blend_direct_u16(
plane: &mut PlaneView<'_>,
image: &OverlayImage<'_>,
src: u32,
alpha: u32,
(x0, y0): (usize, usize),
(xm0, ym0): (usize, usize),
(w, h): (usize, usize),
) {
debug_assert_eq!(plane.pixel_step, 2);
for row in 0..h {
let mask_start = (ym0 + row) * image.stride + xm0;
let mask_row = &image.bitmap[mask_start..mask_start + w];
let dst_start = (y0 + row) * plane.linesize + x0 * 2;
let dst_row = &mut plane.data[dst_start..dst_start + w * 2];
blend_row_u16(dst_row, mask_row, src, alpha);
}
}
pub(super) unsafe fn apply_sums_u8(
plane: &mut PlaneView<'_>,
sums: &[u16],
rect: PooledRect,
src: u32,
alpha: u32,
) {
if alpha == 0 {
return;
}
for row in 0..rect.ph {
let sums_row = &sums[row * rect.pw..(row + 1) * rect.pw];
let dst_start = (rect.py0 + row) * plane.linesize + rect.px0;
let dst_row = &mut plane.data[dst_start..dst_start + rect.pw];
apply_row_u8(dst_row, sums_row, src, alpha, rect.shift);
}
}
pub(super) unsafe fn apply_sums_u16(
plane: &mut PlaneView<'_>,
sums: &[u16],
rect: PooledRect,
src: u32,
alpha: u32,
) {
if alpha == 0 {
return;
}
for row in 0..rect.ph {
let sums_row = &sums[row * rect.pw..(row + 1) * rect.pw];
let dst_start = (rect.py0 + row) * plane.linesize + rect.px0 * 2;
let dst_row = &mut plane.data[dst_start..dst_start + rect.pw * 2];
apply_row_u16(dst_row, sums_row, src, alpha, rect.shift);
}
}
pub(super) unsafe fn pair_sums_one_row(span: &mut [u16], top: &[u8]) {
pair_sums_row::<false>(span, top, top);
}
pub(super) unsafe fn pair_sums_two_rows(span: &mut [u16], top: &[u8], bot: &[u8]) {
pair_sums_row::<true>(span, top, bot);
}
#[target_feature(enable = "avx2")]
unsafe fn blend_row_u8(dst: &mut [u8], mask: &[u8], src: u32, alpha: u32) {
debug_assert_eq!(dst.len(), mask.len());
let alpha_v = _mm256_set1_epi32(alpha as i32);
let src_v = _mm256_set1_epi32(src as i32);
let regroup = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7);
let n = mask.len();
let mut x = 0usize;
while x + 32 <= n {
let m = _mm256_loadu_si256(mask.as_ptr().add(x).cast());
if _mm256_testz_si256(m, m) == 0 {
blend32_u8(dst.as_mut_ptr().add(x), m, alpha_v, src_v, regroup);
}
x += 32;
}
if !mask[x..].iter().all(|&m| m == 0) {
super::blend_span::<SampleU8>(&mut dst[x..], &mask[x..], src, alpha, 1);
}
}
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn blend32_u8(
dst: *mut u8,
mask: __m256i,
alpha_v: __m256i,
src_v: __m256i,
regroup: __m256i,
) {
let m_lo = _mm256_castsi256_si128(mask);
let m_hi = _mm256_extracti128_si256::<1>(mask);
let d = _mm256_loadu_si256(dst.cast());
let d_lo = _mm256_castsi256_si128(d);
let d_hi = _mm256_extracti128_si256::<1>(d);
let b0 = blend8_u8(_mm256_cvtepu8_epi32(m_lo), d_lo, alpha_v, src_v);
let b1 = blend8_u8(
_mm256_cvtepu8_epi32(_mm_srli_si128::<8>(m_lo)),
_mm_srli_si128::<8>(d_lo),
alpha_v,
src_v,
);
let b2 = blend8_u8(_mm256_cvtepu8_epi32(m_hi), d_hi, alpha_v, src_v);
let b3 = blend8_u8(
_mm256_cvtepu8_epi32(_mm_srli_si128::<8>(m_hi)),
_mm_srli_si128::<8>(d_hi),
alpha_v,
src_v,
);
let p01 = _mm256_packus_epi32(b0, b1);
let p23 = _mm256_packus_epi32(b2, b3);
let bytes = _mm256_packus_epi16(p01, p23);
let bytes = _mm256_permutevar8x32_epi32(bytes, regroup);
_mm256_storeu_si256(dst.cast(), bytes);
}
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn blend8_u8(coverage: __m256i, dst8: __m128i, alpha_v: __m256i, src_v: __m256i) -> __m256i {
let a = _mm256_mullo_epi32(coverage, alpha_v);
let inv = _mm256_sub_epi32(_mm256_set1_epi32(0x0101_0101), a);
let d = _mm256_cvtepu8_epi32(dst8);
let num = _mm256_add_epi32(_mm256_mullo_epi32(inv, d), _mm256_mullo_epi32(a, src_v));
_mm256_srli_epi32::<24>(num)
}
#[target_feature(enable = "avx2")]
unsafe fn blend_row_u16(dst: &mut [u8], mask: &[u8], src: u32, alpha: u32) {
debug_assert_eq!(dst.len(), mask.len() * 2);
let alpha_v = _mm256_set1_epi32(alpha as i32);
let src_v = _mm256_set1_epi32(src as i32);
let n = mask.len();
let mut x = 0usize;
while x + 16 <= n {
let m = _mm_loadu_si128(mask.as_ptr().add(x).cast());
if _mm_testz_si128(m, m) == 0 {
let base = dst.as_mut_ptr().add(2 * x);
blend8_u16(base, _mm256_cvtepu8_epi32(m), alpha_v, src_v);
blend8_u16(
base.add(16),
_mm256_cvtepu8_epi32(_mm_srli_si128::<8>(m)),
alpha_v,
src_v,
);
}
x += 16;
}
if !mask[x..].iter().all(|&m| m == 0) {
super::blend_span::<SampleU16Le>(&mut dst[2 * x..], &mask[x..], src, alpha, 2);
}
}
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn blend8_u16(dst: *mut u8, coverage: __m256i, alpha_v: __m256i, src_v: __m256i) {
let a = _mm256_mullo_epi32(coverage, alpha_v);
let inv = _mm256_sub_epi32(_mm256_set1_epi32(0x0001_0001), a);
let d = _mm256_cvtepu16_epi32(_mm_loadu_si128(dst.cast()));
let num = _mm256_add_epi32(_mm256_mullo_epi32(inv, d), _mm256_mullo_epi32(a, src_v));
let b = _mm256_srli_epi32::<16>(num);
let packed = _mm256_packus_epi32(b, b);
let packed = _mm256_permute4x64_epi64::<0b0000_1000>(packed);
_mm_storeu_si128(dst.cast(), _mm256_castsi256_si128(packed));
}
#[target_feature(enable = "avx2")]
unsafe fn apply_row_u8(dst: &mut [u8], sums: &[u16], src: u32, alpha: u32, shift: u32) {
debug_assert_eq!(dst.len(), sums.len());
let alpha_v = _mm256_set1_epi32(alpha as i32);
let src_v = _mm256_set1_epi32(src as i32);
let shift_v = _mm_cvtsi32_si128(shift as i32);
let regroup = _mm256_setr_epi32(0, 4, 1, 5, 0, 0, 0, 0);
let n = sums.len();
let mut x = 0usize;
while x + 16 <= n {
let s = _mm256_loadu_si256(sums.as_ptr().add(x).cast());
if _mm256_testz_si256(s, s) == 0 {
let cov = _mm256_srl_epi16(s, shift_v);
let c0 = _mm256_cvtepu16_epi32(_mm256_castsi256_si128(cov));
let c1 = _mm256_cvtepu16_epi32(_mm256_extracti128_si256::<1>(cov));
let d = _mm_loadu_si128(dst.as_ptr().add(x).cast());
let b0 = blend8_u8(c0, d, alpha_v, src_v);
let b1 = blend8_u8(c1, _mm_srli_si128::<8>(d), alpha_v, src_v);
let p = _mm256_packus_epi32(b0, b1);
let bytes = _mm256_packus_epi16(p, p);
let bytes = _mm256_permutevar8x32_epi32(bytes, regroup);
_mm_storeu_si128(
dst.as_mut_ptr().add(x).cast(),
_mm256_castsi256_si128(bytes),
);
}
x += 16;
}
if !sums[x..].iter().all(|&s| s == 0) {
super::apply_sums_span::<SampleU8>(&mut dst[x..], &sums[x..], src, alpha, shift, 1);
}
}
#[target_feature(enable = "avx2")]
unsafe fn apply_row_u16(dst: &mut [u8], sums: &[u16], src: u32, alpha: u32, shift: u32) {
debug_assert_eq!(dst.len(), sums.len() * 2);
let alpha_v = _mm256_set1_epi32(alpha as i32);
let src_v = _mm256_set1_epi32(src as i32);
let shift_v = _mm_cvtsi32_si128(shift as i32);
let n = sums.len();
let mut x = 0usize;
while x + 16 <= n {
let s = _mm256_loadu_si256(sums.as_ptr().add(x).cast());
if _mm256_testz_si256(s, s) == 0 {
let cov = _mm256_srl_epi16(s, shift_v);
let base = dst.as_mut_ptr().add(2 * x);
blend8_u16(
base,
_mm256_cvtepu16_epi32(_mm256_castsi256_si128(cov)),
alpha_v,
src_v,
);
blend8_u16(
base.add(16),
_mm256_cvtepu16_epi32(_mm256_extracti128_si256::<1>(cov)),
alpha_v,
src_v,
);
}
x += 16;
}
if !sums[x..].iter().all(|&s| s == 0) {
super::apply_sums_span::<SampleU16Le>(&mut dst[2 * x..], &sums[x..], src, alpha, shift, 2);
}
}
#[target_feature(enable = "avx2")]
unsafe fn pair_sums_row<const TWO_ROWS: bool>(span: &mut [u16], top: &[u8], bot: &[u8]) {
debug_assert_eq!(top.len(), span.len() * 2);
if TWO_ROWS {
debug_assert_eq!(bot.len(), span.len() * 2);
}
let ones = _mm256_set1_epi8(1);
let n = span.len();
let mut x = 0usize;
while x + 16 <= n {
let t = _mm256_loadu_si256(top.as_ptr().add(2 * x).cast());
let mut s = _mm256_maddubs_epi16(t, ones);
if TWO_ROWS {
let b = _mm256_loadu_si256(bot.as_ptr().add(2 * x).cast());
s = _mm256_add_epi16(s, _mm256_maddubs_epi16(b, ones));
}
_mm256_storeu_si256(span.as_mut_ptr().add(x).cast(), s);
x += 16;
}
for i in x..n {
let mut sum = u16::from(top[2 * i]) + u16::from(top[2 * i + 1]);
if TWO_ROWS {
sum += u16::from(bot[2 * i]) + u16::from(bot[2 * i + 1]);
}
span[i] = sum;
}
}
#[cfg(test)]
mod tests {
use super::super::{apply_sums, blend_direct, SampleFormat};
use super::*;
fn lcg_bytes(len: usize, seed: &mut u64) -> Vec<u8> {
(0..len)
.map(|_| {
*seed = seed
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
(*seed >> 33) as u8
})
.collect()
}
fn mask_variants(w: usize, seed: &mut u64) -> Vec<Vec<u8>> {
let ramp: Vec<u8> = (0..w).map(|i| (i % 256) as u8).collect();
let mut sparse = lcg_bytes(w, seed);
for (i, byte) in sparse.iter_mut().enumerate() {
if (i / 24) % 2 == 0 {
*byte = 0;
}
}
vec![
vec![0u8; w],
vec![255u8; w],
ramp,
sparse,
lcg_bytes(w, seed),
]
}
fn opacity_alphas(sample: SampleFormat) -> Vec<u32> {
[0u8, 1, 127, 128, 254, 255]
.into_iter()
.map(|o| sample.alpha_fixed(o))
.filter(|&a| a != 0) .collect()
}
#[test]
fn direct_rows_match_scalar_bitexact() {
if !enabled() {
eprintln!("skipping: AVX2 unavailable on this CPU");
return;
}
let mut seed = 0x00C0_FFEE_1234_5678u64;
for w in (1..=67).chain([256]) {
for mask in mask_variants(w, &mut seed) {
let image = OverlayImage {
w,
h: 1,
stride: w,
bitmap: &mask,
color: 0xFFFFFF00,
dst_x: 0,
dst_y: 0,
};
for alpha in opacity_alphas(SampleFormat::U8) {
for src in [0u32, 16, 128, 235, 255] {
for base in [vec![0u8; w], vec![255u8; w], lcg_bytes(w, &mut seed)] {
let mut expected = base.clone();
blend_direct::<SampleU8>(
&mut PlaneView {
data: &mut expected,
linesize: w,
pixel_step: 1,
},
&image,
src,
alpha,
(0, 0),
(0, 0),
(w, 1),
);
let mut got = base.clone();
unsafe {
blend_direct_u8(
&mut PlaneView {
data: &mut got,
linesize: w,
pixel_step: 1,
},
&image,
src,
alpha,
(0, 0),
(0, 0),
(w, 1),
)
};
assert_eq!(got, expected, "u8 w={w} alpha={alpha} src={src}");
}
}
}
for alpha in opacity_alphas(SampleFormat::U16Le) {
for src in [0u32, 514, 943, 60395, 65535] {
for base in [
vec![0u8; w * 2],
vec![255u8; w * 2],
lcg_bytes(w * 2, &mut seed),
] {
let mut expected = base.clone();
blend_direct::<SampleU16Le>(
&mut PlaneView {
data: &mut expected,
linesize: w * 2,
pixel_step: 2,
},
&image,
src,
alpha,
(0, 0),
(0, 0),
(w, 1),
);
let mut got = base.clone();
unsafe {
blend_direct_u16(
&mut PlaneView {
data: &mut got,
linesize: w * 2,
pixel_step: 2,
},
&image,
src,
alpha,
(0, 0),
(0, 0),
(w, 1),
)
};
assert_eq!(got, expected, "u16 w={w} alpha={alpha} src={src}");
}
}
}
}
}
}
#[test]
fn direct_multi_row_geometry_matches_scalar() {
if !enabled() {
eprintln!("skipping: AVX2 unavailable on this CPU");
return;
}
let mut seed = 0xDEAD_BEEF_0BAD_F00Du64;
let (w, h, stride) = (61usize, 9usize, 71usize);
let bitmap = lcg_bytes(stride * (h - 1) + w, &mut seed);
let image = OverlayImage {
w,
h,
stride,
bitmap: &bitmap,
color: 0xFFFFFF00,
dst_x: 0,
dst_y: 0,
};
let coords = ((3usize, 2usize), (1usize, 1usize), (w - 4, h - 3));
for (bytes, src, alpha) in [
(1usize, 235u32, SampleFormat::U8.alpha_fixed(200)),
(2, 943, SampleFormat::U16Le.alpha_fixed(200)),
] {
let linesize = (w + 8) * bytes;
let base = lcg_bytes(linesize * (h + 4), &mut seed);
let mut expected = base.clone();
let mut got = base.clone();
if bytes == 1 {
blend_direct::<SampleU8>(
&mut PlaneView {
data: &mut expected,
linesize,
pixel_step: 1,
},
&image,
src,
alpha,
coords.0,
coords.1,
coords.2,
);
unsafe {
blend_direct_u8(
&mut PlaneView {
data: &mut got,
linesize,
pixel_step: 1,
},
&image,
src,
alpha,
coords.0,
coords.1,
coords.2,
)
};
} else {
blend_direct::<SampleU16Le>(
&mut PlaneView {
data: &mut expected,
linesize,
pixel_step: 2,
},
&image,
src,
alpha,
coords.0,
coords.1,
coords.2,
);
unsafe {
blend_direct_u16(
&mut PlaneView {
data: &mut got,
linesize,
pixel_step: 2,
},
&image,
src,
alpha,
coords.0,
coords.1,
coords.2,
)
};
}
assert_eq!(got, expected, "bytes={bytes}");
}
}
#[test]
fn pair_sums_match_scalar_bitexact() {
if !enabled() {
eprintln!("skipping: AVX2 unavailable on this CPU");
return;
}
let mut seed = 0x5EED_5EED_5EED_5EEDu64;
for n in (0..=67).chain([256]) {
let mut tops = mask_variants((2 * n).max(1), &mut seed);
for top in &mut tops {
top.truncate(2 * n);
}
for top in tops {
let bot = lcg_bytes(2 * n, &mut seed);
for two_rows in [false, true] {
let mut expected = vec![0u16; n];
for (i, sum) in expected.iter_mut().enumerate() {
let mut s = u16::from(top[2 * i]) + u16::from(top[2 * i + 1]);
if two_rows {
s += u16::from(bot[2 * i]) + u16::from(bot[2 * i + 1]);
}
*sum = s;
}
let mut got = vec![0xAAAAu16; n]; unsafe {
if two_rows {
pair_sums_two_rows(&mut got, &top, &bot);
} else {
pair_sums_one_row(&mut got, &top);
}
}
assert_eq!(got, expected, "n={n} two_rows={two_rows}");
}
}
}
}
#[test]
fn apply_sums_match_scalar_bitexact() {
if !enabled() {
eprintln!("skipping: AVX2 unavailable on this CPU");
return;
}
let mut seed = 0x0123_4567_89AB_CDEFu64;
for pw in (1..=67).step_by(3).chain([64, 256]) {
let ph = 3usize;
for shift in [1u32, 2] {
let max_sum = (1u16 << shift) * 255;
let mut sums = Vec::with_capacity(pw * ph);
for i in 0..pw * ph {
let value = match i % 5 {
0 => 0u16,
1 => 1,
2 => max_sum / 2,
3 => max_sum,
_ => {
seed = seed
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
((seed >> 33) % u64::from(max_sum + 1)) as u16
}
};
sums.push(value);
}
for (sample, bytes, srcs) in [
(SampleFormat::U8, 1usize, [16u32, 128, 240]),
(SampleFormat::U16Le, 2, [514, 943, 60395]),
] {
for alpha in opacity_alphas(sample) {
for src in srcs {
let rect = PooledRect {
px0: 2,
py0: 1,
pw,
ph,
shift,
};
let linesize = (pw + 5) * bytes;
let base = lcg_bytes(linesize * (ph + 3), &mut seed);
let mut expected = base.clone();
let mut got = base.clone();
match sample {
SampleFormat::U8 => {
apply_sums::<SampleU8>(
&mut PlaneView {
data: &mut expected,
linesize,
pixel_step: 1,
},
&sums,
rect,
src,
alpha,
);
unsafe {
apply_sums_u8(
&mut PlaneView {
data: &mut got,
linesize,
pixel_step: 1,
},
&sums,
rect,
src,
alpha,
)
};
}
SampleFormat::U16Le => {
apply_sums::<SampleU16Le>(
&mut PlaneView {
data: &mut expected,
linesize,
pixel_step: 2,
},
&sums,
rect,
src,
alpha,
);
unsafe {
apply_sums_u16(
&mut PlaneView {
data: &mut got,
linesize,
pixel_step: 2,
},
&sums,
rect,
src,
alpha,
)
};
}
}
assert_eq!(
got, expected,
"sample={sample:?} pw={pw} shift={shift} alpha={alpha} src={src}"
);
}
}
}
}
}
}
}