#![allow(clippy::undocumented_unsafe_blocks)]
#![allow(unsafe_op_in_unsafe_fn)]
use std::{
arch::x86_64::*,
mem::size_of,
num::{NonZeroU8, NonZeroUsize},
};
use crate::util::Pixel;
#[target_feature(enable = "avx2,avx512f,avx512cd,avx512vl,avx512dq,avx512bw")]
pub(super) fn refine_horizontal_wiener<T: Pixel>(
dest: &mut [T],
src: &[T],
pitch: NonZeroUsize,
width: NonZeroUsize,
height: NonZeroUsize,
bits_per_sample: NonZeroU8,
) {
match size_of::<T>() {
1 => unsafe {
refine_horizontal_wiener_u8(
src.as_ptr().cast(),
dest.as_mut_ptr().cast(),
pitch,
width,
height,
bits_per_sample,
);
},
2 => unsafe {
refine_horizontal_wiener_u16(
src.as_ptr().cast(),
dest.as_mut_ptr().cast(),
pitch,
width,
height,
bits_per_sample,
);
},
_ => unreachable!(),
}
}
#[target_feature(enable = "avx2,avx512f,avx512cd,avx512vl,avx512dq,avx512bw")]
pub(super) fn refine_vertical_wiener<T: Pixel>(
dest: &mut [T],
src: &[T],
pitch: NonZeroUsize,
width: NonZeroUsize,
height: NonZeroUsize,
bits_per_sample: NonZeroU8,
) {
match size_of::<T>() {
1 => unsafe {
refine_vertical_wiener_u8(
src.as_ptr().cast(),
dest.as_mut_ptr().cast(),
pitch,
width,
height,
bits_per_sample,
);
},
2 => unsafe {
refine_vertical_wiener_u16(
src.as_ptr().cast(),
dest.as_mut_ptr().cast(),
pitch,
width,
height,
bits_per_sample,
);
},
_ => unreachable!(),
}
}
#[target_feature(enable = "avx2,avx512f,avx512cd,avx512vl,avx512dq,avx512bw")]
unsafe fn refine_vertical_wiener_u8(
src: *const u8,
dest: *mut u8,
pitch: NonZeroUsize,
width: NonZeroUsize,
height: NonZeroUsize,
_bits_per_sample: NonZeroU8,
) {
let bits_per_sample: u8 = 8;
let pitch = pitch.get();
let width = width.get();
let height = height.get();
let zero = _mm512_setzero_si512();
let avx2_zero = _mm256_setzero_si256();
let one = _mm512_set1_epi16(1);
let avx2_one = _mm256_set1_epi16(1);
let pixel_max = _mm512_set1_epi16(((1i32 << bits_per_sample) - 1) as i16);
let avx2_pixel_max = _mm256_set1_epi16(((1i32 << bits_per_sample) - 1) as i16);
let sixteen = _mm512_set1_epi16(16);
let avx2_sixteen = _mm256_set1_epi16(16);
let w01 = _mm512_set1_epi16(((u16::from((-5i8) as u8) << 8) | 1u16) as i16);
let w23 = _mm512_set1_epi16(((20u16 << 8) | 20u16) as i16);
let w45 = _mm512_set1_epi16(((1u16 << 8) | u16::from((-5i8) as u8)) as i16);
let avx2_w01 = _mm256_set1_epi16(((u16::from((-5i8) as u8) << 8) | 1u16) as i16);
let avx2_w23 = _mm256_set1_epi16(((20u16 << 8) | 20u16) as i16);
let avx2_w45 = _mm256_set1_epi16(((1u16 << 8) | u16::from((-5i8) as u8)) as i16);
let mut offset = 0;
for _ in 0..2.min(height.saturating_sub(1)) {
let mut i = 0;
while i + 64 <= width {
let a_bytes = _mm512_loadu_si512(src.add(offset + i).cast::<__m512i>());
let b_bytes = _mm512_loadu_si512(src.add(offset + i + pitch).cast::<__m512i>());
let a_lo = _mm512_unpacklo_epi8(a_bytes, zero);
let b_lo = _mm512_unpacklo_epi8(b_bytes, zero);
let a_hi = _mm512_unpackhi_epi8(a_bytes, zero);
let b_hi = _mm512_unpackhi_epi8(b_bytes, zero);
let sum_lo = _mm512_add_epi16(_mm512_add_epi16(a_lo, b_lo), one);
let sum_hi = _mm512_add_epi16(_mm512_add_epi16(a_hi, b_hi), one);
let avg_lo = _mm512_srli_epi16(sum_lo, 1);
let avg_hi = _mm512_srli_epi16(sum_hi, 1);
let result = _mm512_packus_epi16(avg_lo, avg_hi);
_mm512_storeu_si512(dest.add(offset + i).cast::<__m512i>(), result);
i += 64;
}
while i + 32 <= width {
let a_bytes = _mm256_loadu_si256(src.add(offset + i).cast::<__m256i>());
let b_bytes = _mm256_loadu_si256(src.add(offset + i + pitch).cast::<__m256i>());
let a_lo = _mm256_unpacklo_epi8(a_bytes, avx2_zero);
let b_lo = _mm256_unpacklo_epi8(b_bytes, avx2_zero);
let a_hi = _mm256_unpackhi_epi8(a_bytes, avx2_zero);
let b_hi = _mm256_unpackhi_epi8(b_bytes, avx2_zero);
let sum_lo = _mm256_add_epi16(_mm256_add_epi16(a_lo, b_lo), avx2_one);
let sum_hi = _mm256_add_epi16(_mm256_add_epi16(a_hi, b_hi), avx2_one);
let avg_lo = _mm256_srli_epi16(sum_lo, 1);
let avg_hi = _mm256_srli_epi16(sum_hi, 1);
let result = _mm256_packus_epi16(avg_lo, avg_hi);
_mm256_storeu_si256(dest.add(offset + i).cast::<__m256i>(), result);
i += 32;
}
while i < width {
let a = *src.add(offset + i) as u16;
let b = *src.add(offset + i + pitch) as u16;
*dest.add(offset + i) = ((a + b + 1) / 2) as u8;
i += 1;
}
offset += pitch;
}
for _ in 2..(height.saturating_sub(4)).max(2) {
let row_m0 = src.add(offset - pitch * 2);
let row_m1 = src.add(offset - pitch);
let row_m2 = src.add(offset);
let row_m3 = src.add(offset + pitch);
let row_m4 = src.add(offset + pitch * 2);
let row_m5 = src.add(offset + pitch * 3);
let row_dest = dest.add(offset);
let mut i = 0;
while i + 64 <= width {
let m0 = _mm512_loadu_si512(row_m0.add(i).cast::<__m512i>());
let m1 = _mm512_loadu_si512(row_m1.add(i).cast::<__m512i>());
let m2 = _mm512_loadu_si512(row_m2.add(i).cast::<__m512i>());
let m3 = _mm512_loadu_si512(row_m3.add(i).cast::<__m512i>());
let m4 = _mm512_loadu_si512(row_m4.add(i).cast::<__m512i>());
let m5 = _mm512_loadu_si512(row_m5.add(i).cast::<__m512i>());
let result = apply_wiener_kernel_u8_fast_zmm(
m0, m1, m2, m3, m4, m5, w01, w23, w45, sixteen, pixel_max,
);
_mm512_storeu_si512(row_dest.add(i).cast::<__m512i>(), result);
i += 64;
}
while i + 32 <= width {
let m0 = _mm256_loadu_si256(row_m0.add(i).cast::<__m256i>());
let m1 = _mm256_loadu_si256(row_m1.add(i).cast::<__m256i>());
let m2 = _mm256_loadu_si256(row_m2.add(i).cast::<__m256i>());
let m3 = _mm256_loadu_si256(row_m3.add(i).cast::<__m256i>());
let m4 = _mm256_loadu_si256(row_m4.add(i).cast::<__m256i>());
let m5 = _mm256_loadu_si256(row_m5.add(i).cast::<__m256i>());
let result = super::avx2::apply_wiener_kernel_u8_fast(
m0,
m1,
m2,
m3,
m4,
m5,
avx2_w01,
avx2_w23,
avx2_w45,
avx2_sixteen,
avx2_pixel_max,
);
_mm256_storeu_si256(row_dest.add(i).cast::<__m256i>(), result);
i += 32;
}
while i < width {
let m0 = *row_m0.add(i) as i16;
let m1 = *row_m1.add(i) as i16;
let mut m2 = *row_m2.add(i) as i16;
let m3 = *row_m3.add(i) as i16;
let m4 = *row_m4.add(i) as i16;
let m5 = *row_m5.add(i) as i16;
m2 = (m2 + m3) * 4;
m2 -= m1 + m4;
m2 *= 5;
let result = (m0 + m5 + m2 + 16) >> 5;
*row_dest.add(i) = result.clamp(0, (1 << bits_per_sample) - 1) as u8;
i += 1;
}
offset += pitch;
}
for _ in (height.saturating_sub(4)).max(2)..height.saturating_sub(1) {
let mut i = 0;
while i + 64 <= width {
let a_bytes = _mm512_loadu_si512(src.add(offset + i).cast::<__m512i>());
let b_bytes = _mm512_loadu_si512(src.add(offset + i + pitch).cast::<__m512i>());
let a_lo = _mm512_unpacklo_epi8(a_bytes, zero);
let b_lo = _mm512_unpacklo_epi8(b_bytes, zero);
let a_hi = _mm512_unpackhi_epi8(a_bytes, zero);
let b_hi = _mm512_unpackhi_epi8(b_bytes, zero);
let sum_lo = _mm512_add_epi16(_mm512_add_epi16(a_lo, b_lo), one);
let sum_hi = _mm512_add_epi16(_mm512_add_epi16(a_hi, b_hi), one);
let avg_lo = _mm512_srli_epi16(sum_lo, 1);
let avg_hi = _mm512_srli_epi16(sum_hi, 1);
let result = _mm512_packus_epi16(avg_lo, avg_hi);
_mm512_storeu_si512(dest.add(offset + i).cast::<__m512i>(), result);
i += 64;
}
while i + 32 <= width {
let a_bytes = _mm256_loadu_si256(src.add(offset + i).cast::<__m256i>());
let b_bytes = _mm256_loadu_si256(src.add(offset + i + pitch).cast::<__m256i>());
let a_lo = _mm256_unpacklo_epi8(a_bytes, avx2_zero);
let b_lo = _mm256_unpacklo_epi8(b_bytes, avx2_zero);
let a_hi = _mm256_unpackhi_epi8(a_bytes, avx2_zero);
let b_hi = _mm256_unpackhi_epi8(b_bytes, avx2_zero);
let sum_lo = _mm256_add_epi16(_mm256_add_epi16(a_lo, b_lo), avx2_one);
let sum_hi = _mm256_add_epi16(_mm256_add_epi16(a_hi, b_hi), avx2_one);
let avg_lo = _mm256_srli_epi16(sum_lo, 1);
let avg_hi = _mm256_srli_epi16(sum_hi, 1);
let result = _mm256_packus_epi16(avg_lo, avg_hi);
_mm256_storeu_si256(dest.add(offset + i).cast::<__m256i>(), result);
i += 32;
}
while i < width {
let a = *src.add(offset + i) as u16;
let b = *src.add(offset + i + pitch) as u16;
*dest.add(offset + i) = ((a + b + 1) / 2) as u8;
i += 1;
}
offset += pitch;
}
if height > 0 {
std::ptr::copy_nonoverlapping(src.add(offset), dest.add(offset), width);
}
}
#[target_feature(enable = "avx2,avx512f,avx512cd,avx512vl,avx512dq,avx512bw")]
unsafe fn refine_vertical_wiener_u16(
src: *const u16,
dest: *mut u16,
pitch: NonZeroUsize,
width: NonZeroUsize,
height: NonZeroUsize,
bits_per_sample: NonZeroU8,
) {
let pitch = pitch.get();
let width = width.get();
let height = height.get();
let use_fast = bits_per_sample.get() <= 15;
let pixel_max = _mm512_set1_epi32((1i32 << bits_per_sample.get()) - 1);
let avx2_pixel_max = _mm256_set1_epi32((1i32 << bits_per_sample.get()) - 1);
let four = _mm256_set1_epi32(4);
let five = _mm256_set1_epi32(5);
let zero = _mm256_setzero_si256();
let zmm_zero = _mm512_setzero_si512();
let zmm_four = _mm512_set1_epi32(4);
let zmm_five = _mm512_set1_epi32(5);
let sixteen = _mm512_set1_epi32(16);
let avx2_sixteen = _mm256_set1_epi32(16);
let w01 = _mm512_set1_epi32(((((-5i32) as u32 & 0xFFFF) << 16) | (1u32 & 0xFFFF)) as i32);
let w23 = _mm512_set1_epi32((((20u32 & 0xFFFF) << 16) | (20u32 & 0xFFFF)) as i32);
let w45 = _mm512_set1_epi32((((1u32 & 0xFFFF) << 16) | ((-5i32) as u32 & 0xFFFF)) as i32);
let avx2_w01 = _mm256_set1_epi32(((((-5i32) as u32 & 0xFFFF) << 16) | (1u32 & 0xFFFF)) as i32);
let avx2_w23 = _mm256_set1_epi32((((20u32 & 0xFFFF) << 16) | (20u32 & 0xFFFF)) as i32);
let avx2_w45 = _mm256_set1_epi32((((1u32 & 0xFFFF) << 16) | ((-5i32) as u32 & 0xFFFF)) as i32);
let permute = _mm512_set_epi64(7, 5, 3, 1, 6, 4, 2, 0);
let mut offset = 0;
for _ in 0..2.min(height.saturating_sub(1)) {
let mut i = 0;
while i + 32 <= width {
let a_words = _mm512_loadu_si512(src.add(offset + i).cast::<__m512i>());
let b_words = _mm512_loadu_si512(src.add(offset + i + pitch).cast::<__m512i>());
_mm512_storeu_si512(
dest.add(offset + i).cast::<__m512i>(),
_mm512_avg_epu16(a_words, b_words),
);
i += 32;
}
while i + 16 <= width {
let a_words = _mm256_loadu_si256(src.add(offset + i).cast::<__m256i>());
let b_words = _mm256_loadu_si256(src.add(offset + i + pitch).cast::<__m256i>());
_mm256_storeu_si256(
dest.add(offset + i).cast::<__m256i>(),
_mm256_avg_epu16(a_words, b_words),
);
i += 16;
}
while i < width {
let a = *src.add(offset + i) as u32;
let b = *src.add(offset + i + pitch) as u32;
*dest.add(offset + i) = ((a + b + 1) / 2) as u16;
i += 1;
}
offset += pitch;
}
for _ in 2..(height.saturating_sub(4)).max(2) {
let row_m0 = src.add(offset - pitch * 2);
let row_m1 = src.add(offset - pitch);
let row_m2 = src.add(offset);
let row_m3 = src.add(offset + pitch);
let row_m4 = src.add(offset + pitch * 2);
let row_m5 = src.add(offset + pitch * 3);
let row_dest = dest.add(offset);
let mut i = 0;
if use_fast {
while i + 32 <= width {
let m0 = _mm512_loadu_si512(row_m0.add(i).cast::<__m512i>());
let m1 = _mm512_loadu_si512(row_m1.add(i).cast::<__m512i>());
let m2 = _mm512_loadu_si512(row_m2.add(i).cast::<__m512i>());
let m3 = _mm512_loadu_si512(row_m3.add(i).cast::<__m512i>());
let m4 = _mm512_loadu_si512(row_m4.add(i).cast::<__m512i>());
let m5 = _mm512_loadu_si512(row_m5.add(i).cast::<__m512i>());
let packed = apply_wiener_kernel_u16_fast_zmm(
m0, m1, m2, m3, m4, m5, w01, w23, w45, sixteen, pixel_max,
);
_mm512_storeu_si512(row_dest.add(i).cast::<__m512i>(), packed);
i += 32;
}
while i + 16 <= width {
let m0 = _mm256_loadu_si256(row_m0.add(i).cast::<__m256i>());
let m1 = _mm256_loadu_si256(row_m1.add(i).cast::<__m256i>());
let m2 = _mm256_loadu_si256(row_m2.add(i).cast::<__m256i>());
let m3 = _mm256_loadu_si256(row_m3.add(i).cast::<__m256i>());
let m4 = _mm256_loadu_si256(row_m4.add(i).cast::<__m256i>());
let m5 = _mm256_loadu_si256(row_m5.add(i).cast::<__m256i>());
let result = super::avx2::apply_wiener_kernel_u16_fast(
m0,
m1,
m2,
m3,
m4,
m5,
avx2_w01,
avx2_w23,
avx2_w45,
avx2_sixteen,
avx2_pixel_max,
);
_mm256_storeu_si256(row_dest.add(i).cast::<__m256i>(), result);
i += 16;
}
} else {
while i + 16 <= width {
let m0 = _mm512_cvtepu16_epi32(_mm256_loadu_si256(row_m0.add(i).cast::<__m256i>()));
let m1 = _mm512_cvtepu16_epi32(_mm256_loadu_si256(row_m1.add(i).cast::<__m256i>()));
let m2 = _mm512_cvtepu16_epi32(_mm256_loadu_si256(row_m2.add(i).cast::<__m256i>()));
let m3 = _mm512_cvtepu16_epi32(_mm256_loadu_si256(row_m3.add(i).cast::<__m256i>()));
let m4 = _mm512_cvtepu16_epi32(_mm256_loadu_si256(row_m4.add(i).cast::<__m256i>()));
let m5 = _mm512_cvtepu16_epi32(_mm256_loadu_si256(row_m5.add(i).cast::<__m256i>()));
let result = apply_wiener_kernel_u16_zmm(
m0, m1, m2, m3, m4, m5, zmm_four, zmm_five, sixteen, pixel_max,
);
let packed = _mm512_packus_epi32(result, zmm_zero);
let packed = _mm512_permutexvar_epi64(permute, packed);
_mm256_storeu_si256(
row_dest.add(i).cast::<__m256i>(),
_mm512_castsi512_si256(packed),
);
i += 16;
}
while i + 8 <= width {
let m0 = _mm256_cvtepu16_epi32(_mm_loadu_si128(row_m0.add(i).cast::<__m128i>()));
let m1 = _mm256_cvtepu16_epi32(_mm_loadu_si128(row_m1.add(i).cast::<__m128i>()));
let m2 = _mm256_cvtepu16_epi32(_mm_loadu_si128(row_m2.add(i).cast::<__m128i>()));
let m3 = _mm256_cvtepu16_epi32(_mm_loadu_si128(row_m3.add(i).cast::<__m128i>()));
let m4 = _mm256_cvtepu16_epi32(_mm_loadu_si128(row_m4.add(i).cast::<__m128i>()));
let m5 = _mm256_cvtepu16_epi32(_mm_loadu_si128(row_m5.add(i).cast::<__m128i>()));
let result = super::avx2::apply_wiener_kernel_u16(
m0,
m1,
m2,
m3,
m4,
m5,
four,
five,
avx2_sixteen,
avx2_pixel_max,
);
let result =
_mm256_permute4x64_epi64(_mm256_packus_epi32(result, zero), 0b11_01_10_00);
_mm_storeu_si128(
row_dest.add(i).cast::<__m128i>(),
_mm256_castsi256_si128(result),
);
i += 8;
}
}
while i < width {
let m0 = *row_m0.add(i) as i32;
let m1 = *row_m1.add(i) as i32;
let mut m2 = *row_m2.add(i) as i32;
let m3 = *row_m3.add(i) as i32;
let m4 = *row_m4.add(i) as i32;
let m5 = *row_m5.add(i) as i32;
m2 = (m2 + m3) * 4;
m2 -= m1 + m4;
m2 *= 5;
let result = (m0 + m5 + m2 + 16) >> 5;
*row_dest.add(i) = result.clamp(0, (1 << bits_per_sample.get()) - 1) as u16;
i += 1;
}
offset += pitch;
}
for _ in (height.saturating_sub(4)).max(2)..height.saturating_sub(1) {
let mut i = 0;
while i + 32 <= width {
let a_words = _mm512_loadu_si512(src.add(offset + i).cast::<__m512i>());
let b_words = _mm512_loadu_si512(src.add(offset + i + pitch).cast::<__m512i>());
_mm512_storeu_si512(
dest.add(offset + i).cast::<__m512i>(),
_mm512_avg_epu16(a_words, b_words),
);
i += 32;
}
while i + 16 <= width {
let a_words = _mm256_loadu_si256(src.add(offset + i).cast::<__m256i>());
let b_words = _mm256_loadu_si256(src.add(offset + i + pitch).cast::<__m256i>());
_mm256_storeu_si256(
dest.add(offset + i).cast::<__m256i>(),
_mm256_avg_epu16(a_words, b_words),
);
i += 16;
}
while i < width {
let a = *src.add(offset + i) as u32;
let b = *src.add(offset + i + pitch) as u32;
*dest.add(offset + i) = ((a + b + 1) / 2) as u16;
i += 1;
}
offset += pitch;
}
if height > 0 {
std::ptr::copy_nonoverlapping(src.add(offset), dest.add(offset), width);
}
}
#[target_feature(enable = "avx2,avx512f,avx512cd,avx512vl,avx512dq,avx512bw")]
unsafe fn refine_horizontal_wiener_u8(
mut src: *const u8,
mut dest: *mut u8,
pitch: NonZeroUsize,
width: NonZeroUsize,
height: NonZeroUsize,
_bits_per_sample: NonZeroU8,
) {
let bits_per_sample: u8 = 8;
let width = width.get();
let height = height.get();
let pitch = pitch.get();
let pixel_max = _mm512_set1_epi16(((1i32 << bits_per_sample) - 1) as i16);
let avx2_pixel_max = _mm256_set1_epi16(((1i32 << bits_per_sample) - 1) as i16);
let sixteen = _mm512_set1_epi16(16);
let avx2_sixteen = _mm256_set1_epi16(16);
let w01 = _mm512_set1_epi16(((u16::from((-5i8) as u8) << 8) | 1u16) as i16);
let w23 = _mm512_set1_epi16(((20u16 << 8) | 20u16) as i16);
let w45 = _mm512_set1_epi16(((1u16 << 8) | u16::from((-5i8) as u8)) as i16);
let avx2_w01 = _mm256_set1_epi16(((u16::from((-5i8) as u8) << 8) | 1u16) as i16);
let avx2_w23 = _mm256_set1_epi16(((20u16 << 8) | 20u16) as i16);
let avx2_w45 = _mm256_set1_epi16(((1u16 << 8) | u16::from((-5i8) as u8)) as i16);
for _y in 0..height {
*dest = ((*src as u16 + *src.add(1) as u16 + 1) >> 1) as u8;
*dest.add(1) = ((*src.add(1) as u16 + *src.add(2) as u16 + 1) >> 1) as u8;
let wiener_start = 2;
let wiener_end = if width >= 4 { width - 4 } else { wiener_start };
let mut x = wiener_start;
while x + 64 <= wiener_end {
let m0 = _mm512_loadu_si512(src.add(x - 2).cast::<__m512i>());
let m1 = _mm512_loadu_si512(src.add(x - 1).cast::<__m512i>());
let m2 = _mm512_loadu_si512(src.add(x).cast::<__m512i>());
let m3 = _mm512_loadu_si512(src.add(x + 1).cast::<__m512i>());
let m4 = _mm512_loadu_si512(src.add(x + 2).cast::<__m512i>());
let m5 = _mm512_loadu_si512(src.add(x + 3).cast::<__m512i>());
let result = apply_wiener_kernel_u8_fast_zmm(
m0, m1, m2, m3, m4, m5, w01, w23, w45, sixteen, pixel_max,
);
_mm512_storeu_si512(dest.add(x).cast::<__m512i>(), result);
x += 64;
}
while x + 32 <= wiener_end {
let m0 = _mm256_loadu_si256(src.add(x - 2).cast::<__m256i>());
let m1 = _mm256_loadu_si256(src.add(x - 1).cast::<__m256i>());
let m2 = _mm256_loadu_si256(src.add(x).cast::<__m256i>());
let m3 = _mm256_loadu_si256(src.add(x + 1).cast::<__m256i>());
let m4 = _mm256_loadu_si256(src.add(x + 2).cast::<__m256i>());
let m5 = _mm256_loadu_si256(src.add(x + 3).cast::<__m256i>());
let result = super::avx2::apply_wiener_kernel_u8_fast(
m0,
m1,
m2,
m3,
m4,
m5,
avx2_w01,
avx2_w23,
avx2_w45,
avx2_sixteen,
avx2_pixel_max,
);
_mm256_storeu_si256(dest.add(x).cast::<__m256i>(), result);
x += 32;
}
while x < wiener_end {
let m0 = *src.add(x - 2) as i16;
let m1 = *src.add(x - 1) as i16;
let mut m2 = *src.add(x) as i16;
let m3 = *src.add(x + 1) as i16;
let m4 = *src.add(x + 2) as i16;
let m5 = *src.add(x + 3) as i16;
m2 = (m2 + m3) * 4;
m2 -= m1 + m4;
m2 *= 5;
let result = (m0 + m5 + m2 + 16) >> 5;
*dest.add(x) = result.clamp(0, (1 << bits_per_sample) - 1) as u8;
x += 1;
}
for x in wiener_end..(width - 1) {
*dest.add(x) = ((*src.add(x) as u16 + *src.add(x + 1) as u16 + 1) >> 1) as u8;
}
*dest.add(width - 1) = *src.add(width - 1);
dest = dest.add(pitch);
src = src.add(pitch);
}
}
#[target_feature(enable = "avx2,avx512f,avx512cd,avx512vl,avx512dq,avx512bw")]
unsafe fn refine_horizontal_wiener_u16(
src: *const u16,
dest: *mut u16,
pitch: NonZeroUsize,
width: NonZeroUsize,
height: NonZeroUsize,
bits_per_sample: NonZeroU8,
) {
let use_fast = bits_per_sample.get() <= 15;
let pixel_max = _mm512_set1_epi32((1i32 << bits_per_sample.get()) - 1);
let avx2_pixel_max = _mm256_set1_epi32((1i32 << bits_per_sample.get()) - 1);
let four = _mm256_set1_epi32(4);
let five = _mm256_set1_epi32(5);
let sixteen = _mm512_set1_epi32(16);
let avx2_sixteen = _mm256_set1_epi32(16);
let w01 = _mm512_set1_epi32(((((-5i32) as u32 & 0xFFFF) << 16) | (1u32 & 0xFFFF)) as i32);
let w23 = _mm512_set1_epi32((((20u32 & 0xFFFF) << 16) | (20u32 & 0xFFFF)) as i32);
let w45 = _mm512_set1_epi32((((1u32 & 0xFFFF) << 16) | ((-5i32) as u32 & 0xFFFF)) as i32);
let avx2_w01 = _mm256_set1_epi32(((((-5i32) as u32 & 0xFFFF) << 16) | (1u32 & 0xFFFF)) as i32);
let avx2_w23 = _mm256_set1_epi32((((20u32 & 0xFFFF) << 16) | (20u32 & 0xFFFF)) as i32);
let avx2_w45 = _mm256_set1_epi32((((1u32 & 0xFFFF) << 16) | ((-5i32) as u32 & 0xFFFF)) as i32);
let zero = _mm256_setzero_si256();
let zmm_zero = _mm512_setzero_si512();
let zmm_four = _mm512_set1_epi32(4);
let zmm_five = _mm512_set1_epi32(5);
let permute = _mm512_set_epi64(7, 5, 3, 1, 6, 4, 2, 0);
let mut offset = 0;
for _j in 0..height.get() {
if width.get() >= 2 {
let a = *src.add(offset) as u32;
let b = *src.add(offset + 1) as u32;
*dest.add(offset) = ((a + b + 1) / 2) as u16;
if width.get() >= 3 {
let c = *src.add(offset + 2) as u32;
*dest.add(offset + 1) = ((b + c + 1) / 2) as u16;
}
}
let wiener_start = 2;
let wiener_end = if width.get() >= 4 {
width.get() - 4
} else {
wiener_start
};
let mut i = wiener_start;
if use_fast {
while i + 32 <= wiener_end {
let m0 = _mm512_loadu_si512(src.add(offset + i - 2).cast::<__m512i>());
let m1 = _mm512_loadu_si512(src.add(offset + i - 1).cast::<__m512i>());
let m2 = _mm512_loadu_si512(src.add(offset + i).cast::<__m512i>());
let m3 = _mm512_loadu_si512(src.add(offset + i + 1).cast::<__m512i>());
let m4 = _mm512_loadu_si512(src.add(offset + i + 2).cast::<__m512i>());
let m5 = _mm512_loadu_si512(src.add(offset + i + 3).cast::<__m512i>());
let packed = apply_wiener_kernel_u16_fast_zmm(
m0, m1, m2, m3, m4, m5, w01, w23, w45, sixteen, pixel_max,
);
_mm512_storeu_si512(dest.add(offset + i).cast::<__m512i>(), packed);
i += 32;
}
while i + 16 <= wiener_end {
let m0 = _mm256_loadu_si256(src.add(offset + i - 2).cast::<__m256i>());
let m1 = _mm256_loadu_si256(src.add(offset + i - 1).cast::<__m256i>());
let m2 = _mm256_loadu_si256(src.add(offset + i).cast::<__m256i>());
let m3 = _mm256_loadu_si256(src.add(offset + i + 1).cast::<__m256i>());
let m4 = _mm256_loadu_si256(src.add(offset + i + 2).cast::<__m256i>());
let m5 = _mm256_loadu_si256(src.add(offset + i + 3).cast::<__m256i>());
let result = super::avx2::apply_wiener_kernel_u16_fast(
m0,
m1,
m2,
m3,
m4,
m5,
avx2_w01,
avx2_w23,
avx2_w45,
avx2_sixteen,
avx2_pixel_max,
);
_mm256_storeu_si256(dest.add(offset + i).cast::<__m256i>(), result);
i += 16;
}
} else {
while i + 16 <= wiener_end {
let src_i = offset + i;
let m0 =
_mm512_cvtepu16_epi32(_mm256_loadu_si256(src.add(src_i - 2).cast::<__m256i>()));
let m1 =
_mm512_cvtepu16_epi32(_mm256_loadu_si256(src.add(src_i - 1).cast::<__m256i>()));
let m2 =
_mm512_cvtepu16_epi32(_mm256_loadu_si256(src.add(src_i).cast::<__m256i>()));
let m3 =
_mm512_cvtepu16_epi32(_mm256_loadu_si256(src.add(src_i + 1).cast::<__m256i>()));
let m4 =
_mm512_cvtepu16_epi32(_mm256_loadu_si256(src.add(src_i + 2).cast::<__m256i>()));
let m5 =
_mm512_cvtepu16_epi32(_mm256_loadu_si256(src.add(src_i + 3).cast::<__m256i>()));
let result = apply_wiener_kernel_u16_zmm(
m0, m1, m2, m3, m4, m5, zmm_four, zmm_five, sixteen, pixel_max,
);
let packed = _mm512_packus_epi32(result, zmm_zero);
let packed = _mm512_permutexvar_epi64(permute, packed);
_mm256_storeu_si256(
dest.add(src_i).cast::<__m256i>(),
_mm512_castsi512_si256(packed),
);
i += 16;
}
while i + 8 <= wiener_end {
let src_i = offset + i;
let m0 =
_mm256_cvtepu16_epi32(_mm_loadu_si128(src.add(src_i - 2).cast::<__m128i>()));
let m1 =
_mm256_cvtepu16_epi32(_mm_loadu_si128(src.add(src_i - 1).cast::<__m128i>()));
let m2 = _mm256_cvtepu16_epi32(_mm_loadu_si128(src.add(src_i).cast::<__m128i>()));
let m3 =
_mm256_cvtepu16_epi32(_mm_loadu_si128(src.add(src_i + 1).cast::<__m128i>()));
let m4 =
_mm256_cvtepu16_epi32(_mm_loadu_si128(src.add(src_i + 2).cast::<__m128i>()));
let m5 =
_mm256_cvtepu16_epi32(_mm_loadu_si128(src.add(src_i + 3).cast::<__m128i>()));
let result = super::avx2::apply_wiener_kernel_u16(
m0,
m1,
m2,
m3,
m4,
m5,
four,
five,
avx2_sixteen,
avx2_pixel_max,
);
let result =
_mm256_permute4x64_epi64(_mm256_packus_epi32(result, zero), 0b11_01_10_00);
_mm_storeu_si128(
dest.add(src_i).cast::<__m128i>(),
_mm256_castsi256_si128(result),
);
i += 8;
}
}
while i < wiener_end {
let m0 = *src.add(offset + i - 2) as i32;
let m1 = *src.add(offset + i - 1) as i32;
let mut m2 = *src.add(offset + i) as i32;
let m3 = *src.add(offset + i + 1) as i32;
let m4 = *src.add(offset + i + 2) as i32;
let m5 = *src.add(offset + i + 3) as i32;
m2 = (m2 + m3) * 4;
m2 -= m1 + m4;
m2 *= 5;
let result = (m0 + m5 + m2 + 16) >> 5;
*dest.add(offset + i) = result.clamp(0, (1 << bits_per_sample.get()) - 1) as u16;
i += 1;
}
for i in wiener_end..(width.get() - 1).min(width.get()) {
let a = *src.add(offset + i) as u32;
let b = *src.add(offset + i + 1) as u32;
*dest.add(offset + i) = ((a + b + 1) / 2) as u16;
}
if width.get() > 0 {
*dest.add(offset + width.get() - 1) = *src.add(offset + width.get() - 1);
}
offset += pitch.get();
}
}
#[inline]
#[target_feature(enable = "avx2,avx512f,avx512cd,avx512vl,avx512dq,avx512bw")]
unsafe fn apply_wiener_kernel_u8_fast_zmm(
m0_bytes: __m512i,
m1_bytes: __m512i,
m2_bytes: __m512i,
m3_bytes: __m512i,
m4_bytes: __m512i,
m5_bytes: __m512i,
w01: __m512i,
w23: __m512i,
w45: __m512i,
sixteen: __m512i,
pixel_max: __m512i,
) -> __m512i {
let pair01_lo = _mm512_unpacklo_epi8(m0_bytes, m1_bytes);
let pair23_lo = _mm512_unpacklo_epi8(m2_bytes, m3_bytes);
let pair45_lo = _mm512_unpacklo_epi8(m4_bytes, m5_bytes);
let pair01_hi = _mm512_unpackhi_epi8(m0_bytes, m1_bytes);
let pair23_hi = _mm512_unpackhi_epi8(m2_bytes, m3_bytes);
let pair45_hi = _mm512_unpackhi_epi8(m4_bytes, m5_bytes);
let lo = _mm512_add_epi16(
_mm512_add_epi16(
_mm512_maddubs_epi16(pair01_lo, w01),
_mm512_maddubs_epi16(pair23_lo, w23),
),
_mm512_add_epi16(_mm512_maddubs_epi16(pair45_lo, w45), sixteen),
);
let hi = _mm512_add_epi16(
_mm512_add_epi16(
_mm512_maddubs_epi16(pair01_hi, w01),
_mm512_maddubs_epi16(pair23_hi, w23),
),
_mm512_add_epi16(_mm512_maddubs_epi16(pair45_hi, w45), sixteen),
);
let res_lo = _mm512_srai_epi16(lo, 5);
let res_hi = _mm512_srai_epi16(hi, 5);
let zero = _mm512_setzero_si512();
let clamped_lo = _mm512_max_epi16(zero, _mm512_min_epi16(res_lo, pixel_max));
let clamped_hi = _mm512_max_epi16(zero, _mm512_min_epi16(res_hi, pixel_max));
_mm512_packus_epi16(clamped_lo, clamped_hi)
}
#[inline]
#[target_feature(enable = "avx2,avx512f,avx512cd,avx512vl,avx512dq,avx512bw")]
unsafe fn apply_wiener_kernel_u16_fast_zmm(
m0: __m512i,
m1: __m512i,
m2: __m512i,
m3: __m512i,
m4: __m512i,
m5: __m512i,
w01: __m512i,
w23: __m512i,
w45: __m512i,
sixteen: __m512i,
pixel_max: __m512i,
) -> __m512i {
let pair01_lo = _mm512_unpacklo_epi16(m0, m1);
let pair23_lo = _mm512_unpacklo_epi16(m2, m3);
let pair45_lo = _mm512_unpacklo_epi16(m4, m5);
let pair01_hi = _mm512_unpackhi_epi16(m0, m1);
let pair23_hi = _mm512_unpackhi_epi16(m2, m3);
let pair45_hi = _mm512_unpackhi_epi16(m4, m5);
let lo = _mm512_add_epi32(
_mm512_add_epi32(
_mm512_madd_epi16(pair01_lo, w01),
_mm512_madd_epi16(pair23_lo, w23),
),
_mm512_add_epi32(_mm512_madd_epi16(pair45_lo, w45), sixteen),
);
let hi = _mm512_add_epi32(
_mm512_add_epi32(
_mm512_madd_epi16(pair01_hi, w01),
_mm512_madd_epi16(pair23_hi, w23),
),
_mm512_add_epi32(_mm512_madd_epi16(pair45_hi, w45), sixteen),
);
let res_lo = _mm512_srai_epi32(lo, 5);
let res_hi = _mm512_srai_epi32(hi, 5);
let zero = _mm512_setzero_si512();
let clamped_lo = _mm512_max_epi32(zero, _mm512_min_epi32(res_lo, pixel_max));
let clamped_hi = _mm512_max_epi32(zero, _mm512_min_epi32(res_hi, pixel_max));
_mm512_packus_epi32(clamped_lo, clamped_hi)
}
#[inline]
#[target_feature(enable = "avx2,avx512f,avx512cd,avx512vl,avx512dq,avx512bw")]
unsafe fn apply_wiener_kernel_u16_zmm(
m0: __m512i,
m1: __m512i,
m2: __m512i,
m3: __m512i,
m4: __m512i,
m5: __m512i,
four: __m512i,
five: __m512i,
sixteen: __m512i,
pixel_max: __m512i,
) -> __m512i {
let center = _mm512_mullo_epi32(_mm512_add_epi32(m2, m3), four);
let neighbors = _mm512_add_epi32(m1, m4);
let weighted = _mm512_mullo_epi32(_mm512_sub_epi32(center, neighbors), five);
let sum = _mm512_add_epi32(
_mm512_add_epi32(m0, m5),
_mm512_add_epi32(weighted, sixteen),
);
let result = _mm512_srai_epi32(sum, 5);
let zero = _mm512_setzero_si512();
_mm512_max_epi32(zero, _mm512_min_epi32(result, pixel_max))
}