use std::arch::aarch64::*;
#[allow(clippy::too_many_arguments)]
#[target_feature(enable = "neon")]
pub(crate) fn quantize_block_ac_neon(
block_in: &[f32],
c: usize,
qm: &[f32],
quant: i32,
scale: f32,
qm_multiplier: f32,
xsize: usize,
ysize: usize,
block_out: &mut [i32],
) {
let width = xsize * 8;
let height = ysize * 8;
let n = width * height;
debug_assert_eq!(qm.len(), n);
debug_assert!(block_in.len() >= n, "block_in too small");
debug_assert!(block_out.len() >= n, "block_out too small");
debug_assert!(width.is_multiple_of(4));
let qm = &qm[..n];
let block_in = &block_in[..n];
let block_out = &mut block_out[..n];
let thr = crate::enc_group::quantize_ac_thresholds(c, xsize, ysize);
let q_scaled = crate::enc_group::quantize_ac_q_scaled(quant, scale, qm_multiplier);
let half = width / 2;
let qs = vdupq_n_f32(q_scaled);
let zero = vdupq_n_s32(0);
for (y, ((qm_row, in_row), out_row)) in qm
.chunks_exact(width)
.zip(block_in.chunks_exact(width))
.zip(block_out.chunks_exact_mut(width))
.take(height)
.enumerate()
{
let yfix = if y >= height / 2 { 2 } else { 0 };
let thr_lo = thr[yfix];
let thr_hi = thr[yfix + 1];
for (x0, ((qm4, in4), out4)) in qm_row
.as_chunks::<4>()
.0
.iter()
.zip(in_row.as_chunks::<4>().0.iter())
.zip(out_row.as_chunks_mut::<4>().0.iter_mut())
.enumerate()
{
let x = x0 * 4;
let threshold = if x >= half { thr_hi } else { thr_lo };
let qmv = unsafe { vld1q_f32(qm4.as_ptr()) };
let inv = unsafe { vld1q_f32(in4.as_ptr()) };
let val = vmulq_f32(vmulq_f32(qmv, qs), inv);
let keep = vcgeq_f32(vabsq_f32(val), vdupq_n_f32(threshold));
let q = vcvtaq_s32_f32(val);
let q = vbslq_s32(keep, q, zero);
unsafe { vst1q_s32(out4.as_mut_ptr(), q) };
}
}
}