use crate::filter_weights::FilterWeights;
use crate::neon::utils::{load_3b_as_u16x4, vxmlal_high_lane_s16, vxmlal_lane_s16, vxmlal_s16};
use crate::support::PRECISION;
use std::arch::aarch64::*;
#[must_use]
#[inline(always)]
fn conv_horiz_rgb_4_u8<const D: bool>(
start_x: usize,
src: &[u8],
weights: int16x4_t,
store: int32x4_t,
shuffle: uint8x16_t,
) -> int32x4_t {
unsafe {
const CN: usize = 3;
let src_ptr = src.get_unchecked((start_x * CN)..);
let mut rgb_pixel = vcombine_u8(vld1_u8(src_ptr.as_ptr()), vdup_n_u8(0));
rgb_pixel = vreinterpretq_u8_u32(vld1q_lane_u32::<2>(
src_ptr.get_unchecked(8..).as_ptr().cast(),
vreinterpretq_u32_u8(rgb_pixel),
));
rgb_pixel = vqtbl1q_u8(rgb_pixel, shuffle);
let hi = vreinterpretq_s16_u16(vmovl_high_u8(rgb_pixel));
let lo = vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(rgb_pixel)));
let acc = vxmlal_high_lane_s16::<D, 3>(store, hi, weights);
let acc = vxmlal_lane_s16::<D, 2>(acc, vget_low_s16(hi), weights);
let acc = vxmlal_high_lane_s16::<D, 1>(acc, lo, weights);
vxmlal_lane_s16::<D, 0>(acc, vget_low_s16(lo), weights)
}
}
#[must_use]
#[inline(always)]
fn conv_horiz_rgba_2_u8<const D: bool>(
start_x: usize,
src: &[u8],
weights: int16x4_t,
store: int32x4_t,
shuffle: uint8x8_t,
) -> int32x4_t {
unsafe {
const CN: usize = 3;
let src_ptr = src.get_unchecked((start_x * CN)..);
let mut rgb_pixel = vld1_lane_u32::<0>(src_ptr.as_ptr().cast(), vdup_n_u32(0));
rgb_pixel = vreinterpret_u32_u16(vld1_lane_u16::<2>(
src_ptr.get_unchecked(4..).as_ptr().cast(),
vreinterpret_u16_u32(rgb_pixel),
));
rgb_pixel = vreinterpret_u32_u8(vtbl1_u8(vreinterpret_u8_u32(rgb_pixel), shuffle));
let wide = vreinterpretq_s16_u16(vmovl_u8(vreinterpret_u8_u32(rgb_pixel)));
let acc = vxmlal_high_lane_s16::<D, 1>(store, wide, weights);
vxmlal_lane_s16::<D, 0>(acc, vget_low_s16(wide), weights)
}
}
#[must_use]
#[inline(always)]
fn conv_horiz_rgba_1_u8<const D: bool>(
start_x: usize,
src: &[u8],
w0: int16x4_t,
store: int32x4_t,
) -> int32x4_t {
unsafe {
const CN: usize = 3;
let src_ptr = src.get_unchecked((start_x * CN)..);
let rgb_pixel = load_3b_as_u16x4(src_ptr.as_ptr());
let lo = vreinterpret_s16_u16(rgb_pixel);
vxmlal_s16::<D>(store, lo, w0)
}
}
#[inline(always)]
fn write_accumulator_u8<const PRECISION: i32>(store: int32x4_t, dst: &mut [u8]) {
unsafe {
let store_16 = vqshrun_n_s32::<PRECISION>(store);
let store_16_8 = vqmovn_u16(vcombine_u16(store_16, store_16));
vst1_lane_u16::<0>(
dst.as_mut_ptr() as *mut u16,
vreinterpret_u16_u8(store_16_8),
);
vst1_lane_u8::<2>(dst.as_mut_ptr().add(2), store_16_8);
}
}
pub(crate) fn convolve_horizontal_rgb_neon_rows_4(
src: &[u8],
src_stride: usize,
dst: &mut [u8],
dst_stride: usize,
filter_weights: &FilterWeights<i16>,
_: u32,
) {
convolve_horizontal_rgb_neon_rows_4_impl::<false, PRECISION>(
src,
src_stride,
dst,
dst_stride,
filter_weights,
);
}
fn convolve_horizontal_rgb_neon_rows_4_impl<const D: bool, const PRECISION: i32>(
src: &[u8],
src_stride: usize,
dst: &mut [u8],
dst_stride: usize,
filter_weights: &FilterWeights<i16>,
) {
unsafe {
let shuf_table_1: [u8; 8] = [0, 1, 2, 255, 3, 4, 5, 255];
let shuffle_1 = vld1_u8(shuf_table_1.as_ptr());
let shuf_table_2: [u8; 8] = [6, 7, 8, 255, 9, 10, 11, 255];
let shuffle_2 = vld1_u8(shuf_table_2.as_ptr());
let shuffle = vcombine_u8(shuffle_1, shuffle_2);
let rnd_const: i32 = 1 << (PRECISION - 1);
const CN: usize = 3;
let init = vdupq_n_s32(rnd_const);
let (row0_ref, rest) = dst.split_at_mut(dst_stride);
let (row1_ref, rest) = rest.split_at_mut(dst_stride);
let (row2_ref, row3_ref) = rest.split_at_mut(dst_stride);
let iter_row0 = row0_ref.as_chunks_mut::<CN>().0;
let iter_row1 = row1_ref.as_chunks_mut::<CN>().0;
let iter_row2 = row2_ref.as_chunks_mut::<CN>().0;
let iter_row3 = row3_ref.as_chunks_mut::<CN>().0;
for (((((chunk0, chunk1), chunk2), chunk3), &bounds), weights) in iter_row0
.iter_mut()
.zip(iter_row1.iter_mut())
.zip(iter_row2.iter_mut())
.zip(iter_row3.iter_mut())
.zip(filter_weights.bounds.iter())
.zip(
filter_weights
.weights
.chunks_exact(filter_weights.aligned_size),
)
{
let mut jx = 0usize;
let mut store_0 = init;
let mut store_1 = init;
let mut store_2 = init;
let mut store_3 = init;
let src0 = src;
let src1 = src0.get_unchecked(src_stride..);
let src2 = src1.get_unchecked(src_stride..);
let src3 = src2.get_unchecked(src_stride..);
while jx + 4 <= bounds.size {
let bounds_start = bounds.start + jx;
let w_ptr = weights.get_unchecked(jx..);
let weights = vld1_s16(w_ptr.as_ptr());
store_0 = conv_horiz_rgb_4_u8::<D>(bounds_start, src0, weights, store_0, shuffle);
store_1 = conv_horiz_rgb_4_u8::<D>(bounds_start, src1, weights, store_1, shuffle);
store_2 = conv_horiz_rgb_4_u8::<D>(bounds_start, src2, weights, store_2, shuffle);
store_3 = conv_horiz_rgb_4_u8::<D>(bounds_start, src3, weights, store_3, shuffle);
jx += 4;
}
while jx + 2 <= bounds.size {
let w_ptr = weights.get_unchecked(jx..);
let bnds = bounds.start + jx;
let mut v_weight = vld1_dup_s16(w_ptr.as_ptr());
v_weight = vld1_lane_s16::<1>(w_ptr.as_ptr().add(1), v_weight);
store_0 = conv_horiz_rgba_2_u8::<D>(bnds, src0, v_weight, store_0, shuffle_1);
store_1 = conv_horiz_rgba_2_u8::<D>(bnds, src1, v_weight, store_1, shuffle_1);
store_2 = conv_horiz_rgba_2_u8::<D>(bnds, src2, v_weight, store_2, shuffle_1);
store_3 = conv_horiz_rgba_2_u8::<D>(bnds, src3, v_weight, store_3, shuffle_1);
jx += 2;
}
while jx < bounds.size {
let w_ptr = weights.get_unchecked(jx..);
let bnds = bounds.start + jx;
let weight0 = vld1_dup_s16(w_ptr.as_ptr());
store_0 = conv_horiz_rgba_1_u8::<D>(bnds, src0, weight0, store_0);
store_1 = conv_horiz_rgba_1_u8::<D>(bnds, src1, weight0, store_1);
store_2 = conv_horiz_rgba_1_u8::<D>(bnds, src2, weight0, store_2);
store_3 = conv_horiz_rgba_1_u8::<D>(bnds, src3, weight0, store_3);
jx += 1;
}
write_accumulator_u8::<PRECISION>(store_0, chunk0);
write_accumulator_u8::<PRECISION>(store_1, chunk1);
write_accumulator_u8::<PRECISION>(store_2, chunk2);
write_accumulator_u8::<PRECISION>(store_3, chunk3);
}
}
}
pub(crate) fn convolve_horizontal_rgb_neon_row_one(
src: &[u8],
dst: &mut [u8],
filter_weights: &FilterWeights<i16>,
_: u32,
) {
convolve_horizontal_rgb_neon_row_one_impl::<false, PRECISION>(src, dst, filter_weights);
}
fn convolve_horizontal_rgb_neon_row_one_impl<const D: bool, const PRECISION: i32>(
src: &[u8],
dst: &mut [u8],
filter_weights: &FilterWeights<i16>,
) {
unsafe {
const CN: usize = 3;
let shuf_table_1: [u8; 8] = [0, 1, 2, 255, 3, 4, 5, 255];
let shuffle_1 = vld1_u8(shuf_table_1.as_ptr());
let shuf_table_2: [u8; 8] = [6, 7, 8, 255, 9, 10, 11, 255];
let shuffle_2 = vld1_u8(shuf_table_2.as_ptr());
let shuffle = vcombine_u8(shuffle_1, shuffle_2);
let rnd_const: i32 = 1i32 << (PRECISION - 1);
for ((dst, bounds), weights) in dst
.as_chunks_mut::<CN>()
.0
.iter_mut()
.zip(filter_weights.bounds.iter())
.zip(
filter_weights
.weights
.chunks_exact(filter_weights.aligned_size),
)
{
let bounds_size = bounds.size;
let mut jx = 0usize;
let mut store = vdupq_n_s32(rnd_const);
while jx + 4 <= bounds_size {
let bounds_start = bounds.start + jx;
let w_ptr = weights.get_unchecked(jx..);
let weights = vld1_s16(w_ptr.as_ptr());
store = conv_horiz_rgb_4_u8::<D>(bounds_start, src, weights, store, shuffle);
jx += 4;
}
while jx + 2 <= bounds_size {
let w_ptr = weights.get_unchecked(jx..);
let bounds_start = bounds.start + jx;
let v_weight = vreinterpret_s16_s32(vld1_dup_s32(w_ptr.as_ptr() as *const _));
store = conv_horiz_rgba_2_u8::<D>(bounds_start, src, v_weight, store, shuffle_1);
jx += 2;
}
while jx < bounds_size {
let w_ptr = weights.get_unchecked(jx..);
let weight0 = vld1_dup_s16(w_ptr.as_ptr());
let bnds = bounds.start + jx;
store = conv_horiz_rgba_1_u8::<D>(bnds, src, weight0, store);
jx += 1;
}
write_accumulator_u8::<PRECISION>(store, dst);
}
}
}