use std::arch::aarch64::*;
#[inline]
#[target_feature(enable = "neon")]
fn neon_log2p1_f32(x: float32x4_t) -> float32x4_t {
let y = vaddq_f32(x, vdupq_n_f32(1.0));
let bits = vreinterpretq_u32_f32(y);
let exp_u = vshrq_n_u32::<23>(bits);
let exp_i = vsubq_s32(vreinterpretq_s32_u32(exp_u), vdupq_n_s32(127));
let e = vcvtq_f32_s32(exp_i);
let mant_bits = vorrq_u32(
vandq_u32(bits, vdupq_n_u32(0x007f_ffff)),
vdupq_n_u32(0x3f80_0000),
);
let m = vreinterpretq_f32_u32(mant_bits);
let t = vsubq_f32(m, vdupq_n_f32(1.0));
let c0 = vdupq_n_f32(1.4426934719085693359375);
let c1 = vdupq_n_f32(-0.721179187297821044921875);
let c2 = vdupq_n_f32(0.477900087833404541015625);
let c3 = vdupq_n_f32(-0.340080082416534423828125);
let c4 = vdupq_n_f32(0.21719777584075927734375);
let c5 = vdupq_n_f32(-9.749893844127655029296875e-2);
let c6 = vdupq_n_f32(2.096841670572757720947265625e-2);
let mut p = c6;
p = vfmaq_f32(c5, t, p);
p = vfmaq_f32(c4, t, p);
p = vfmaq_f32(c3, t, p);
p = vfmaq_f32(c2, t, p);
p = vfmaq_f32(c1, t, p);
p = vfmaq_f32(c0, t, p);
vaddq_f32(e, vmulq_f32(t, p))
}
#[target_feature(enable = "neon")]
pub(crate) fn coeff_rate_f32_neon(lev: &[f32]) -> f32 {
let mut sum = vdupq_n_f32(0.0);
let two = vdupq_n_f32(2.0);
let zero = vdupq_n_f32(0.0);
let chunks = lev.as_chunks::<4>();
let rem = chunks.1;
for chunk in chunks.0.iter() {
let v = unsafe { vld1q_f32(chunk.as_ptr()) };
let a = vabsq_f32(v);
let nz = vcgtq_f32(a, zero);
let cost = vfmaq_f32(two, two, neon_log2p1_f32(a));
sum = vaddq_f32(sum, vbslq_f32(nz, cost, zero));
}
let mut out = vaddvq_f32(sum);
for &v in rem {
if v != 0.0 {
out += 2.0 + 2.0 * crate::av2::helpers::log2p1_approx_f32(v.abs());
}
}
out
}
#[target_feature(enable = "neon")]
pub(crate) fn coeff_abs_rate_f32_neon(lev: &[f32]) -> f32 {
let mut sum = vdupq_n_f32(0.0);
let chunks = lev.as_chunks::<4>();
let rem = chunks.1;
for chunk in chunks.0.iter() {
let v = unsafe { vld1q_f32(chunk.as_ptr()) };
sum = vaddq_f32(sum, vabsq_f32(v));
}
let mut out = vaddvq_f32(sum);
for &v in rem {
out += v.abs();
}
out
}