ferrox-quant 0.19.1

use super::{sdot, sdot_lane, vmmla_s32};
use crate::repack::common::*;
use crate::repack::q5_kx8::*;
use crate::{Q8KActivations, Q5_K_BLOCK_ELEMS};
use std::arch::aarch64::*;

/// NEON DotProd GEMV for interleave-4 `block_q5_Kx8` (llama
/// `ggml_gemv_q5_K_8x4_q8_K`).
#[target_feature(enable = "neon,dotprod")]
pub unsafe fn gemv_q5_kx8_q8_k_neon_sdot(
    packed: &[u8],
    act: &Q8KActivations,
    n_cols: usize,
    n_row_groups: usize,
    out: &mut [f32],
) {
    let nb = n_cols / Q5_K_BLOCK_ELEMS;
    let m4b = vdupq_n_u8(0x0f);
    let mone = vdupq_n_u8(1);
    let mtwo = vdupq_n_u8(2);

    for x in 0..n_row_groups {
        let mut acc_f32 = [vdupq_n_f32(0.0), vdupq_n_f32(0.0)];
        let group_off = x * nb * Q5_KX8_BLOCK_BYTES;

        for b in 0..nb {
            let blk = packed.as_ptr().add(group_off + b * Q5_KX8_BLOCK_BYTES);
            let mut d_arr = [0f32; 8];
            let mut dmin_arr = [0f32; 8];
            for j in 0..8 {
                d_arr[j] = f16_from_bytes(std::slice::from_raw_parts(blk.add(j * 2), 2));
                dmin_arr[j] = f16_from_bytes(std::slice::from_raw_parts(blk.add(16 + j * 2), 2));
            }
            let q8_d = act.d[b];
            let sb_scale_0123 = vmulq_n_f32(vld1q_f32(d_arr.as_ptr()), q8_d);
            let sb_scale_4567 = vmulq_n_f32(vld1q_f32(d_arr.as_ptr().add(4)), q8_d);
            let sb_min_0123 = vmulq_n_f32(vld1q_f32(dmin_arr.as_ptr()), q8_d);
            let sb_min_4567 = vmulq_n_f32(vld1q_f32(dmin_arr.as_ptr().add(4)), q8_d);

            let mut bias_acc = [vdupq_n_s32(0), vdupq_n_s32(0)];
            let q8_base = act.q.as_ptr().add(b * Q5_K_BLOCK_ELEMS);
            let bsums_ptr = act.bsums.as_ptr().add(b * 16);
            let mut bsums_arr = [0i16; 8];
            for (i, slot) in bsums_arr.iter_mut().enumerate() {
                *slot = *bsums_ptr.add(2 * i) + *bsums_ptr.add(2 * i + 1);
            }

            let scales_base = blk.add(32);
            let qh_base = blk.add(128);
            let qs_base = blk.add(384);

            // qh[c][i]: 2 col-groups × 8 vectors; shift down 2 bits each sb.
            let mut qh = [[vdupq_n_u8(0); 8]; 2];
            for (c, qh_c) in qh.iter_mut().enumerate() {
                for (i, slot) in qh_c.iter_mut().enumerate() {
                    *slot = vld1q_u8(qh_base.add(i * 32 + 16 * c));
                }
            }

            for sb in 0..4 {
                let mut acc_lo = [vdupq_n_s32(0), vdupq_n_s32(0)];
                let mut acc_hi = [vdupq_n_s32(0), vdupq_n_s32(0)];

                let mut q5sb_mins = [vdupq_n_s16(0); 2];
                let mut q5sb_scales = [vdupq_n_s16(0); 2];
                for i in 0..2 {
                    let mut sc = [0u8; 8];
                    let mut mn = [0u8; 8];
                    let offset = sb * 24 + i * 12;
                    decode_scales_mins(
                        std::slice::from_raw_parts(scales_base.add(offset), 12),
                        &mut sc,
                        &mut mn,
                    );
                    let mut sc_i8 = [0i8; 8];
                    let mut mn_i8 = [0i8; 8];
                    for t in 0..8 {
                        sc_i8[t] = sc[t] as i8;
                        mn_i8[t] = mn[t] as i8;
                    }
                    q5sb_scales[i] = vmovl_s8(vld1_s8(sc_i8.as_ptr()));
                    q5sb_mins[i] = vmovl_s8(vld1_s8(mn_i8.as_ptr()));
                }

                let mut q8_qs = [vdupq_n_s8(0); 4];
                for (i, slot) in q8_qs.iter_mut().enumerate() {
                    *slot = vld1q_s8(q8_base.add(sb * 64 + i * 16));
                }

                for c in 0..2 {
                    let mut q5_lo = [vdupq_n_s8(0); 8];
                    let mut q5_hi = [vdupq_n_s8(0); 8];
                    for i in 0..8 {
                        let q5_cols =
                            vld1q_u8(qs_base.add(sb * Q5_K_BLOCK_ELEMS + i * 32 + 16 * c));
                        let hbit_lo = vandq_u8(qh[c][i], mone);
                        let hbit_hi = vshlq_n_u8(vandq_u8(qh[c][i], mtwo), 3);
                        qh[c][i] = vshrq_n_u8(qh[c][i], 2);
                        q5_lo[i] =
                            vreinterpretq_s8_u8(vsliq_n_u8(vandq_u8(q5_cols, m4b), hbit_lo, 4));
                        q5_hi[i] = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q5_cols, 4), hbit_hi));
                    }
                    acc_lo[c] = sdot_lane(acc_lo[c], q5_lo[0], q8_qs[0], 0);
                    acc_lo[c] = sdot_lane(acc_lo[c], q5_lo[1], q8_qs[0], 1);
                    acc_lo[c] = sdot_lane(acc_lo[c], q5_lo[2], q8_qs[0], 2);
                    acc_lo[c] = sdot_lane(acc_lo[c], q5_lo[3], q8_qs[0], 3);
                    acc_lo[c] = sdot_lane(acc_lo[c], q5_lo[4], q8_qs[1], 0);
                    acc_lo[c] = sdot_lane(acc_lo[c], q5_lo[5], q8_qs[1], 1);
                    acc_lo[c] = sdot_lane(acc_lo[c], q5_lo[6], q8_qs[1], 2);
                    acc_lo[c] = sdot_lane(acc_lo[c], q5_lo[7], q8_qs[1], 3);
                    acc_hi[c] = sdot_lane(acc_hi[c], q5_hi[0], q8_qs[2], 0);
                    acc_hi[c] = sdot_lane(acc_hi[c], q5_hi[1], q8_qs[2], 1);
                    acc_hi[c] = sdot_lane(acc_hi[c], q5_hi[2], q8_qs[2], 2);
                    acc_hi[c] = sdot_lane(acc_hi[c], q5_hi[3], q8_qs[2], 3);
                    acc_hi[c] = sdot_lane(acc_hi[c], q5_hi[4], q8_qs[3], 0);
                    acc_hi[c] = sdot_lane(acc_hi[c], q5_hi[5], q8_qs[3], 1);
                    acc_hi[c] = sdot_lane(acc_hi[c], q5_hi[6], q8_qs[3], 2);
                    acc_hi[c] = sdot_lane(acc_hi[c], q5_hi[7], q8_qs[3], 3);
                }

                let sc_0123_lo = vget_low_s16(q5sb_scales[0]);
                let sc_0123_hi = vget_low_s16(q5sb_scales[1]);
                let sumf_0123 = vcvtq_f32_s32(vaddq_s32(
                    vmulq_s32(vmovl_s16(sc_0123_lo), acc_lo[0]),
                    vmulq_s32(vmovl_s16(sc_0123_hi), acc_hi[0]),
                ));
                acc_f32[0] = vfmaq_f32(acc_f32[0], sb_scale_0123, sumf_0123);

                let sc_4567_lo = vget_high_s16(q5sb_scales[0]);
                let sc_4567_hi = vget_high_s16(q5sb_scales[1]);
                let sumf_4567 = vcvtq_f32_s32(vaddq_s32(
                    vmulq_s32(vmovl_s16(sc_4567_lo), acc_lo[1]),
                    vmulq_s32(vmovl_s16(sc_4567_hi), acc_hi[1]),
                ));
                acc_f32[1] = vfmaq_f32(acc_f32[1], sb_scale_4567, sumf_4567);

                let bsums_vec_lo = vdup_n_s16(bsums_arr[2 * sb]);
                let bsums_vec_hi = vdup_n_s16(bsums_arr[2 * sb + 1]);
                bias_acc[0] = vmlal_s16(bias_acc[0], bsums_vec_lo, vget_low_s16(q5sb_mins[0]));
                bias_acc[0] = vmlal_s16(bias_acc[0], bsums_vec_hi, vget_low_s16(q5sb_mins[1]));
                bias_acc[1] = vmlal_s16(bias_acc[1], bsums_vec_lo, vget_high_s16(q5sb_mins[0]));
                bias_acc[1] = vmlal_s16(bias_acc[1], bsums_vec_hi, vget_high_s16(q5sb_mins[1]));
            }

            acc_f32[0] = vmlsq_f32(acc_f32[0], vcvtq_f32_s32(bias_acc[0]), sb_min_0123);
            acc_f32[1] = vmlsq_f32(acc_f32[1], vcvtq_f32_s32(bias_acc[1]), sb_min_4567);
        }

        let base = x * Q5_KX8_NROWS;
        vst1q_f32(out.as_mut_ptr().add(base), acc_f32[0]);
        vst1q_f32(out.as_mut_ptr().add(base + 4), acc_f32[1]);
    }
}

/// NEON DotProd GEMV for interleave-8 packed Q5_K weights (llama.cpp
/// `ggml_gemv_q5_K_8x8_q8_K` in `arch/arm/repack.cpp`). Each 8-byte q8
/// run is broadcast to both vector halves so one `sdot` covers two
/// interleaved columns at once.
#[target_feature(enable = "neon,dotprod")]
pub unsafe fn gemv_q5_kx8_q8_k_neon_8x8(
    packed: &[u8],
    act: &Q8KActivations,
    n_cols: usize,
    n_row_groups: usize,
    out: &mut [f32],
) {
    let nb = n_cols / Q5_K_BLOCK_ELEMS;
    let m4b = vdupq_n_u8(0x0f);
    let mone = vdupq_n_u8(1);
    let mtwo = vdupq_n_u8(2);

    for x in 0..n_row_groups {
        let mut acc_f32 = [vdupq_n_f32(0.0), vdupq_n_f32(0.0)];
        let group_off = x * nb * Q5_KX8_BLOCK_BYTES;

        for b in 0..nb {
            let blk = packed.as_ptr().add(group_off + b * Q5_KX8_BLOCK_BYTES);
            let mut d_arr = [0f32; 8];
            let mut dmin_arr = [0f32; 8];
            for j in 0..8 {
                d_arr[j] = f16_from_bytes(std::slice::from_raw_parts(blk.add(j * 2), 2));
                dmin_arr[j] = f16_from_bytes(std::slice::from_raw_parts(blk.add(16 + j * 2), 2));
            }
            let q8_d = act.d[b];
            let sb_scale = [
                vmulq_n_f32(vld1q_f32(d_arr.as_ptr()), q8_d),
                vmulq_n_f32(vld1q_f32(d_arr.as_ptr().add(4)), q8_d),
            ];
            let sb_min = [
                vmulq_n_f32(vld1q_f32(dmin_arr.as_ptr()), q8_d),
                vmulq_n_f32(vld1q_f32(dmin_arr.as_ptr().add(4)), q8_d),
            ];

            let q8_base = act.q.as_ptr().add(b * Q5_K_BLOCK_ELEMS);
            let bsums_ptr = act.bsums.as_ptr().add(b * 16);
            let mut bsums_arr = [0i16; 8];
            for (i, slot) in bsums_arr.iter_mut().enumerate() {
                *slot = *bsums_ptr.add(2 * i) + *bsums_ptr.add(2 * i + 1);
            }

            let scales_base = blk.add(32);
            let qh_base = blk.add(128);
            let qs_base = blk.add(384);

            // qh state per column pair; two bits consumed per sub-block.
            let mut qh = [[vdupq_n_u8(0); 4]; 4];
            for (cp, qh_cp) in qh.iter_mut().enumerate() {
                for (m, slot) in qh_cp.iter_mut().enumerate() {
                    *slot = vld1q_u8(qh_base.add(16 * cp + 64 * m));
                }
            }

            for sb in 0..4 {
                let mut acc_lo = [vdupq_n_s32(0); 4];
                let mut acc_hi = [vdupq_n_s32(0); 4];

                let mut q5sb_scales = [vdupq_n_s16(0); 2];
                let mut q5sb_mins = [vdupq_n_s16(0); 2];
                for i in 0..2 {
                    let mut sc = [0u8; 8];
                    let mut mn = [0u8; 8];
                    let offset = sb * 24 + i * 12;
                    decode_scales_mins(
                        std::slice::from_raw_parts(scales_base.add(offset), 12),
                        &mut sc,
                        &mut mn,
                    );
                    let mut sc_i8 = [0i8; 8];
                    let mut mn_i8 = [0i8; 8];
                    for t in 0..8 {
                        sc_i8[t] = sc[t] as i8;
                        mn_i8[t] = mn[t] as i8;
                    }
                    q5sb_scales[i] = vmovl_s8(vld1_s8(sc_i8.as_ptr()));
                    q5sb_mins[i] = vmovl_s8(vld1_s8(mn_i8.as_ptr()));
                }

                let q8_sb = q8_base.add(sb * 64);
                let mut q8_qs = [vdupq_n_s8(0); 8];
                for (i, slot) in q8_qs.iter_mut().enumerate() {
                    *slot = vreinterpretq_s8_s64(vld1q_dup_s64(q8_sb.add(i * 8) as *const i64));
                }

                for cp in 0..4 {
                    let q5_qs = [
                        vld1q_u8(qs_base.add(sb * Q5_K_BLOCK_ELEMS + 16 * cp)),
                        vld1q_u8(qs_base.add(sb * Q5_K_BLOCK_ELEMS + 16 * cp + 64)),
                        vld1q_u8(qs_base.add(sb * Q5_K_BLOCK_ELEMS + 16 * cp + 128)),
                        vld1q_u8(qs_base.add(sb * Q5_K_BLOCK_ELEMS + 16 * cp + 192)),
                    ];
                    for m in 0..4 {
                        let hbit_lo = vandq_u8(qh[cp][m], mone);
                        let hbit_hi = vshlq_n_u8(vandq_u8(qh[cp][m], mtwo), 3);
                        qh[cp][m] = vshrq_n_u8(qh[cp][m], 2);
                        let q5_lo =
                            vreinterpretq_s8_u8(vsliq_n_u8(vandq_u8(q5_qs[m], m4b), hbit_lo, 4));
                        let q5_hi = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q5_qs[m], 4), hbit_hi));
                        acc_lo[cp] = sdot(acc_lo[cp], q5_lo, q8_qs[m]);
                        acc_hi[cp] = sdot(acc_hi[cp], q5_hi, q8_qs[m + 4]);
                    }
                }

                let bsums_vec_lo = vdup_n_s16(bsums_arr[2 * sb]);
                let bsums_vec_hi = vdup_n_s16(bsums_arr[2 * sb + 1]);
                for i in 0..2 {
                    let p = i * 2;
                    let (scales_lo, scales_hi, mins_lo, mins_hi) = if i == 0 {
                        (
                            vget_low_s16(q5sb_scales[0]),
                            vget_low_s16(q5sb_scales[1]),
                            vget_low_s16(q5sb_mins[0]),
                            vget_low_s16(q5sb_mins[1]),
                        )
                    } else {
                        (
                            vget_high_s16(q5sb_scales[0]),
                            vget_high_s16(q5sb_scales[1]),
                            vget_high_s16(q5sb_mins[0]),
                            vget_high_s16(q5sb_mins[1]),
                        )
                    };
                    let sumf_0 = vcvtq_f32_s32(vmulq_s32(
                        vmovl_s16(scales_lo),
                        vpaddq_s32(acc_lo[p], acc_lo[p + 1]),
                    ));
                    acc_f32[i] = vfmaq_f32(acc_f32[i], sb_scale[i], sumf_0);
                    let sumf_1 = vcvtq_f32_s32(vmulq_s32(
                        vmovl_s16(scales_hi),
                        vpaddq_s32(acc_hi[p], acc_hi[p + 1]),
                    ));
                    acc_f32[i] = vfmaq_f32(acc_f32[i], sb_scale[i], sumf_1);

                    let mut bias = vmull_s16(bsums_vec_lo, mins_lo);
                    bias = vmlal_s16(bias, bsums_vec_hi, mins_hi);
                    acc_f32[i] = vmlsq_f32(acc_f32[i], sb_min[i], vcvtq_f32_s32(bias));
                }
            }
        }

        let base = x * Q5_KX8_NROWS;
        vst1q_f32(out.as_mut_ptr().add(base), acc_f32[0]);
        vst1q_f32(out.as_mut_ptr().add(base + 4), acc_f32[1]);
    }
}

/// NEON DotProd **GEMM** for interleave-4 `block_q5_Kx8` — weight unpack
/// once per act tile (llama `ggml_gemm_q5_K_8x4_q8_K` motivation).
#[target_feature(enable = "neon,dotprod")]
pub unsafe fn gemm_q5_kx8_q8_k_neon_sdot(
    packed: &[u8],
    acts: &[Q8KActivations],
    n_cols: usize,
    out: &mut [f32],
) {
    let na = acts.len();
    debug_assert!(na <= Q5_KX8_GEMM_NC);
    let nb = n_cols / Q5_K_BLOCK_ELEMS;
    let m4b = vdupq_n_u8(0x0f);
    let mone = vdupq_n_u8(1);
    let mtwo = vdupq_n_u8(2);

    let mut acc_f32 = [[vdupq_n_f32(0.0); 2]; Q5_KX8_GEMM_NC];
    let mut bias_acc = [[vdupq_n_s32(0); 2]; Q5_KX8_GEMM_NC];

    for b in 0..nb {
        let blk = packed.as_ptr().add(b * Q5_KX8_BLOCK_BYTES);
        let mut d_arr = [0f32; 8];
        let mut dmin_arr = [0f32; 8];
        for j in 0..8 {
            d_arr[j] = f16_from_bytes(std::slice::from_raw_parts(blk.add(j * 2), 2));
            dmin_arr[j] = f16_from_bytes(std::slice::from_raw_parts(blk.add(16 + j * 2), 2));
        }
        let d_lo = vld1q_f32(d_arr.as_ptr());
        let d_hi = vld1q_f32(d_arr.as_ptr().add(4));
        let dmin_lo = vld1q_f32(dmin_arr.as_ptr());
        let dmin_hi = vld1q_f32(dmin_arr.as_ptr().add(4));

        let mut sb_scale = [[vdupq_n_f32(0.0); 2]; Q5_KX8_GEMM_NC];
        let mut sb_min = [[vdupq_n_f32(0.0); 2]; Q5_KX8_GEMM_NC];
        let mut bsums_arr = [[0i16; 8]; Q5_KX8_GEMM_NC];
        for (a, act) in acts.iter().enumerate() {
            let q8_d = act.d[b];
            sb_scale[a] = [vmulq_n_f32(d_lo, q8_d), vmulq_n_f32(d_hi, q8_d)];
            sb_min[a] = [vmulq_n_f32(dmin_lo, q8_d), vmulq_n_f32(dmin_hi, q8_d)];
            let bsums_ptr = act.bsums.as_ptr().add(b * 16);
            for (i, slot) in bsums_arr[a].iter_mut().enumerate() {
                *slot = *bsums_ptr.add(2 * i) + *bsums_ptr.add(2 * i + 1);
            }
        }

        let scales_base = blk.add(32);
        let qh_base = blk.add(128);
        let qs_base = blk.add(384);

        let mut qh = [[vdupq_n_u8(0); 8]; 2];
        for (c, qh_c) in qh.iter_mut().enumerate() {
            for (i, slot) in qh_c.iter_mut().enumerate() {
                *slot = vld1q_u8(qh_base.add(i * 32 + 16 * c));
            }
        }

        for sb in 0..4 {
            let mut q5sb_mins = [vdupq_n_s16(0); 2];
            let mut q5sb_scales = [vdupq_n_s16(0); 2];
            for i in 0..2 {
                let mut sc = [0u8; 8];
                let mut mn = [0u8; 8];
                let offset = sb * 24 + i * 12;
                decode_scales_mins(
                    std::slice::from_raw_parts(scales_base.add(offset), 12),
                    &mut sc,
                    &mut mn,
                );
                let mut sc_i8 = [0i8; 8];
                let mut mn_i8 = [0i8; 8];
                for t in 0..8 {
                    sc_i8[t] = sc[t] as i8;
                    mn_i8[t] = mn[t] as i8;
                }
                q5sb_scales[i] = vmovl_s8(vld1_s8(sc_i8.as_ptr()));
                q5sb_mins[i] = vmovl_s8(vld1_s8(mn_i8.as_ptr()));
            }

            for c in 0..2 {
                let mut q5_lo = [vdupq_n_s8(0); 8];
                let mut q5_hi = [vdupq_n_s8(0); 8];
                for i in 0..8 {
                    let q5_cols = vld1q_u8(qs_base.add(sb * Q5_K_BLOCK_ELEMS + i * 32 + 16 * c));
                    let hbit_lo = vandq_u8(qh[c][i], mone);
                    let hbit_hi = vshlq_n_u8(vandq_u8(qh[c][i], mtwo), 3);
                    qh[c][i] = vshrq_n_u8(qh[c][i], 2);
                    q5_lo[i] = vreinterpretq_s8_u8(vsliq_n_u8(vandq_u8(q5_cols, m4b), hbit_lo, 4));
                    q5_hi[i] = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q5_cols, 4), hbit_hi));
                }
                let (sc_lo, sc_hi) = if c == 0 {
                    (vget_low_s16(q5sb_scales[0]), vget_low_s16(q5sb_scales[1]))
                } else {
                    (vget_high_s16(q5sb_scales[0]), vget_high_s16(q5sb_scales[1]))
                };
                let sc_lo_w = vmovl_s16(sc_lo);
                let sc_hi_w = vmovl_s16(sc_hi);

                for a in 0..na {
                    let q8_base = acts[a].q.as_ptr().add(b * Q5_K_BLOCK_ELEMS);
                    let y0 = vld1q_s8(q8_base.add(sb * 64));
                    let y1 = vld1q_s8(q8_base.add(sb * 64 + 16));
                    let y2 = vld1q_s8(q8_base.add(sb * 64 + 32));
                    let y3 = vld1q_s8(q8_base.add(sb * 64 + 48));
                    let mut acc_lo = vdupq_n_s32(0);
                    let mut acc_hi = vdupq_n_s32(0);
                    acc_lo = sdot_lane(acc_lo, q5_lo[0], y0, 0);
                    acc_lo = sdot_lane(acc_lo, q5_lo[1], y0, 1);
                    acc_lo = sdot_lane(acc_lo, q5_lo[2], y0, 2);
                    acc_lo = sdot_lane(acc_lo, q5_lo[3], y0, 3);
                    acc_lo = sdot_lane(acc_lo, q5_lo[4], y1, 0);
                    acc_lo = sdot_lane(acc_lo, q5_lo[5], y1, 1);
                    acc_lo = sdot_lane(acc_lo, q5_lo[6], y1, 2);
                    acc_lo = sdot_lane(acc_lo, q5_lo[7], y1, 3);
                    acc_hi = sdot_lane(acc_hi, q5_hi[0], y2, 0);
                    acc_hi = sdot_lane(acc_hi, q5_hi[1], y2, 1);
                    acc_hi = sdot_lane(acc_hi, q5_hi[2], y2, 2);
                    acc_hi = sdot_lane(acc_hi, q5_hi[3], y2, 3);
                    acc_hi = sdot_lane(acc_hi, q5_hi[4], y3, 0);
                    acc_hi = sdot_lane(acc_hi, q5_hi[5], y3, 1);
                    acc_hi = sdot_lane(acc_hi, q5_hi[6], y3, 2);
                    acc_hi = sdot_lane(acc_hi, q5_hi[7], y3, 3);
                    let sumf = vcvtq_f32_s32(vaddq_s32(
                        vmulq_s32(sc_lo_w, acc_lo),
                        vmulq_s32(sc_hi_w, acc_hi),
                    ));
                    acc_f32[a][c] = vfmaq_f32(acc_f32[a][c], sb_scale[a][c], sumf);
                }
            }

            for a in 0..na {
                let bs_lo = vdup_n_s16(bsums_arr[a][2 * sb]);
                let bs_hi = vdup_n_s16(bsums_arr[a][2 * sb + 1]);
                bias_acc[a][0] = vmlal_s16(bias_acc[a][0], bs_lo, vget_low_s16(q5sb_mins[0]));
                bias_acc[a][0] = vmlal_s16(bias_acc[a][0], bs_hi, vget_low_s16(q5sb_mins[1]));
                bias_acc[a][1] = vmlal_s16(bias_acc[a][1], bs_lo, vget_high_s16(q5sb_mins[0]));
                bias_acc[a][1] = vmlal_s16(bias_acc[a][1], bs_hi, vget_high_s16(q5sb_mins[1]));
            }
        }

        for a in 0..na {
            for c in 0..2 {
                acc_f32[a][c] =
                    vmlsq_f32(acc_f32[a][c], vcvtq_f32_s32(bias_acc[a][c]), sb_min[a][c]);
                bias_acc[a][c] = vdupq_n_s32(0);
            }
        }
    }

    for a in 0..na {
        let mut row = [0f32; Q5_KX8_NROWS];
        vst1q_f32(row.as_mut_ptr(), acc_f32[a][0]);
        vst1q_f32(row.as_mut_ptr().add(4), acc_f32[a][1]);
        for (r, v) in row.iter().enumerate() {
            out[r * na + a] = *v;
        }
    }
}

/// NEON i8mm **GEMM** for interleave-8 packed Q5_K weights (llama.cpp
/// `ggml_gemm_q5_K_8x8_q8_K` in `arch/arm/repack.cpp`). Same 2×8×8
/// `vmmlaq_s32` tiling as [`gemm_q4_kx8_q8_k_neon_i8mm`]; the only
/// difference is splicing the fifth bit in from `qh` before the MMLAs,
/// two bits consumed per sub-block.
#[target_feature(enable = "neon,i8mm")]
pub unsafe fn gemm_q5_kx8_q8_k_neon_i8mm(
    packed: &[u8],
    tile: &Q8KActsX4,
    n_cols: usize,
    out: &mut [f32],
) {
    let na = tile.na;
    debug_assert!(na <= Q5_KX8_GEMM_NC);
    let nb = n_cols / Q5_K_BLOCK_ELEMS;
    debug_assert_eq!(tile.n_blocks, nb);
    let m4b = vdupq_n_u8(0x0f);
    let mone = vdupq_n_u8(1);
    let mtwo = vdupq_n_u8(2);
    const Q8_K_BLOCKLEN: usize = 4;

    let mut acc_f32 = [vdupq_n_f32(0.0); Q5_KX8_GEMM_NC * 2];

    for b in 0..nb {
        let blk = packed.as_ptr().add(b * Q5_KX8_BLOCK_BYTES);
        let bsums_base = tile.bsums.as_ptr().add(b * Q8_K_BLOCKLEN * 8);

        let mut acc = [vdupq_n_s32(0); 8];
        let mut bias_acc = [vdupq_n_s32(0); 8];

        let scales_base = blk.add(32);
        let qh_base = blk.add(128);
        let qs_base = blk.add(384);
        let q8_base = tile.qs.as_ptr().add(b * Q5_K_BLOCK_ELEMS * 4);

        // qh state per column pair; two bits consumed per sub-block.
        let mut qh = [[vdupq_n_u8(0); 4]; 4];
        for (cp, qh_cp) in qh.iter_mut().enumerate() {
            for (m, slot) in qh_cp.iter_mut().enumerate() {
                *slot = vld1q_u8(qh_base.add(16 * cp + 64 * m));
            }
        }

        for sb in 0..4 {
            let mut q5sb_scales = [[0i8; 8]; 2];
            let mut q5sb_mins = [vdupq_n_s16(0); 2];
            for i in 0..2 {
                let mut sc = [0u8; 8];
                let mut mn = [0u8; 8];
                let offset = sb * 24 + i * 12;
                decode_scales_mins(
                    std::slice::from_raw_parts(scales_base.add(offset), 12),
                    &mut sc,
                    &mut mn,
                );
                let mut mn_i8 = [0i8; 8];
                for t in 0..8 {
                    q5sb_scales[i][t] = sc[t] as i8;
                    mn_i8[t] = mn[t] as i8;
                }
                q5sb_mins[i] = vmovl_s8(vld1_s8(mn_i8.as_ptr()));
            }

            let q8_sb = q8_base.add(sb * 256);
            let mut q8_qs_01 = [vdupq_n_s8(0); 8];
            let mut q8_qs_23 = [vdupq_n_s8(0); 8];
            for i in 0..8 {
                q8_qs_01[i] = vld1q_s8(q8_sb.add(i * 32));
                q8_qs_23[i] = vld1q_s8(q8_sb.add(i * 32 + 16));
            }
            let q8s = [q8_qs_01, q8_qs_23];

            for cp in 0..4 {
                let mut sb_acc = [vdupq_n_s32(0); 4];

                let q5_qs = [
                    vld1q_u8(qs_base.add(sb * Q5_K_BLOCK_ELEMS + 16 * cp)),
                    vld1q_u8(qs_base.add(sb * Q5_K_BLOCK_ELEMS + 16 * cp + 64)),
                    vld1q_u8(qs_base.add(sb * Q5_K_BLOCK_ELEMS + 16 * cp + 128)),
                    vld1q_u8(qs_base.add(sb * Q5_K_BLOCK_ELEMS + 16 * cp + 192)),
                ];
                let mut q5_lo = [vdupq_n_s8(0); 4];
                let mut q5_hi = [vdupq_n_s8(0); 4];
                for m in 0..4 {
                    let hbit_lo = vandq_u8(qh[cp][m], mone);
                    let hbit_hi = vshlq_n_u8(vandq_u8(qh[cp][m], mtwo), 3);
                    qh[cp][m] = vshrq_n_u8(qh[cp][m], 2);
                    q5_lo[m] = vreinterpretq_s8_u8(vsliq_n_u8(vandq_u8(q5_qs[m], m4b), hbit_lo, 4));
                    q5_hi[m] = vreinterpretq_s8_u8(vorrq_u8(vshrq_n_u8(q5_qs[m], 4), hbit_hi));
                }
                let q5_vals = [q5_lo, q5_hi];

                for rp in 0..2 {
                    for half in 0..2 {
                        let q8 = &q8s[rp][4 * half..4 * half + 4];
                        let q5 = &q5_vals[half];
                        let mut tile_acc = sb_acc[2 * rp + half];
                        for m in 0..4 {
                            tile_acc = vmmla_s32(tile_acc, q5[m], q8[m]);
                        }
                        sb_acc[2 * rp + half] = tile_acc;
                    }
                }

                let scale_offset = cp * 2;
                let block_scale_0 = vcombine_s32(
                    vdup_n_s32(i32::from(q5sb_scales[0][scale_offset])),
                    vdup_n_s32(i32::from(q5sb_scales[0][scale_offset + 1])),
                );
                let block_scale_1 = vcombine_s32(
                    vdup_n_s32(i32::from(q5sb_scales[1][scale_offset])),
                    vdup_n_s32(i32::from(q5sb_scales[1][scale_offset + 1])),
                );

                acc[cp] = vmlaq_s32(acc[cp], sb_acc[0], block_scale_0);
                acc[cp + 4] = vmlaq_s32(acc[cp + 4], sb_acc[2], block_scale_0);
                acc[cp] = vmlaq_s32(acc[cp], sb_acc[1], block_scale_1);
                acc[cp + 4] = vmlaq_s32(acc[cp + 4], sb_acc[3], block_scale_1);
            }

            for q8_row in 0..Q8_K_BLOCKLEN {
                let bs_lo = vdup_n_s16(*bsums_base.add(q8_row * 8 + 2 * sb));
                let bs_hi = vdup_n_s16(*bsums_base.add(q8_row * 8 + 2 * sb + 1));
                bias_acc[2 * q8_row] =
                    vmlal_s16(bias_acc[2 * q8_row], bs_lo, vget_low_s16(q5sb_mins[0]));
                bias_acc[2 * q8_row] =
                    vmlal_s16(bias_acc[2 * q8_row], bs_hi, vget_low_s16(q5sb_mins[1]));
                bias_acc[2 * q8_row + 1] =
                    vmlal_s16(bias_acc[2 * q8_row + 1], bs_lo, vget_high_s16(q5sb_mins[0]));
                bias_acc[2 * q8_row + 1] =
                    vmlal_s16(bias_acc[2 * q8_row + 1], bs_hi, vget_high_s16(q5sb_mins[1]));
            }
        }

        for lane in acc.iter_mut() {
            let aux = vzip_s32(vget_low_s32(*lane), vget_high_s32(*lane));
            *lane = vcombine_s32(aux.0, aux.1);
        }
        let reorder_acc = [
            vcombine_s32(vget_low_s32(acc[0]), vget_low_s32(acc[1])),
            vcombine_s32(vget_low_s32(acc[2]), vget_low_s32(acc[3])),
            vcombine_s32(vget_high_s32(acc[0]), vget_high_s32(acc[1])),
            vcombine_s32(vget_high_s32(acc[2]), vget_high_s32(acc[3])),
            vcombine_s32(vget_low_s32(acc[4]), vget_low_s32(acc[5])),
            vcombine_s32(vget_low_s32(acc[6]), vget_low_s32(acc[7])),
            vcombine_s32(vget_high_s32(acc[4]), vget_high_s32(acc[5])),
            vcombine_s32(vget_high_s32(acc[6]), vget_high_s32(acc[7])),
        ];

        let mut d_arr = [0f32; 8];
        let mut dmin_arr = [0f32; 8];
        for j in 0..8 {
            d_arr[j] = f16_from_bytes(std::slice::from_raw_parts(blk.add(j * 2), 2));
            dmin_arr[j] = f16_from_bytes(std::slice::from_raw_parts(blk.add(16 + j * 2), 2));
        }

        for i in 0..na {
            for j in 0..2 {
                let q8_d = vdupq_n_f32(*tile.d.as_ptr().add(b * Q8_K_BLOCKLEN + i));
                let dmins = vmulq_f32(vld1q_f32(dmin_arr.as_ptr().add(j * 4)), q8_d);
                let scale = vmulq_f32(vld1q_f32(d_arr.as_ptr().add(j * 4)), q8_d);
                let idx = 2 * i + j;
                acc_f32[idx] = vmlsq_f32(acc_f32[idx], vcvtq_f32_s32(bias_acc[idx]), dmins);
                acc_f32[idx] = vmlaq_f32(acc_f32[idx], vcvtq_f32_s32(reorder_acc[idx]), scale);
            }
        }
    }

    for a in 0..na {
        let mut row = [0f32; Q5_KX8_NROWS];
        vst1q_f32(row.as_mut_ptr(), acc_f32[2 * a]);
        vst1q_f32(row.as_mut_ptr().add(4), acc_f32[2 * a + 1]);
        for (r, v) in row.iter().enumerate() {
            out[r * na + a] = *v;
        }
    }
}