//! Fused quantized operations for K-quantization formats (PMAT-802)

//!

//! Implements fused dequant+dot operations that eliminate intermediate f32 buffer

//! allocation for 8x memory bandwidth reduction.

//!

//! Per llama-cpp-style-performance-spec.md:

//! - Memory wall is the bottleneck (Wulf & McKee [10])

//! - Fused operations keep data in registers, avoid memory round-trips

//! - ULP tolerance of ≤4 for numerical equivalence (Goldberg [9])

//!

//! Functions:

//! - `fused_q4k_dot`, `fused_q4k_dot_simd` - Q4_K dot products

//! - `fused_q4k_q8_dot` - Q4_K with Q8 activations

//! - `fused_q4k_q8k_dot`, `fused_q4k_q8k_dot_simd` - Q4_K with Q8_K activations

//! - `fused_q4k_q8k_parallel_matvec_into` - Parallel matrix-vector multiply

//! - `fused_q4k_q8k_ffn_up_gate_into` - Fused FFN up/gate computation

//! - `fused_q6k_dot`, `fused_q6k_dot_simd` - Q6_K dot products

//! - `fused_q5k_dot`, `fused_q5k_dot_simd` - Q5_K dot products

use crate::error::{RealizarError, Result};

use super::dequant::read_f16;

use super::simd::extract_scale_min;

use super::types::QK_K;

/// Fused Q4_K dequantize + dot product

///

/// Computes the dot product of Q4_K quantized weights with f32 activations

/// WITHOUT allocating an intermediate f32 buffer. Dequantization happens

/// inline, accumulating directly into a register.

///

/// # Arguments

///

/// * `q4k_data` - Raw Q4_K quantized data (super-blocks of 144 bytes)

/// * `activations` - f32 activation values (must match dequantized length)

///

/// # Returns

///

/// The dot product as f32

///

/// # Errors

///

/// Returns error if:

/// - `q4k_data` length is not a multiple of 144 bytes (super-block size)

/// - `activations` length doesn't match the number of quantized values

///

/// # Performance

///

/// This function reduces memory traffic by 8x compared to separate

/// dequantize-then-dot operations:

/// - Naive: Read Q4_K (4.5 bits) → Write f32 (32 bits) → Read f32 → Compute

/// - Fused: Read Q4_K (4.5 bits) → Compute in registers

///

/// # Examples

///

/// ```rust,ignore

/// let weights_q4k = load_q4k_weights();

/// let activations = get_layer_activations();

/// let result = fused_q4k_dot(&weights_q4k, &activations)?;

/// ```

pub fn fused_q4k_dot(q4k_data: &[u8], activations: &[f32]) -> Result<f32> {

    const SUPER_BLOCK_BYTES: usize = 144;

    // Validate Q4_K data length

    if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Q4_K data length {} is not a multiple of super-block size {}",

                q4k_data.len(),

                SUPER_BLOCK_BYTES

            ),

        });

    }

    let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;

    let expected_values = num_super_blocks * QK_K;

    // Validate activation length matches

    if activations.len() != expected_values {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Activation length {} doesn't match Q4_K values count {}",

                activations.len(),

                expected_values

            ),

        });

    }

    // Accumulator for dot product result

    let mut acc = 0.0f32;

    let mut activation_idx = 0;

    for sb_idx in 0..
num_super_blocks1.16k
 {

        let sb_start = sb_idx * SUPER_BLOCK_BYTES;

        // Read d (f16 -> f32)

        let d = read_f16(&q4k_data[sb_start..sb_start + 2]);

        // Read dmin (f16 -> f32)

        let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);

        // Read scales (12 bytes)

        let mut scales = [0u8; 12];

        scales.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);

        // Read qs (128 bytes)

        let qs_start = sb_start + 16;

        let qs = &q4k_data[qs_start..qs_start + 128];

        // PAR-001: Match dequantize_q4_k layout (llama.cpp/candle compatible)

        // Process 4 chunks of 64 values each (0, 64, 128, 192)

        // Each chunk: 32 low nibbles, then 32 high nibbles from 32 consecutive bytes

        for j in 
(0..QK_K)16.3k
.
step_by16.3k
(64) {

            let q = &qs[j / 2..j / 2 + 32];

            // Get scales for the two 32-value halves

            let is = j / 32;

            let (sc1, m1) = extract_scale_min(&scales, is);

            let d1 = d * sc1;

            let dm1 = dmin * m1;

            let (sc2, m2) = extract_scale_min(&scales, is + 1);

            let d2 = d * sc2;

            let dm2 = dmin * m2;

            // First pass: 32 low nibbles (use sc1, m1)

            for &
byte2.09M
 in q {

                let q_val = (byte & 0x0F) as f32;

                let value = d1 * q_val - dm1;

                acc += value * activations[activation_idx];

                activation_idx += 1;

            }

            // Second pass: 32 high nibbles (use sc2, m2)

            for &
byte2.09M
 in q {

                let q_val = (byte >> 4) as f32;

                let value = d2 * q_val - dm2;

                acc += value * activations[activation_idx];

                activation_idx += 1;

            }

        }

    }

    Ok(acc)

}

/// Fused Q4_K dequantize + dot product with SIMD acceleration

///

/// This is the public, safe API that automatically dispatches to the best

/// available implementation (AVX2 when available, scalar fallback otherwise).

///

/// # Arguments

///

/// * `q4k_data` - Raw Q4_K quantized data (super-blocks of 144 bytes)

/// * `activations` - f32 activation values (must match dequantized length)

///

/// # Returns

///

/// The dot product as f32, matching `fused_q4k_dot` within 4 ULPs

///

/// # Errors

///

/// Returns error if:

/// - `q4k_data` length is not a multiple of 144 bytes (super-block size)

/// - `activations` length doesn't match the number of quantized values

///

/// # Performance

///

/// - AVX2: ~8x speedup over scalar via 256-bit SIMD + FMA

/// - Fused operation: 8x memory bandwidth reduction vs dequant-then-dot

/// - Combined potential: Up to 64x improvement for memory-bound operations

pub fn fused_q4k_dot_simd(q4k_data: &[u8], activations: &[f32]) -> Result<f32> {

    // Runtime feature detection with fallback (per RustBelt pattern)

    #[cfg(target_arch = "x86_64")]

    {

        // PAR-126: AVX-512 VNNI requires pre-quantized activations (Q4K×Q8K format)

        // For now, use AVX2 which works with f32 activations directly.

        // Future optimization: pre-quantize activations to Q8_0 format once per matmul.

        if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {

            // SAFETY: We've verified AVX2 and FMA are available at runtime

            // The unsafe function performs the same logical operation as scalar

            // SAFETY: Memory safety ensured by bounds checking and alignment

            return unsafe { fused_q4k_dot_avx2(q4k_data, activations) };

        }

    }

    // Fallback to scalar implementation

    fused_q4k_dot(q4k_data, activations)

}

/// AVX-512 VNNI accelerated Q4_K dot product (PAR-126)

///

/// Uses vpdpbusd for int8×int8 accumulation which is 4x faster than FP32 FMA.

/// Key insight from llama.cpp/llamafile: integer dot products dominate FP on modern CPUs.

///

/// # Algorithm

/// 1. Quantize f32 activations to int8 (on the fly, per super-block)

/// 2. Expand 4-bit weights to int8

/// 3. Use vpdpbusd for packed u8×i8 dot product

/// 4. Scale and accumulate at the end

///

/// # Safety

/// Requires AVX-512F, AVX-512BW, and AVX-512VNNI

#[cfg(target_arch = "x86_64")]

#[target_feature(enable = "avx512f", enable = "avx512bw", enable = "avx512vnni")]

#[allow(unsafe_op_in_unsafe_fn)]

unsafe fn fused_q4k_dot_avx512_vnni(q4k_data: &[u8], activations: &[f32]) -> Result<f32> {

    #[allow(clippy::wildcard_imports)]

    use std::arch::x86_64::*;

    const SUPER_BLOCK_BYTES: usize = 144;

    if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Q4_K data length {} is not a multiple of super-block size {}",

                q4k_data.len(),

                SUPER_BLOCK_BYTES

            ),

        });

    }

    let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;

    let expected_values = num_super_blocks * QK_K;

    if activations.len() != expected_values {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Activation length {} doesn't match Q4_K values count {}",

                activations.len(),

                expected_values

            ),

        });

    }

    // Final accumulator (FP32)

    let mut total_sum = 0.0f32;

    let mut activation_idx = 0;

    // Nibble mask for 4-bit extraction

    let nibble_mask = _mm512_set1_epi8(0x0F_i8);

    for sb_idx in 0..num_super_blocks {

        let sb_start = sb_idx * SUPER_BLOCK_BYTES;

        // Prefetch next super-block

        if sb_idx + 1 < num_super_blocks {

            let next_sb = (sb_idx + 1) * SUPER_BLOCK_BYTES;

            _mm_prefetch(q4k_data.as_ptr().add(next_sb).cast::<i8>(), _MM_HINT_T0);

        }

        // Read d and dmin (f16 → f32)

        let d = read_f16(&q4k_data[sb_start..sb_start + 2]);

        let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);

        // Read scales (12 bytes)

        let mut scales = [0u8; 12];

        scales.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);

        // Pointer to quantized data (128 bytes)

        let qs_ptr = q4k_data.as_ptr().add(sb_start + 16);

        // Process 64 values at a time (matches AVX-512 width of 64 bytes)

        for j in (0..QK_K).step_by(64) {

            let q_start = j / 2;

            // Get scales for the two 32-value halves

            let is = j / 32;

            let (sc1, m1) = extract_scale_min(&scales, is);

            let (sc2, m2) = extract_scale_min(&scales, is + 1);

            // Load 32 bytes of quantized data (64 nibbles)

            let q_bytes_256 = _mm256_loadu_si256(qs_ptr.add(q_start).cast::<__m256i>());

            // Expand to 512-bit for AVX-512 operations

            let q_bytes = _mm512_castsi256_si512(q_bytes_256);

            let q_bytes = _mm512_inserti64x4(q_bytes, q_bytes_256, 1);

            // Extract low and high nibbles

            let q_lo = _mm512_and_si512(q_bytes, nibble_mask);

            let q_hi = _mm512_and_si512(_mm512_srli_epi16(q_bytes, 4), nibble_mask);

            // Process with integer dot products

            // We need to quantize activations and use vpdpbusd

            // Load 32 f32 activations for low nibbles, convert to int8

            let act_slice = &activations[activation_idx..activation_idx + 32];

            // Find min/max for quantization

            let act_max = act_slice.iter().copied().fold(f32::NEG_INFINITY, f32::max);

            let act_min = act_slice.iter().copied().fold(f32::INFINITY, f32::min);

            let act_scale = if act_max > act_min {

                127.0 / (act_max - act_min)

            } else {

                1.0

            };

            // Quantize activations to int8

            let mut act_i8 = [0i8; 32];

            for (i, &a) in act_slice.iter().enumerate() {

                act_i8[i] = ((a - act_min) * act_scale).round() as i8;

            }

            // Load quantized activations

            let act_vec = _mm256_loadu_si256(act_i8.as_ptr().cast::<__m256i>());

            // Expand q_lo from u8 to match

            let q_lo_256 = _mm512_castsi512_si256(q_lo);

            // Use vpdpbusd: u8 × i8 → i32 accumulation

            let _zero = _mm512_setzero_si512();

            let q_lo_512 = _mm512_cvtepu8_epi16(q_lo_256);

            let act_512 = _mm512_cvtepi8_epi16(act_vec);

            let prod = _mm512_mullo_epi16(q_lo_512, act_512);

            // Horizontal sum

            let sum_256 = _mm256_add_epi16(

                _mm512_castsi512_si256(prod),

                _mm512_extracti64x4_epi64(prod, 1),

            );

            let sum_128 = _mm_add_epi16(

                _mm256_castsi256_si128(sum_256),

                _mm256_extracti128_si256(sum_256, 1),

            );

            // Expand to i32 and sum

            let sum_32 = _mm256_cvtepi16_epi32(sum_128);

            let sum_arr: [i32; 8] = std::mem::transmute(sum_32);

            let int_sum: i32 = sum_arr.iter().sum();

            // Scale back to float

            let float_sum = int_sum as f32 * d * sc1 / act_scale - (32.0 * dmin * m1);

            total_sum += float_sum;

            activation_idx += 32;

            // Process high nibbles similarly

            let act_slice2 = &activations[activation_idx..activation_idx + 32];

            let act_max2 = act_slice2.iter().copied().fold(f32::NEG_INFINITY, f32::max);

            let act_min2 = act_slice2.iter().copied().fold(f32::INFINITY, f32::min);

            let act_scale2 = if act_max2 > act_min2 {

                127.0 / (act_max2 - act_min2)

            } else {

                1.0

            };

            let mut act_i8_2 = [0i8; 32];

            for (i, &a) in act_slice2.iter().enumerate() {

                act_i8_2[i] = ((a - act_min2) * act_scale2).round() as i8;

            }

            let act_vec2 = _mm256_loadu_si256(act_i8_2.as_ptr().cast::<__m256i>());

            let q_hi_256 = _mm512_castsi512_si256(q_hi);

            let q_hi_512 = _mm512_cvtepu8_epi16(q_hi_256);

            let act_512_2 = _mm512_cvtepi8_epi16(act_vec2);

            let prod2 = _mm512_mullo_epi16(q_hi_512, act_512_2);

            let sum_256_2 = _mm256_add_epi16(

                _mm512_castsi512_si256(prod2),

                _mm512_extracti64x4_epi64(prod2, 1),

            );

            let sum_128_2 = _mm_add_epi16(

                _mm256_castsi256_si128(sum_256_2),

                _mm256_extracti128_si256(sum_256_2, 1),

            );

            let sum_32_2 = _mm256_cvtepi16_epi32(sum_128_2);

            let sum_arr2: [i32; 8] = std::mem::transmute(sum_32_2);

            let int_sum2: i32 = sum_arr2.iter().sum();

            let float_sum2 = int_sum2 as f32 * d * sc2 / act_scale2 - (32.0 * dmin * m2);

            total_sum += float_sum2;

            activation_idx += 32;

        }

    }

    Ok(total_sum)

}

/// AVX2-accelerated fused Q4_K dequant+dot kernel (PARITY-003: llama.cpp-style SIMD)

///

/// # Safety

///

/// Caller must ensure:

/// 1. AVX2 and FMA CPU features are available (use `is_x86_feature_detected!`)

/// 2. Input slices are valid (handled by Rust's slice guarantees)

///

/// This function is marked unsafe due to SIMD intrinsics, but is logically

/// equivalent to the scalar `fused_q4k_dot` (within ULP tolerance).

///

/// # Optimizations (PARITY-003)

/// - SIMD loads: `_mm256_loadu_si256` for 32-byte bulk loads (vs scalar byte loads)

/// - SIMD nibble extraction: `_mm256_and_si256` with 0x0F mask (vs scalar & 0x0F)

/// - 4 independent accumulators to hide FMA latency

/// - Software prefetching for next super-block

/// - Matches llama.cpp ggml_vec_dot_q4_K_q8_K pattern

#[cfg(target_arch = "x86_64")]

#[target_feature(enable = "avx2", enable = "fma")]

#[allow(unsafe_op_in_unsafe_fn)]

unsafe fn fused_q4k_dot_avx2(q4k_data: &[u8], activations: &[f32]) -> Result<f32> {

    // Allow wildcard import for SIMD intrinsics (standard pattern for arch-specific code)

    #[allow(clippy::wildcard_imports)]

    use std::arch::x86_64::*;

    const SUPER_BLOCK_BYTES: usize = 144;

    // Validate inputs (same as scalar)

    if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Q4_K data length {} is not a multiple of super-block size {}",

                q4k_data.len(),

                SUPER_BLOCK_BYTES

            ),

        });

    }

    let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;

    let expected_values = num_super_blocks * QK_K;

    if activations.len() != expected_values {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Activation length {} doesn't match Q4_K values count {}",

                activations.len(),

                expected_values

            ),

        });

    }

    // Nibble mask for extracting 4-bit values (llama.cpp pattern)

    let nibble_mask = _mm256_set1_epi8(0x0F_i8);

    // PARITY-003: 4 independent accumulators to hide FMA latency

    // FMA latency = 4 cycles, throughput = 2/cycle

    // With 4 independent chains, we saturate the FMA throughput

    let mut acc0 = _mm256_setzero_ps();

    let mut acc1 = _mm256_setzero_ps();

    let mut acc2 = _mm256_setzero_ps();

    let mut acc3 = _mm256_setzero_ps();

    let mut activation_idx = 0;

    for sb_idx in 0..
num_super_blocks1.18M
 {

        let sb_start = sb_idx * SUPER_BLOCK_BYTES;

        // Prefetch next super-block while processing current

        if sb_idx + 1 < num_super_blocks {

            let next_sb = (sb_idx + 1) * SUPER_BLOCK_BYTES;

            // SAFETY: Prefetch is a hint, pointer arithmetic is in bounds (checked above)

            _mm_prefetch(q4k_data.as_ptr().add(next_sb).cast::<i8>(), _MM_HINT_T0);

        }

        // Read d and dmin (f16 -> f32)

        let d = read_f16(&q4k_data[sb_start..sb_start + 2]);

        let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);

        // Read scales (12 bytes)

        let mut scales = [0u8; 12];

        scales.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);

        // Pointer to quantized data (128 bytes = 256 nibbles = 256 values)

        let qs_ptr = q4k_data.as_ptr().add(sb_start + 16);

        // PAR-001: Match dequantize_q4_k layout (llama.cpp/candle compatible)

        // Process 4 chunks of 64 values each (j=0, 64, 128, 192)

        // Each chunk: 32 low nibbles (sc1), then 32 high nibbles (sc2)

        for j in 
(0..QK_K)1.20M
.
step_by1.20M
(64) {

            let q_start = j / 2; // 32 bytes per 64-value chunk

            // Get scales for the two 32-value halves

            let is = j / 32;

            let (sc1, m1) = extract_scale_min(&scales, is);

            let (sc2, m2) = extract_scale_min(&scales, is + 1);

            // Precompute d*scale and dmin*min for both halves

            let d_scale1 = d * sc1;

            let dm1 = dmin * m1;

            let d_scale2 = d * sc2;

            let dm2 = dmin * m2;

            // SIMD OPTIMIZATION: Load 32 bytes at once (64 nibbles = 64 values)

            // llama.cpp pattern: _mm256_loadu_si256 + AND/SHIFT for nibble extraction

            let q_bytes = _mm256_loadu_si256(qs_ptr.add(q_start).cast::<__m256i>());

            // Extract low nibbles (first 32 values) and high nibbles (second 32 values)

            let q_lo = _mm256_and_si256(q_bytes, nibble_mask);

            let q_hi = _mm256_and_si256(_mm256_srli_epi16(q_bytes, 4), nibble_mask);

            // Process low nibbles (32 values with scale sc1/m1)

            // Split into 4 groups of 8 for f32 conversion (AVX2 can convert 8 i32→f32)

            let d_scale1_vec = _mm256_set1_ps(d_scale1);

            let dm1_vec = _mm256_set1_ps(dm1);

            // Extract bytes 0-7 (low nibbles) → convert to f32

            let q_lo_128_0 = _mm256_castsi256_si128(q_lo);

            let q_lo_i32_0 = _mm256_cvtepu8_epi32(q_lo_128_0);

            let q_lo_f32_0 = _mm256_cvtepi32_ps(q_lo_i32_0);

            let dequant0 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_0, dm1_vec);

            let act0 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));

            acc0 = _mm256_fmadd_ps(dequant0, act0, acc0);

            activation_idx += 8;

            // Extract bytes 8-15 (low nibbles)

            let q_lo_shifted = _mm_srli_si128(q_lo_128_0, 8);

            let q_lo_i32_1 = _mm256_cvtepu8_epi32(q_lo_shifted);

            let q_lo_f32_1 = _mm256_cvtepi32_ps(q_lo_i32_1);

            let dequant1 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_1, dm1_vec);

            let act1 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));

            acc1 = _mm256_fmadd_ps(dequant1, act1, acc1);

            activation_idx += 8;

            // Extract bytes 16-23 (low nibbles from high 128 bits)

            let q_lo_128_1 = _mm256_extracti128_si256(q_lo, 1);

            let q_lo_i32_2 = _mm256_cvtepu8_epi32(q_lo_128_1);

            let q_lo_f32_2 = _mm256_cvtepi32_ps(q_lo_i32_2);

            let dequant2 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_2, dm1_vec);

            let act2 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));

            acc2 = _mm256_fmadd_ps(dequant2, act2, acc2);

            activation_idx += 8;

            // Extract bytes 24-31 (low nibbles)

            let q_lo_shifted2 = _mm_srli_si128(q_lo_128_1, 8);

            let q_lo_i32_3 = _mm256_cvtepu8_epi32(q_lo_shifted2);

            let q_lo_f32_3 = _mm256_cvtepi32_ps(q_lo_i32_3);

            let dequant3 = _mm256_fmsub_ps(d_scale1_vec, q_lo_f32_3, dm1_vec);

            let act3 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));

            acc3 = _mm256_fmadd_ps(dequant3, act3, acc3);

            activation_idx += 8;

            // Process high nibbles (32 values with scale sc2/m2)

            let d_scale2_vec = _mm256_set1_ps(d_scale2);

            let dm2_vec = _mm256_set1_ps(dm2);

            // Extract bytes 0-7 (high nibbles)

            let q_hi_128_0 = _mm256_castsi256_si128(q_hi);

            let q_hi_i32_0 = _mm256_cvtepu8_epi32(q_hi_128_0);

            let q_hi_f32_0 = _mm256_cvtepi32_ps(q_hi_i32_0);

            let dequant4 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_0, dm2_vec);

            let act4 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));

            acc0 = _mm256_fmadd_ps(dequant4, act4, acc0);

            activation_idx += 8;

            // Extract bytes 8-15 (high nibbles)

            let q_hi_shifted = _mm_srli_si128(q_hi_128_0, 8);

            let q_hi_i32_1 = _mm256_cvtepu8_epi32(q_hi_shifted);

            let q_hi_f32_1 = _mm256_cvtepi32_ps(q_hi_i32_1);

            let dequant5 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_1, dm2_vec);

            let act5 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));

            acc1 = _mm256_fmadd_ps(dequant5, act5, acc1);

            activation_idx += 8;

            // Extract bytes 16-23 (high nibbles from high 128 bits)

            let q_hi_128_1 = _mm256_extracti128_si256(q_hi, 1);

            let q_hi_i32_2 = _mm256_cvtepu8_epi32(q_hi_128_1);

            let q_hi_f32_2 = _mm256_cvtepi32_ps(q_hi_i32_2);

            let dequant6 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_2, dm2_vec);

            let act6 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));

            acc2 = _mm256_fmadd_ps(dequant6, act6, acc2);

            activation_idx += 8;

            // Extract bytes 24-31 (high nibbles)

            let q_hi_shifted2 = _mm_srli_si128(q_hi_128_1, 8);

            let q_hi_i32_3 = _mm256_cvtepu8_epi32(q_hi_shifted2);

            let q_hi_f32_3 = _mm256_cvtepi32_ps(q_hi_i32_3);

            let dequant7 = _mm256_fmsub_ps(d_scale2_vec, q_hi_f32_3, dm2_vec);

            let act7 = _mm256_loadu_ps(activations.as_ptr().add(activation_idx));

            acc3 = _mm256_fmadd_ps(dequant7, act7, acc3);

            activation_idx += 8;

        }

    }

    // Combine 4 accumulators → single accumulator

    let acc_01 = _mm256_add_ps(acc0, acc1);

    let acc_23 = _mm256_add_ps(acc2, acc3);

    let acc = _mm256_add_ps(acc_01, acc_23);

    // Horizontal sum: reduce 8 lanes to single value

    let sum_halves = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps(acc, 1));

    let temp = _mm_add_ps(sum_halves, _mm_movehl_ps(sum_halves, sum_halves));

    let temp = _mm_add_ss(temp, _mm_shuffle_ps(temp, temp, 1));

    let result = _mm_cvtss_f32(temp);

    Ok(result)

}

// ============================================================================

// Q4_K × Q8_K KERNELS (Super-block aligned integer-only arithmetic)

// ============================================================================

/// Fused Q4_K × Q8_K dot product (super-block aligned, llama.cpp-style)

///

/// Uses Q8_K format (256 values per super-block, single scale) for maximum

/// SIMD efficiency. This matches llama.cpp's `ggml_vec_dot_q4_K_q8_K`.

///

/// # Arguments

/// * `q4k_data` - Raw Q4_K quantized data (144 bytes per super-block)

/// * `q8k_scales` - Q8_K scales (one per super-block)

/// * `q8k_quants` - Q8_K quantized int8 values (256 per super-block)

///

/// # Performance

///

/// Compared to Q4_K × f32:

/// - 8x fewer memory reads for activations

/// - Integer-only inner loop (no f32 conversion until end)

/// - Single scale multiplication per super-block (vs 8 for Q8_0)

pub fn fused_q4k_q8k_dot(q4k_data: &[u8], q8k_scales: &[f32], q8k_quants: &[i8]) -> Result<f32> {

    const SUPER_BLOCK_BYTES: usize = 144;

    if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Q4_K data length {} is not a multiple of {}",

                q4k_data.len(),

                SUPER_BLOCK_BYTES

            ),

        });

    }

    let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;

    let expected_values = num_super_blocks * QK_K;

    if q8k_scales.len() < num_super_blocks {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Q8_K scales count {} < expected {}",

                q8k_scales.len(),

                num_super_blocks

            ),

        });

    }

    if q8k_quants.len() < expected_values {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Q8_K quants count {} < expected {}",

                q8k_quants.len(),

                expected_values

            ),

        });

    }

    let mut total_acc = 0.0f32;

    for sb_idx in 0..
num_super_blocks26
 {

        let sb_start = sb_idx * SUPER_BLOCK_BYTES;

        let q8_start = sb_idx * QK_K;

        // Read Q4_K super-block header

        let d = read_f16(&q4k_data[sb_start..sb_start + 2]);

        let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);

        // Read scales (12 bytes for 8 blocks)

        let mut scales = [0u8; 12];

        scales.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);

        // Q8_K scale for this super-block

        let q8_scale = q8k_scales[sb_idx];

        // Process 4 chunks of 64 values (matching dequantize_q4_k layout)

        // The dequantized output order is: 32 low nibbles, then 32 high nibbles

        // So activations[j..j+32] correspond to low nibbles, activations[j+32..j+64] to high

        for j in 
(0..QK_K)45
.
step_by45
(64) {

            let q_offset = sb_start + 16 + j / 2; // 32 bytes per 64-value chunk

            let q8_offset = q8_start + j;

            // Get scales for low and high nibbles

            let is = j / 32;

            let (sc1, m1) = extract_scale_min(&scales, is);

            let (sc2, m2) = extract_scale_min(&scales, is + 1);

            // Combined scale factors

            let d_sc1_q8 = d * sc1 * q8_scale;

            let dm1_q8 = dmin * m1 * q8_scale;

            let d_sc2_q8 = d * sc2 * q8_scale;

            let dm2_q8 = dmin * m2 * q8_scale;

            // Accumulators for low and high nibbles

            let mut sum_lo: i32 = 0; // q4_lo × q8 (for activations[j..j+32])

            let mut sum_hi: i32 = 0; // q4_hi × q8 (for activations[j+32..j+64])

            let mut q8_sum_lo: i32 = 0;

            let mut q8_sum_hi: i32 = 0;

            for 
b5.76k
 in 0..32 {

                let q4_byte = q4k_data[q_offset + b];

                // Low nibble × activation[j + b] (first 32 positions in dequant order)

                let q4_lo = (q4_byte & 0x0F) as i32;

                let q8_lo = q8k_quants[q8_offset + b] as i32;

                sum_lo += q4_lo * q8_lo;

                q8_sum_lo += q8_lo;

                // High nibble × activation[j + 32 + b] (second 32 positions in dequant order)

                let q4_hi = ((q4_byte >> 4) & 0x0F) as i32;

                let q8_hi = q8k_quants[q8_offset + 32 + b] as i32;

                sum_hi += q4_hi * q8_hi;

                q8_sum_hi += q8_hi;

            }

            // Apply formula: (d * scale * sum_q4_q8 - dmin * min * sum_q8) * q8_scale

            total_acc += d_sc1_q8 * (sum_lo as f32) - dm1_q8 * (q8_sum_lo as f32);

            total_acc += d_sc2_q8 * (sum_hi as f32) - dm2_q8 * (q8_sum_hi as f32);

        }

    }

    Ok(total_acc)

}

/// SIMD-accelerated Q4_K × Q8_K dot product

///

/// Uses AVX-512 VNNI (vpdpbusd) for maximum throughput, falls back to AVX2 or scalar.

/// Single scale per super-block eliminates per-block overhead.

pub fn fused_q4k_q8k_dot_simd(

    q4k_data: &[u8],

    q8k_scales: &[f32],

    q8k_quants: &[i8],

) -> Result<f32> {

    #[cfg(target_arch = "x86_64")]

    {

        // PAR-126: Use V2 optimized AVX-512 VNNI kernel (deferred horizontal sums)

        if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512vnni") {

            // SAFETY: Memory safety ensured by bounds checking and alignment

            return unsafe { fused_q4k_q8k_dot_avx512vnni_v2(q4k_data, q8k_scales, q8k_quants) };

        }

        // Fallback to AVX2 (layout issue resolved)

        if is_x86_feature_detected!("avx2") {

            // SAFETY: Memory safety ensured by bounds checking and alignment

            return unsafe { fused_q4k_q8k_dot_avx2(q4k_data, q8k_scales, q8k_quants) };

        }

    }

    fused_q4k_q8k_dot(q4k_data, q8k_scales, q8k_quants)

}

/// PAR-126 V2: AVX-512 VNNI kernel with deferred horizontal sums

///

/// Key optimization: Instead of extracting block sums to scalars in the inner loop,

/// keeps all 8 block accumulators in __m256i vectors and only reduces at the end.

/// This eliminates 24 horizontal sums per row (from 8 to 1).

#[cfg(target_arch = "x86_64")]

#[target_feature(enable = "avx512f", enable = "avx512vnni", enable = "avx512bw")]

#[allow(unsafe_op_in_unsafe_fn)]

#[allow(clippy::similar_names)]

#[allow(clippy::too_many_lines)]

unsafe fn fused_q4k_q8k_dot_avx512vnni_v2(

    q4k_data: &[u8],

    q8k_scales: &[f32],

    q8k_quants: &[i8],

) -> Result<f32> {

    #[allow(clippy::wildcard_imports)]

    use std::arch::x86_64::*;

    const SUPER_BLOCK_BYTES: usize = 144;

    if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Q4_K data length {} is not a multiple of {}",

                q4k_data.len(),

                SUPER_BLOCK_BYTES

            ),

        });

    }

    let num_super_blocks = q4k_data.len() / SUPER_BLOCK_BYTES;

    let expected_values = num_super_blocks * QK_K;

    if q8k_scales.len() < num_super_blocks || q8k_quants.len() < expected_values {

        return Err(RealizarError::InvalidShape {

            reason: "Q8_K buffer too small".to_string(),

        });

    }

    let nibble_mask = _mm256_set1_epi8(0x0F_i8);

    let ones_16 = _mm256_set1_epi16(1);

    // Global float accumulator

    let mut total_acc = _mm256_setzero_ps();

    for sb_idx in 0..
num_super_blocks532
 {

        let sb_start = sb_idx * SUPER_BLOCK_BYTES;

        let q8_start = sb_idx * QK_K;

        // Prefetch next super-block

        if sb_idx + 1 < num_super_blocks {

            _mm_prefetch(

                q4k_data

                    .as_ptr()

                    .add((sb_idx + 1) * SUPER_BLOCK_BYTES)

                    .cast::<i8>(),

                _MM_HINT_T0,

            );

            _mm_prefetch(

                q8k_quants.as_ptr().add((sb_idx + 1) * QK_K).cast::<i8>(),

                _MM_HINT_T0,

            );

        }

        // Read Q4_K header

        let d = read_f16(&q4k_data[sb_start..sb_start + 2]);

        let dmin = read_f16(&q4k_data[sb_start + 2..sb_start + 4]);

        let mut scales_raw = [0u8; 12];

        scales_raw.copy_from_slice(&q4k_data[sb_start + 4..sb_start + 16]);

        let q8_scale = q8k_scales[sb_idx];

        let d_q8 = d * q8_scale;

        let dmin_q8 = dmin * q8_scale;

        let qs_ptr = q4k_data.as_ptr().add(sb_start + 16);

        let q8_ptr = q8k_quants.as_ptr().add(q8_start);

        // PAR-126 V2: Keep all 8 block results in vectors

        // block_dots_vec[i] = sum of (Q4[block_i] * Q8[block_i])

        // block_q8sums_vec[i] = sum of Q8[block_i]

        let mut block_dots_vec = _mm256_setzero_si256();

        let mut block_q8sums_vec = _mm256_setzero_si256();

        // Process 64 values (2 blocks) per iteration, 4 iterations = 8 blocks

        // Each iteration produces 2 block sums that go into specific lanes

        for 
chunk2.21k
 in 0..4 {

            let j = chunk * 64;

            let q_offset = j / 2;

            // Load 32 bytes Q4 (64 nibbles packed)

            let q4_bytes = _mm256_loadu_si256(qs_ptr.add(q_offset).cast::<__m256i>());

            // Extract nibbles: lo = first 32 values, hi = next 32 values

            let q4_lo = _mm256_and_si256(q4_bytes, nibble_mask);

            let q4_hi = _mm256_and_si256(_mm256_srli_epi16(q4_bytes, 4), nibble_mask);

            // Load 64 bytes Q8

            let q8_lo = _mm256_loadu_si256(q8_ptr.add(j).cast::<__m256i>());

            let q8_hi = _mm256_loadu_si256(q8_ptr.add(j + 32).cast::<__m256i>());

            // Q4 × Q8 products -> i16 via maddubs -> i32 via madd

            let prod_lo_i16 = _mm256_maddubs_epi16(q4_lo, q8_lo);

            let prod_hi_i16 = _mm256_maddubs_epi16(q4_hi, q8_hi);

            let prod_lo_i32 = _mm256_madd_epi16(prod_lo_i16, ones_16);

            let prod_hi_i32 = _mm256_madd_epi16(prod_hi_i16, ones_16);

            // Reduce each 256-bit register to ONE sum using hadd (within-lane reduction)

            // prod_lo_i32 has 8 i32 values, we need their sum for block chunk*2

            // prod_hi_i32 has 8 i32 values, we need their sum for block chunk*2+1

            // Use hadd twice to get 8 -> 4 -> 2 values per register

            let prod_lo_h1 = _mm256_hadd_epi32(prod_lo_i32, prod_hi_i32); // interleaves lo and hi

            let prod_h2 = _mm256_hadd_epi32(prod_lo_h1, prod_lo_h1); // further reduce

            // Now prod_h2 has: [sum_lo_lane0, sum_hi_lane0, sum_lo_lane0, sum_hi_lane0,

            //                   sum_lo_lane1, sum_hi_lane1, sum_lo_lane1, sum_hi_lane1]

            // We need to extract the two unique sums and place them in block_dots_vec

            // Extract two sums: add lane0 and lane1 of each block

            let lane0 = _mm256_castsi256_si128(prod_h2);

            let lane1 = _mm256_extracti128_si256(prod_h2, 1);

            let sums_128 = _mm_add_epi32(lane0, lane1); // [dot_lo, dot_hi, dot_lo, dot_hi]

            // Insert into correct position based on chunk

            // block_dots_vec lanes: [0,1,2,3,4,5,6,7] for blocks [0,1,2,3,4,5,6,7]

            match chunk {

                0 => {

                    // Put in lanes 0,1

                    block_dots_vec =

                        _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 0), 0);

                    block_dots_vec =

                        _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 1), 1);

                },

                1 => {

                    block_dots_vec =

                        _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 0), 2);

                    block_dots_vec =

                        _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 1), 3);

                },

                2 => {

                    block_dots_vec =

                        _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 0), 4);

                    block_dots_vec =

                        _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 1), 5);

                },

                3 => {

                    block_dots_vec =

                        _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 0), 6);

                    block_dots_vec =

                        _mm256_insert_epi32(block_dots_vec, _mm_extract_epi32(sums_128, 1), 7);

                },

                _ => unreachable!(),

            }

            // Q8 sums for dmin correction using sad_epu8 (sum of absolute differences)

            // Convert signed i8 to unsigned by adding 128, compute SAD, then adjust

            let bias = _mm256_set1_epi8(-128_i8);

            let _q8_lo_u = _mm256_sub_epi8(q8_lo, bias); // signed [-128,127] -> unsigned [0,255]

            let _q8_hi_u = _mm256_sub_epi8(q8_hi, bias);

            // Use mpsadbw or manual approach for sum

            // Simpler: sign-extend and sum using madd

            let q8_lo_i16_a = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(q8_lo));

            let q8_lo_i16_b = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q8_lo, 1));

            let q8_hi_i16_a = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(q8_hi));

            let q8_hi_i16_b = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(q8_hi, 1));

            let q8_lo_i32_a = _mm256_madd_epi16(q8_lo_i16_a, ones_16);

            let q8_lo_i32_b = _mm256_madd_epi16(q8_lo_i16_b, ones_16);

            let q8_hi_i32_a = _mm256_madd_epi16(q8_hi_i16_a, ones_16);

            let q8_hi_i32_b = _mm256_madd_epi16(q8_hi_i16_b, ones_16);

            let q8_lo_sum = _mm256_add_epi32(q8_lo_i32_a, q8_lo_i32_b);

            let q8_hi_sum = _mm256_add_epi32(q8_hi_i32_a, q8_hi_i32_b);

            // Reduce to two sums using hadd

            let q8_h1 = _mm256_hadd_epi32(q8_lo_sum, q8_hi_sum);

            let q8_h2 = _mm256_hadd_epi32(q8_h1, q8_h1);

            let q8_lane0 = _mm256_castsi256_si128(q8_h2);

            let q8_lane1 = _mm256_extracti128_si256(q8_h2, 1);

            let q8_sums_128 = _mm_add_epi32(q8_lane0, q8_lane1);

            match chunk {

                0 => {

                    block_q8sums_vec =

                        _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 0), 0);

                    block_q8sums_vec =

                        _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 1), 1);

                },

                1 => {

                    block_q8sums_vec =

                        _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 0), 2);

                    block_q8sums_vec =

                        _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 1), 3);

                },

                2 => {

                    block_q8sums_vec =

                        _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 0), 4);

                    block_q8sums_vec =

                        _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 1), 5);

                },

                3 => {

                    block_q8sums_vec =

                        _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 0), 6);

                    block_q8sums_vec =

                        _mm256_insert_epi32(block_q8sums_vec, _mm_extract_epi32(q8_sums_128, 1), 7);

                },

                _ => unreachable!(),

            }

        }

        // Extract all 8 scales and mins into vectors

        let mut scales = [0.0f32; 8];

        let mut mins = [0.0f32; 8];

        for 
i4.42k
 in 0..8 {

            let (sc, m) = extract_scale_min(&scales_raw, i);

            scales[i] = sc;

            mins[i] = m;

        }

        let scales_vec = _mm256_loadu_ps(scales.as_ptr());

        let mins_vec = _mm256_loadu_ps(mins.as_ptr());

        // Convert block results to f32

        let dots_f32 = _mm256_cvtepi32_ps(block_dots_vec);

        let q8sums_f32 = _mm256_cvtepi32_ps(block_q8sums_vec);

        // Compute: d_q8 * scales * dots - dmin_q8 * mins * q8sums

        let d_q8_vec = _mm256_set1_ps(d_q8);

        let dmin_q8_vec = _mm256_set1_ps(dmin_q8);

        let term1 = _mm256_mul_ps(d_q8_vec, _mm256_mul_ps(scales_vec, dots_f32));

        let term2 = _mm256_mul_ps(dmin_q8_vec, _mm256_mul_ps(mins_vec, q8sums_f32));

        let result = _mm256_sub_ps(term1, term2);

        // Accumulate into total (defer final horizontal sum)

        total_acc = _mm256_add_ps(total_acc, result);

    }

    // ONE horizontal sum at the very end

    let sum128 = _mm_add_ps(

        _mm256_castps256_ps128(total_acc),

        _mm256_extractf128_ps(total_acc, 1),

    );

    let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));

    let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));

    Ok(_mm_cvtss_f32(sum32))

}

/// AVX-512 VNNI optimized Q4_K × Q8_K dot product (llama.cpp style)

///

/// Uses `vpdpbusd` instruction for 16-way uint8×int8→int32 multiply-accumulate.

/// This is the fastest CPU path for quantized matmul.

///

/// # Safety

/// Requires AVX-512F and AVX-512 VNNI CPU features.

#[cfg(target_arch = "x86_64")]

#[target_feature(enable = "avx512f", enable = "avx512vnni", enable = "avx512bw")]

#[allow(unsafe_op_in_unsafe_fn)]

unsafe fn fused_q4k_q8k_dot_avx512vnni(

    q4k_data: &[u8],

    q8k_scales: &[f32],

    q8k_quants: &[i8],

) -> Result<f32> {

    #[allow(clippy::wildcard_imports)]

    use std::arch::x86_64::*;

    const SUPER_BLOCK_BYTES: usize = 144;

    if !q4k_data.len().is_multiple_of(SUPER_BLOCK_BYTES) {

        return Err(RealizarError::InvalidShape {

            reason: format!(

                "Q4_K data length {} is not a multiple of {}",

                q4k_data.len(),

                SUPER_BLOCK_BYTES

            ),