// File: kernel.cpp
// Reference-scalar mha bf16 (fp32 accumulate/output).
// LLM target: replace this file with an optimised inner_mha.
// All per-definition constants live in mha_def:: (mha.h).
// Assumes NHeads == KvHeads (standard MHA, no GQA head-repeat) — matches this
// definition's reference implementation exactly.
#include "mha.h"
#include <cmath>
#include <cstring>
#include <vector>
using namespace mha_def;

namespace {
// bf16 shares fp32's exponent field, so widening is an exact bit-shift — no
// rounding/subnormal handling needed (unlike fp16->fp32).
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;
}
} // namespace

extern "C" void inner_mha(const uint16_t* Q, float* output, const uint16_t* K,
                           const uint16_t* V, int M, int S)
{
    const float scale = 1.0f / std::sqrt((float)HeadDim);
    std::vector<float> scores(S);

    for (int h = 0; h < NHeads; ++h) {
        for (int m = 0; m < M; ++m) {
            const uint16_t* q_row = Q + ((long)m * NHeads + h) * HeadDim;

            float max_score = -INFINITY;
            for (int s = 0; s < S; ++s) {
                const uint16_t* k_row = K + ((long)s * KvHeads + h) * HeadDim;
                float dot = 0.0f;
                for (int d = 0; d < HeadDim; ++d) {
                    dot += bf16_to_f32(q_row[d]) * bf16_to_f32(k_row[d]);
                }
                scores[s] = dot * scale;
                if (scores[s] > max_score) max_score = scores[s];
            }

            float sum = 0.0f;
            for (int s = 0; s < S; ++s) {
                scores[s] = std::exp(scores[s] - max_score);
                sum += scores[s];
            }

            float* out_row = output + ((long)m * NHeads + h) * HeadDim;
            for (int d = 0; d < HeadDim; ++d) out_row[d] = 0.0f;
            for (int s = 0; s < S; ++s) {
                const float w = scores[s] / sum;
                const uint16_t* v_row = V + ((long)s * KvHeads + h) * HeadDim;
                for (int d = 0; d < HeadDim; ++d) {
                    out_row[d] += w * bf16_to_f32(v_row[d]);
                }
            }
        }
    }
}