use crate::neon::{load_f32_f16, store_f32_f16};
use crate::stackblur::stack_blur_pass::StackBlurWorkingPass;
use crate::unsafe_slice::UnsafeSlice;
use crate::util::ScratchBuffer;
use core::f16;
use std::arch::aarch64::*;
pub(crate) struct VerticalNeonStackBlurPassFloat16<const CN: usize> {}
impl<const CN: usize> Default for VerticalNeonStackBlurPassFloat16<CN> {
fn default() -> Self {
VerticalNeonStackBlurPassFloat16::<CN> {}
}
}
impl<const CN: usize> VerticalNeonStackBlurPassFloat16<CN> {
#[target_feature(enable = "neon")]
fn pass_impl(
&self,
pixels: &UnsafeSlice<f16>,
stride: u32,
width: u32,
height: u32,
radius: u32,
thread: usize,
total_threads: usize,
) {
unsafe {
let div = ((radius * 2) + 1) as usize;
let mut yp;
let mut sp;
let mut stack_start;
let mut scratch_buffer = ScratchBuffer::<f32, 2048>::new(4 * div);
let stacks = scratch_buffer.as_mut_slice();
let v_mul_value = vdupq_n_f32(1. / ((radius as f32 + 1.) * (radius as f32 + 1.)));
let hm = height - 1;
let div = (radius * 2) + 1;
let mut src_ptr;
let mut dst_ptr;
let min_x = thread * width as usize / total_threads;
let max_x = (thread + 1) * width as usize / total_threads;
for x in min_x..max_x {
let mut sums = vdupq_n_f32(0.);
let mut sum_in = vdupq_n_f32(0.);
let mut sum_out = vdupq_n_f32(0.);
src_ptr = CN * x;
let src_pixel = load_f32_f16::<CN>(pixels.get_ptr(src_ptr).cast());
for i in 0..=radius {
let stack_ptr = stacks.as_mut_ptr().add(i as usize * 4);
vst1q_f32(stack_ptr.cast(), src_pixel);
sums = vfmaq_n_f32(sums, src_pixel, i as f32 + 1.);
sum_out = vaddq_f32(sum_out, src_pixel);
}
for i in 1..=radius {
if i <= hm {
src_ptr += stride as usize;
}
let stack_ptr = stacks.as_mut_ptr().add((i + radius) as usize * 4);
let src_pixel = load_f32_f16::<CN>(pixels.get_ptr(src_ptr));
vst1q_f32(stack_ptr.cast(), src_pixel);
sums = vfmaq_n_f32(sums, src_pixel, radius as f32 + 1. - i as f32);
sum_in = vaddq_f32(sum_in, src_pixel);
}
sp = radius;
yp = radius;
if yp > hm {
yp = hm;
}
src_ptr = CN * x + yp as usize * stride as usize;
dst_ptr = CN * x;
for _ in 0..height {
let blurred = vmulq_f32(sums, v_mul_value);
store_f32_f16::<CN>(pixels.get_ptr(dst_ptr), blurred);
dst_ptr += stride as usize;
sums = vsubq_f32(sums, sum_out);
stack_start = sp + div - radius;
if stack_start >= div {
stack_start -= div;
}
let stack_ptr = stacks.as_mut_ptr().add(stack_start as usize * 4);
let stack_val = vld1q_f32(stack_ptr.cast());
sum_out = vsubq_f32(sum_out, stack_val);
if yp < hm {
src_ptr += stride as usize; yp += 1;
}
let src_pixel = load_f32_f16::<CN>(pixels.get_ptr(src_ptr));
vst1q_f32(stack_ptr.cast(), src_pixel);
sum_in = vaddq_f32(sum_in, src_pixel);
sums = vaddq_f32(sums, sum_in);
sp += 1;
if sp >= div {
sp = 0;
}
let stack_ptr = stacks.as_mut_ptr().add(sp as usize * 4);
let stack_val = vld1q_f32(stack_ptr.cast());
sum_out = vaddq_f32(sum_out, stack_val);
sum_in = vsubq_f32(sum_in, stack_val);
}
}
}
}
}
impl<const CN: usize> StackBlurWorkingPass<f16, CN> for VerticalNeonStackBlurPassFloat16<CN> {
fn pass(
&self,
pixels: &UnsafeSlice<f16>,
stride: u32,
width: u32,
height: u32,
radius: u32,
thread: usize,
total_threads: usize,
) {
unsafe {
self.pass_impl(pixels, stride, width, height, radius, thread, total_threads);
}
}
}