// File: kernel.cpp
// Reference-scalar gemm Q8_0 (activations + weights both int8, per-32-block fp16
// scale). LLM target: replace this file with an optimised inner_gemm.
// All per-definition constants live in gemm_def:: (gemm.h).
#include "gemm.h"
#include <cstring>
#include <cstdint>
#include <vector>
using namespace gemm_def;

namespace {

// Proper IEEE-754 half->single conversion (fp16, not bf16).
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;
}

// Dequantize a [rows, K] int8 matrix (32-wide blocks, one fp16 scale per block)
// to float32.
void dequant(const int8_t* x, const uint16_t* scales, int rows, std::vector<float>& out) {
    out.resize((size_t)rows * K);
    for (int r = 0; r < rows; ++r) {
        const int8_t* xr = x + (size_t)r * K;
        const uint16_t* sr = scales + (size_t)r * K_blk;
        float* outr = &out[(size_t)r * K];
        for (int blk = 0; blk < K_blk; ++blk) {
            const float d = f16_to_f32(sr[blk]);
            for (int l = 0; l < 32; ++l) outr[blk * 32 + l] = d * (float)xr[blk * 32 + l];
        }
    }
}

} // namespace

extern "C" void inner_gemm(const int8_t* A, float* output, const uint16_t* A_scales,
                            const int8_t* B, const uint16_t* B_scales, int M)
{
    std::vector<float> Af, Bf;
    dequant(A, A_scales, M, Af);
    dequant(B, B_scales, N, Bf);
    for (int m = 0; m < M; ++m) {
        const float* ar = &Af[(size_t)m * K];
        float* out_row = output + (size_t)m * N;
        for (int n = 0; n < N; ++n) {
            const float* br = &Bf[(size_t)n * K];
            double acc = 0.0;
            for (int k = 0; k < K; ++k) acc += (double)ar[k] * br[k];
            out_row[n] = (float)acc;
        }
    }
}