// Fused Q4_K Matmul Kernel - W4A16
// Directly computes C = A * dequant(B_q4k) without separate dequant pass
//
// For M=1 decode (GEMV): C[N] = A[K] * B_q4k^T[N,K]
// B_q4k is stored in Q4_K format: [N * ceil(K/256) * 144 bytes]
//
// Key optimizations:
// 1. Fused dequant + matmul - eliminates memory round-trip (2-3x speedup)
// 2. Subgroup operations for reduction
// 3. On-the-fly dequantization in registers
//
// A is f32 (activations), B_q4k is Q4_K quantized weights, C is f32.

enable subgroups;

// Q4_K constants
const QK_K: u32 = 256u;           // Elements per super-block
const SUBBLOCK_SIZE: u32 = 32u;   // Elements per sub-block
const NUM_SUBBLOCKS: u32 = QK_K / SUBBLOCK_SIZE;  // 8 sub-blocks per super-block

override WORKGROUP_SIZE: u32 = 256u;
const MAX_WORKGROUP_SIZE: u32 = 256u;

struct Uniforms {
    M: u32,                   // Always 1 for GEMV
    N: u32,                   // Output dimension
    K: u32,                   // Inner dimension (may be non-256-aligned)
    alpha: f32,
    num_blocks_per_row: u32,  // ceil(K / 256)
    _pad0: u32,               // 16-byte alignment padding
    _pad1: u32,
    _pad2: u32,
}

// Q4_K block structure (144 bytes)
// Layout: d(2) + dmin(2) + scales(12) + qs(128)
struct Q4KBlock {
    d_dmin: u32,          // d (f16) and dmin (f16) packed
    scales: array<u32, 3>, // 12 bytes of packed 6-bit scales
    qs: array<u32, 32>,   // 128 bytes of 4-bit quantized values
}

@group(0) @binding(0) var<uniform> u: Uniforms;
@group(0) @binding(1) var<storage, read> A: array<f32>;
@group(0) @binding(2) var<storage, read> B_q4k: array<Q4KBlock>;
@group(0) @binding(3) var<storage, read_write> C: array<f32>;

// Extract f16 from packed u32
fn unpack_f16_lo(packed: u32) -> f32 {
    return unpack2x16float(packed).x;
}

fn unpack_f16_hi(packed: u32) -> f32 {
    return unpack2x16float(packed).y;
}

// Get byte from scales array
fn get_scale_byte(scales: array<u32, 3>, byte_idx: u32) -> u32 {
    let word_idx = byte_idx / 4u;
    let byte_in_word = byte_idx % 4u;
    return (scales[word_idx] >> (byte_in_word * 8u)) & 0xFFu;
}

// llama.cpp Q4_K scale/min extraction
fn get_scale_min_k4(scales: array<u32, 3>, j: u32) -> vec2<u32> {
    var sc: u32;
    var mn: u32;

    if (j < 4u) {
        sc = get_scale_byte(scales, j) & 63u;
        mn = get_scale_byte(scales, j + 4u) & 63u;
    } else {
        let q_j = get_scale_byte(scales, j + 4u);
        let q_lo = get_scale_byte(scales, j - 4u);
        let q_hi = get_scale_byte(scales, j);
        sc = (q_j & 0xFu) | ((q_lo >> 6u) << 4u);
        mn = (q_j >> 4u) | ((q_hi >> 6u) << 4u);
    }
    return vec2<u32>(sc, mn);
}

// Extract 4-bit quantized value from qs array
fn get_q4(qs: array<u32, 32>, idx: u32) -> u32 {
    let chunk = idx / 64u;
    let pos_in_chunk = idx % 64u;
    let use_upper = pos_in_chunk >= 32u;
    let byte_in_range = select(pos_in_chunk, pos_in_chunk - 32u, use_upper);
    let byte_idx = chunk * 32u + byte_in_range;

    let word_idx = byte_idx / 4u;
    let byte_in_word = byte_idx % 4u;
    let byte_val = (qs[word_idx] >> (byte_in_word * 8u)) & 0xFFu;

    if (use_upper) {
        return (byte_val >> 4u) & 0xFu;
    } else {
        return byte_val & 0xFu;
    }
}


// ============================================================================
// Multi-column GEMV for large vocab (LM head)
// ============================================================================
// For large N (e.g., vocab=262144), single-column-per-workgroup is inefficient:
// - 262K workgroups with only ~5 blocks per thread
// - GPU launch overhead dominates
//
// This variant processes COLS_PER_WG columns per workgroup:
// - 262144/32 = 8192 workgroups (32x fewer)
// - Each thread processes multiple columns, amortizing A loads
// - Much better GPU occupancy and throughput
//
// Workgroup layout: 256 threads = 8 threads per column × 32 columns
override COLS_PER_WG: u32 = 32u;
override THREADS_PER_COL_GEMV: u32 = 8u;  // 256 / 32 = 8
override USE_FULL_BLOCK_FAST_PATH: bool = false;

// Shared memory for reduction (one slot per thread)
var<workgroup> multicol_sums: array<f32, MAX_WORKGROUP_SIZE>;

@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
fn main_multicol(
    @builtin(local_invocation_id) lid: vec3<u32>,
    @builtin(workgroup_id) wg_id: vec3<u32>,
    @builtin(subgroup_invocation_id) sg_id: u32,
    @builtin(subgroup_size) sg_size: u32
) {
    let local_id = lid.x;

    // Which column within this workgroup (0..31)
    let col_in_wg = local_id / THREADS_PER_COL_GEMV;
    // Which thread within the column (0..7)
    let tid_in_col = local_id % THREADS_PER_COL_GEMV;

    // Global column index
    let col = wg_id.x * COLS_PER_WG + col_in_wg;

    // Track validity
    let is_valid = col < u.N;

    var partial_sum: f32 = 0.0;

    if (is_valid) {
        let num_blocks = u.num_blocks_per_row;
        let tail_size = select(u.K & 255u, 0u, USE_FULL_BLOCK_FAST_PATH);
        let full_blocks = select(
            num_blocks - select(0u, 1u, tail_size > 0u),
            num_blocks,
            USE_FULL_BLOCK_FAST_PATH
        );

        // B_q4k layout: row-major [N, K/256] - block b for column col is at col * num_blocks + b
        // Each of the 8 threads processes every 8th block
        for (var b: u32 = tid_in_col; b < full_blocks; b = b + THREADS_PER_COL_GEMV) {
            let block = B_q4k[col * num_blocks + b];
            let d = unpack_f16_lo(block.d_dmin);
            let dmin = unpack_f16_hi(block.d_dmin);
            let k_base = b * QK_K;

            // Process all 8 sub-blocks
            for (var sb: u32 = 0u; sb < NUM_SUBBLOCKS; sb = sb + 1u) {
                let sm = get_scale_min_k4(block.scales, sb);
                let scale = d * f32(sm.x);
                let min_val = dmin * f32(sm.y);
                let sb_base = sb * SUBBLOCK_SIZE;

                // Unroll by 4 for better ILP
                for (var i: u32 = 0u; i < SUBBLOCK_SIZE; i = i + 4u) {
                    let k0 = k_base + sb_base + i;
                    let k1 = k0 + 1u;
                    let k2 = k0 + 2u;
                    let k3 = k0 + 3u;

                    let a0 = A[k0];
                    let a1 = A[k1];
                    let a2 = A[k2];
                    let a3 = A[k3];

                    let q0 = get_q4(block.qs, sb_base + i);
                    let q1 = get_q4(block.qs, sb_base + i + 1u);
                    let q2 = get_q4(block.qs, sb_base + i + 2u);
                    let q3 = get_q4(block.qs, sb_base + i + 3u);

                    let w0 = scale * f32(q0) - min_val;
                    let w1 = scale * f32(q1) - min_val;
                    let w2 = scale * f32(q2) - min_val;
                    let w3 = scale * f32(q3) - min_val;

                    partial_sum = partial_sum + a0 * w0 + a1 * w1 + a2 * w2 + a3 * w3;
                }
            }
        }

        if (!USE_FULL_BLOCK_FAST_PATH && tail_size > 0u) {
            let tail_block = full_blocks;
            if (tail_block % THREADS_PER_COL_GEMV == tid_in_col) {
                let block = B_q4k[col * num_blocks + tail_block];
                let d = unpack_f16_lo(block.d_dmin);
                let dmin = unpack_f16_hi(block.d_dmin);
                let k_base = tail_block * QK_K;

                for (var sb: u32 = 0u; sb < NUM_SUBBLOCKS; sb = sb + 1u) {
                    let sb_base = sb * SUBBLOCK_SIZE;
                    if (sb_base >= tail_size) {
                        break;
                    }
                    let sm = get_scale_min_k4(block.scales, sb);
                    let scale = d * f32(sm.x);
                    let min_val = dmin * f32(sm.y);

                    for (var i: u32 = 0u; i < SUBBLOCK_SIZE; i = i + 4u) {
                        let k0 = k_base + sb_base + i;
                        let k1 = k0 + 1u;
                        let k2 = k0 + 2u;
                        let k3 = k0 + 3u;

                        var a0: f32 = 0.0;
                        var a1: f32 = 0.0;
                        var a2: f32 = 0.0;
                        var a3: f32 = 0.0;
                        if (k0 < u.K) { a0 = A[k0]; }
                        if (k1 < u.K) { a1 = A[k1]; }
                        if (k2 < u.K) { a2 = A[k2]; }
                        if (k3 < u.K) { a3 = A[k3]; }

                        let q0 = get_q4(block.qs, sb_base + i);
                        let q1 = get_q4(block.qs, sb_base + i + 1u);
                        let q2 = get_q4(block.qs, sb_base + i + 2u);
                        let q3 = get_q4(block.qs, sb_base + i + 3u);

                        let w0 = scale * f32(q0) - min_val;
                        let w1 = scale * f32(q1) - min_val;
                        let w2 = scale * f32(q2) - min_val;
                        let w3 = scale * f32(q3) - min_val;

                        partial_sum = partial_sum + a0 * w0 + a1 * w1 + a2 * w2 + a3 * w3;
                    }
                }
            }
        }
    }

    // Reduction within each column's 8 threads
    // Use shared memory since threads for one column may span multiple subgroups
    multicol_sums[local_id] = partial_sum;
    workgroupBarrier();

    // Thread 0 of each column reduces its 8 values
    if (tid_in_col == 0u && is_valid) {
        var final_sum: f32 = 0.0;
        let base = col_in_wg * THREADS_PER_COL_GEMV;
        for (var i: u32 = 0u; i < THREADS_PER_COL_GEMV; i = i + 1u) {
            final_sum = final_sum + multicol_sums[base + i];
        }
        C[col] = final_sum * u.alpha;
    }
}


override SHARED_A_MAX: u32 = 3584u;


// ============================================================================
// Optimised GEMV: shared-A cooperative load + fast nibble extraction.
// Uses shared A for the hot prefix and falls back to global A reads when
// K exceeds SHARED_A_MAX, which keeps large FFN down-projections correct.
// ============================================================================
var<workgroup> gemv_shared_A: array<f32, SHARED_A_MAX>;
var<workgroup> gemv_sums: array<f32, MAX_WORKGROUP_SIZE>;

fn gemv_load_a(k: u32) -> f32 {
    if (k < SHARED_A_MAX) {
        return gemv_shared_A[k];
    }
    return A[k];
}

@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
fn main_gemv(
    @builtin(local_invocation_id) lid: vec3<u32>,
    @builtin(workgroup_id) wg_id: vec3<u32>,
    @builtin(subgroup_invocation_id) sg_id: u32,
    @builtin(subgroup_size) sg_size: u32
) {
    let local_id = lid.x;

    // Cooperative load: all 256 threads fill shared A from global memory once.
    for (var idx: u32 = local_id; idx < min(u.K, SHARED_A_MAX); idx = idx + WORKGROUP_SIZE) {
        gemv_shared_A[idx] = A[idx];
    }
    workgroupBarrier();

    let col_in_wg = local_id / THREADS_PER_COL_GEMV;
    let tid_in_col = local_id % THREADS_PER_COL_GEMV;
    let col = wg_id.x * COLS_PER_WG + col_in_wg;
    let is_valid = col < u.N;

    var partial_sum: f32 = 0.0;

    if (is_valid) {
        let num_blocks = u.num_blocks_per_row;
        let tail_size = select(u.K & 255u, 0u, USE_FULL_BLOCK_FAST_PATH);
        let full_blocks = select(
            num_blocks - select(0u, 1u, tail_size > 0u),
            num_blocks,
            USE_FULL_BLOCK_FAST_PATH
        );

        for (var b: u32 = tid_in_col; b < full_blocks; b = b + THREADS_PER_COL_GEMV) {
            let block = B_q4k[col * num_blocks + b];
            let d = unpack_f16_lo(block.d_dmin);
            let dmin = unpack_f16_hi(block.d_dmin);
            let k_base = b * QK_K;

            for (var sb: u32 = 0u; sb < NUM_SUBBLOCKS; sb = sb + 1u) {
                let sm = get_scale_min_k4(block.scales, sb);
                let scale = d * f32(sm.x);
                let min_val = dmin * f32(sm.y);
                let sb_base = sb * SUBBLOCK_SIZE;

                // Fast nibble extraction: even sub-blocks use lower nibble,
                // odd sub-blocks use upper nibble.
                let chunk = sb >> 1u;
                let nibble_shift = (sb & 1u) * 4u;
                let word_base = chunk * 8u;

                for (var i: u32 = 0u; i < SUBBLOCK_SIZE; i = i + 4u) {
                    let k0 = k_base + sb_base + i;

                    let a0 = gemv_load_a(k0);
                    let a1 = gemv_load_a(k0 + 1u);
                    let a2 = gemv_load_a(k0 + 2u);
                    let a3 = gemv_load_a(k0 + 3u);

                    let word = block.qs[word_base + (i >> 2u)];
                    let q0 = (word >> nibble_shift) & 0xFu;
                    let q1 = (word >> (nibble_shift + 8u)) & 0xFu;
                    let q2 = (word >> (nibble_shift + 16u)) & 0xFu;
                    let q3 = (word >> (nibble_shift + 24u)) & 0xFu;

                    let w0 = scale * f32(q0) - min_val;
                    let w1 = scale * f32(q1) - min_val;
                    let w2 = scale * f32(q2) - min_val;
                    let w3 = scale * f32(q3) - min_val;

                    partial_sum = partial_sum + a0 * w0 + a1 * w1 + a2 * w2 + a3 * w3;
                }
            }
        }

        if (!USE_FULL_BLOCK_FAST_PATH && tail_size > 0u) {
            let tail_block = full_blocks;
            if (tail_block % THREADS_PER_COL_GEMV == tid_in_col) {
                let block = B_q4k[col * num_blocks + tail_block];
                let d = unpack_f16_lo(block.d_dmin);
                let dmin = unpack_f16_hi(block.d_dmin);
                let k_base = tail_block * QK_K;

                for (var sb: u32 = 0u; sb < NUM_SUBBLOCKS; sb = sb + 1u) {
                    let sb_base = sb * SUBBLOCK_SIZE;
                    if (sb_base >= tail_size) {
                        break;
                    }
                    let sm = get_scale_min_k4(block.scales, sb);
                    let scale = d * f32(sm.x);
                    let min_val = dmin * f32(sm.y);

                    let chunk = sb >> 1u;
                    let nibble_shift = (sb & 1u) * 4u;
                    let word_base = chunk * 8u;

                    for (var i: u32 = 0u; i < SUBBLOCK_SIZE; i = i + 4u) {
                        let k0 = k_base + sb_base + i;

                        var a0: f32 = 0.0;
                        var a1: f32 = 0.0;
                        var a2: f32 = 0.0;
                        var a3: f32 = 0.0;
                        if (k0 < u.K) { a0 = gemv_load_a(k0); }
                        if (k0 + 1u < u.K) { a1 = gemv_load_a(k0 + 1u); }
                        if (k0 + 2u < u.K) { a2 = gemv_load_a(k0 + 2u); }
                        if (k0 + 3u < u.K) { a3 = gemv_load_a(k0 + 3u); }

                        let word = block.qs[word_base + (i >> 2u)];
                        let q0 = (word >> nibble_shift) & 0xFu;
                        let q1 = (word >> (nibble_shift + 8u)) & 0xFu;
                        let q2 = (word >> (nibble_shift + 16u)) & 0xFu;
                        let q3 = (word >> (nibble_shift + 24u)) & 0xFu;

                        let w0 = scale * f32(q0) - min_val;
                        let w1 = scale * f32(q1) - min_val;
                        let w2 = scale * f32(q2) - min_val;
                        let w3 = scale * f32(q3) - min_val;

                        partial_sum = partial_sum + a0 * w0 + a1 * w1 + a2 * w2 + a3 * w3;
                    }
                }
            }
        }
    }

    gemv_sums[local_id] = partial_sum;
    workgroupBarrier();

    if (tid_in_col == 0u && is_valid) {
        var final_sum: f32 = 0.0;
        let base = col_in_wg * THREADS_PER_COL_GEMV;
        for (var i: u32 = 0u; i < THREADS_PER_COL_GEMV; i = i + 1u) {
            final_sum = final_sum + gemv_sums[base + i];
        }
        C[col] = final_sum * u.alpha;
    }
}