// File: kernel.cpp
// Reference-scalar gemm Q4_K (weight-only, genuine low-bit dot product).
// LLM target: replace this file with an optimised inner_gemm (SVE2/NEON
// nibble-unpack + int8 dot-product intrinsics are the intended optimization
// surface here).
//
// Unlike Definition.reference (which computes a clean, implementation-
// agnostic dequant-then-fp32-multiply ground truth), this kernel mirrors
// real GGML's ggml_vec_dot_q4_K_q8_K: the bf16 activation is dynamically
// quantized into 256-element blocks (one scale, int8 values -- same idea as
// ggml's Q8_K), then for each 32-element sub-block an integer dot product
// against the unpacked Q4_K nibbles is computed and scale-corrected before
// accumulating into a running fp32 total.
//
// All per-definition constants live in gemm_def:: (gemm.h).
#include "gemm.h"
#include <cmath>
#include <cstdint>
#include <cstring>
#include <vector>
using namespace gemm_def;

namespace {

// bf16 shares fp32's exponent field, so widening is an exact bit-shift.
inline float bf16_to_f32(uint16_t bits) {
    uint32_t b = (uint32_t)bits << 16;
    float f;
    std::memcpy(&f, &b, sizeof(f));
    return f;
}

// Proper IEEE-754 half->single conversion (fp16, not bf16 -- B_scales/B_mins
// are real fp16, whose exponent field does NOT alias fp32's, so this needs
// full mantissa/exponent remapping, unlike bf16_to_f32 above).
inline float f16_to_f32(uint16_t h) {
    uint32_t sign = (uint32_t)(h & 0x8000) << 16;
    uint32_t exp  = (h >> 10) & 0x1F;
    uint32_t mant = h & 0x3FF;
    uint32_t bits;
    if (exp == 0) {
        if (mant == 0) {
            bits = sign;
        } else {
            exp = 1;
            while ((mant & 0x400) == 0) { mant <<= 1; --exp; }
            mant &= 0x3FF;
            bits = sign | ((exp + 112u) << 23) | (mant << 13);
        }
    } else if (exp == 0x1F) {
        bits = sign | 0x7F800000u | (mant << 13);
    } else {
        bits = sign | ((exp + 112u) << 23) | (mant << 13);
    }
    float f;
    std::memcpy(&f, &bits, sizeof(f));
    return f;
}

} // namespace

extern "C" void inner_gemm(const uint16_t* A, float* output, const uint8_t* B_q4,
                            const uint16_t* B_scales, const uint16_t* B_mins, int M)
{
    constexpr int K_super = K / 256;

    std::vector<float> h(K);
    std::vector<float> d_act(K_super);
    std::vector<int8_t> q8(K);

    for (int m = 0; m < M; ++m) {
        const uint16_t* a_row = A + (long)m * K;
        for (int k = 0; k < K; ++k) h[k] = bf16_to_f32(a_row[k]);

        // Dynamically quantize the activation row into 256-element blocks
        // (Q8_K-style: one scale, int8 values in [-127, 127]).
        for (int sb = 0; sb < K_super; ++sb) {
            float amax = 0.0f;
            for (int l = 0; l < 256; ++l) {
                float v = std::fabs(h[sb * 256 + l]);
                if (v > amax) amax = v;
            }
            const float d = amax / 127.0f;
            d_act[sb] = d;
            const float id = d > 0.0f ? 1.0f / d : 0.0f;
            for (int l = 0; l < 256; ++l) {
                int q = (int)std::lround(h[sb * 256 + l] * id);
                if (q > 127) q = 127;
                if (q < -127) q = -127;
                q8[sb * 256 + l] = (int8_t)q;
            }
        }

        float* out_row = output + (long)m * N;
        for (int n = 0; n < N; ++n) {
            const uint8_t* b_row = B_q4 + (long)n * K_half;
            const uint16_t* sc_row = B_scales + (long)n * K_sub;
            const uint16_t* mn_row = B_mins + (long)n * K_sub;

            float acc = 0.0f;
            for (int sb32 = 0; sb32 < K_sub; ++sb32) {
                const float scale = f16_to_f32(sc_row[sb32]);
                const float minv = f16_to_f32(mn_row[sb32]);
                const float dact = d_act[sb32 / 8];  // 8 sub-blocks per 256-elem superblock

                int32_t sumi = 0;
                int32_t bsum = 0;
                const int base = sb32 * 32;
                for (int l = 0; l < 32; ++l) {
                    const int k = base + l;
                    const uint8_t byte = b_row[k / 2];
                    const int nib = (k % 2 == 0) ? (byte & 0xF) : (byte >> 4);
                    const int8_t qv = q8[k];
                    sumi += nib * qv;
                    bsum += qv;
                }
                acc += dact * (scale * (float)sumi - minv * (float)bsum);
            }
            out_row[n] = acc;
        }
    }
}