// CPU reference correctness tests for the sparkinfer kernel algorithms.
//
// These re-implement each CUDA kernel's exact numerical algorithm in plain C++
// and check it against an INDEPENDENT double-precision ground truth (different
// loop order / higher precision). A match is real evidence the algorithm is
// correct; the device-side sm_120 compile (see .cudaverify) separately proves
// the same code targets the RTX 5090. Together they cover "valid for the 5090"
// and "computes the right thing" — the two halves a GPU-less environment allows.
//
// Build: g++ -O2 -std=c++17 cpu_reference_test.cpp -o cpu_reference_test

#include <cstdio>
#include <cmath>
#include <cstring>
#include <cstdint>
#include <vector>
#include <random>
#include <algorithm>
#include <utility>

using std::vector;
static std::mt19937 rng(1234);
static float frand() { return std::uniform_real_distribution<float>(-1.f, 1.f)(rng); }

// Round a float to bf16 (round-to-nearest-even), returned as float. Models the
// __float2bfloat16 round-trip the kernels do when reading/writing bf16 rows.
static float to_bf16(float f) {
    uint32_t u; std::memcpy(&u, &f, 4);
    if ((u & 0x7fffffffu) > 0x7f800000u) return f;  // NaN: leave as-is
    const uint32_t lsb = (u >> 16) & 1u;
    u += 0x7fffu + lsb;
    u &= 0xffff0000u;
    float r; std::memcpy(&r, &u, 4);
    return r;
}

static int g_fail = 0;
static void check(const char* name, double max_err, double tol) {
    bool ok = max_err <= tol;
    printf("  [%s] %-34s max_err=%.3e (tol=%.0e)\n", ok ? "PASS" : "FAIL", name, max_err, tol);
    if (!ok) g_fail++;
}

static float silu(float x) { return x / (1.f + std::exp(-x)); }

// ---------------------------------------------------------------------------
// 1. Flash decode: online-softmax (kernel algorithm) vs naive full softmax.
// ---------------------------------------------------------------------------
static double test_attention(int HD, int kvlen) {
    vector<float> q(HD), K(kvlen * HD), V(kvlen * HD);
    for (auto& x : q) x = frand();
    for (auto& x : K) x = frand();
    for (auto& x : V) x = frand();
    const float scale = 1.f / std::sqrt((float)HD);

    // Ground truth (double precision, two-pass softmax).
    vector<double> scores(kvlen);
    double mx = -1e300;
    for (int t = 0; t < kvlen; t++) {
        double d = 0; for (int i = 0; i < HD; i++) d += (double)q[i] * K[t * HD + i];
        scores[t] = d * scale; mx = std::max(mx, scores[t]);
    }
    double denom = 0; for (int t = 0; t < kvlen; t++) denom += std::exp(scores[t] - mx);
    vector<double> ref(HD, 0);
    for (int t = 0; t < kvlen; t++) {
        double p = std::exp(scores[t] - mx) / denom;
        for (int i = 0; i < HD; i++) ref[i] += p * V[t * HD + i];
    }

    // Kernel algorithm: single-pass online softmax in float.
    float m = -1e30f, l = 0.f; vector<float> acc(HD, 0.f);
    for (int t = 0; t < kvlen; t++) {
        float d = 0; for (int i = 0; i < HD; i++) d += q[i] * K[t * HD + i];
        float score = d * scale;
        float m_new = std::max(m, score);
        float corr = std::exp(m - m_new), p = std::exp(score - m_new);
        l = l * corr + p;
        for (int i = 0; i < HD; i++) acc[i] = acc[i] * corr + p * V[t * HD + i];
        m = m_new;
    }
    double err = 0; for (int i = 0; i < HD; i++) err = std::max(err, std::abs(acc[i] / l - ref[i]));
    return err;
}

// ---------------------------------------------------------------------------
// 2. Router top-k: kernel mask-argmax algorithm vs sort-based reference.
// ---------------------------------------------------------------------------
static double test_router(int E, int K) {
    vector<float> logits(E); for (auto& x : logits) x = frand();

    // Reference: stable sort by (value desc, index asc), take K; softmax over them.
    vector<int> idx(E); for (int i = 0; i < E; i++) idx[i] = i;
    std::stable_sort(idx.begin(), idx.end(), [&](int a, int b) {
        return logits[a] > logits[b] || (logits[a] == logits[b] && a < b); });
    vector<int> ref_id(idx.begin(), idx.begin() + K);
    double rmx = logits[ref_id[0]], rden = 0;
    for (int j = 0; j < K; j++) rden += std::exp((double)logits[ref_id[j]] - rmx);
    vector<double> ref_w(K);
    for (int j = 0; j < K; j++) ref_w[j] = std::exp((double)logits[ref_id[j]] - rmx) / rden;

    // Kernel algorithm: K passes of arg-max with masking, then softmax over picks.
    vector<float> s = logits; vector<int> sel(K); vector<float> sl(K);
    for (int j = 0; j < K; j++) {
        float best = -1e30f; int bi = -1;
        for (int e = 0; e < E; e++) if (s[e] > best || (s[e] == best && e < bi)) { best = s[e]; bi = e; }
        sel[j] = bi; sl[j] = best; s[bi] = -1e30f;
    }
    float kmx = sl[0]; for (int j = 1; j < K; j++) kmx = std::max(kmx, sl[j]);
    float kden = 0; for (int j = 0; j < K; j++) kden += std::exp(sl[j] - kmx);

    double err = 0;
    for (int j = 0; j < K; j++) {
        if (sel[j] != ref_id[j]) err = std::max(err, 1.0);
        err = std::max(err, std::abs(std::exp(sl[j] - kmx) / kden - ref_w[j]));
    }
    return err;
}

// ---------------------------------------------------------------------------
// 3. SwiGLU expert FFN: kernel math (float) vs double ground truth.
// ---------------------------------------------------------------------------
static double test_swiglu(int H, int F) {
    vector<float> X(H), gate(H * F), up(H * F), down(F * H);
    for (auto& x : X) x = frand();
    for (auto& x : gate) x = frand() * 0.1f;
    for (auto& x : up) x = frand() * 0.1f;
    for (auto& x : down) x = frand() * 0.1f;
    const float w = 0.37f;

    vector<double> hbuf_d(F), ref(H, 0);
    for (int f = 0; f < F; f++) {
        double g = 0, u = 0;
        for (int h = 0; h < H; h++) { g += (double)X[h] * gate[h * F + f]; u += (double)X[h] * up[h * F + f]; }
        hbuf_d[f] = (g / (1.0 + std::exp(-g))) * u;
    }
    for (int h = 0; h < H; h++) { double y = 0; for (int f = 0; f < F; f++) y += hbuf_d[f] * down[f * H + h]; ref[h] = w * y; }

    vector<float> hbuf(F), acc(H, 0.f);
    for (int f = 0; f < F; f++) {
        float g = 0, u = 0;
        for (int h = 0; h < H; h++) { g += X[h] * gate[h * F + f]; u += X[h] * up[h * F + f]; }
        hbuf[f] = silu(g) * u;
    }
    for (int h = 0; h < H; h++) { float y = 0; for (int f = 0; f < F; f++) y += hbuf[f] * down[f * H + h]; acc[h] = w * y; }

    double err = 0; for (int h = 0; h < H; h++) err = std::max(err, std::abs((double)acc[h] - ref[h]));
    return err;
}

// ---------------------------------------------------------------------------
// 4. GEMM: tiled accumulation order vs double triple-loop.
// ---------------------------------------------------------------------------
static double test_gemm(int M, int N, int Kd) {
    vector<float> A(M * Kd), B(Kd * N);
    for (auto& x : A) x = frand();
    for (auto& x : B) x = frand();
    vector<double> ref(M * N, 0);
    for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) { double s = 0; for (int k = 0; k < Kd; k++) s += (double)A[i*Kd+k]*B[k*N+j]; ref[i*N+j] = s; }

    const int TILE = 16; vector<float> C(M * N, 0.f);
    for (int i = 0; i < M; i++) for (int j = 0; j < N; j++) {
        float acc = 0.f;
        for (int k0 = 0; k0 < Kd; k0 += TILE) { float t = 0.f; for (int k = k0; k < std::min(k0+TILE,Kd); k++) t += A[i*Kd+k]*B[k*N+j]; acc += t; }
        C[i*N+j] = acc;
    }
    double err = 0; for (int i = 0; i < M*N; i++) err = std::max(err, std::abs((double)C[i] - ref[i]));
    return err;
}

// ---------------------------------------------------------------------------
// 5. RMSNorm: kernel math vs double ground truth.
// ---------------------------------------------------------------------------
static double test_rmsnorm(int cols) {
    vector<float> x(cols), wt(cols); for (auto& v : x) v = frand(); for (auto& v : wt) v = frand();
    const float eps = 1e-6f;
    double ss = 0; for (int c = 0; c < cols; c++) ss += (double)x[c]*x[c];
    double inv = 1.0 / std::sqrt(ss / cols + eps);
    vector<double> ref(cols); for (int c = 0; c < cols; c++) ref[c] = x[c]*inv*wt[c];

    float fss = 0; for (int c = 0; c < cols; c++) fss += x[c]*x[c];
    float finv = 1.f/std::sqrt(fss/cols + eps);
    double err = 0; for (int c = 0; c < cols; c++) err = std::max(err, std::abs((double)(x[c]*finv*wt[c]) - ref[c]));
    return err;
}

// ---------------------------------------------------------------------------
// 5b. Vectorized RMSNorm (PR #44): the kernel reads the row in 8-wide (uint4 =
//     8 bf16) packs and accumulates the sum-of-squares per-pack with FMA. Only
//     the per-thread element grouping in the SS reduction changes (FP assoc.);
//     this checks the 8-wide grouped reduction still matches the fp64 ground
//     truth. cols here are multiples of 8, as on every real RMSNorm call site.
// ---------------------------------------------------------------------------
static double test_rmsnorm_vec8(int cols) {
    vector<float> x(cols), wt(cols);
    for (auto& v : x) v = frand();
    for (auto& v : wt) v = frand();
    const float eps = 1e-6f;

    double ss = 0; for (int c = 0; c < cols; c++) ss += (double)x[c]*x[c];
    double inv = 1.0 / std::sqrt(ss / cols + eps);
    vector<double> ref(cols); for (int c = 0; c < cols; c++) ref[c] = x[c]*inv*wt[c];

    // 8-wide packs, FMA accumulation (mirrors rn_unpack8 + __fmaf_rn).
    const int npack = cols >> 3;
    float fss = 0.f;
    for (int p = 0; p < npack; p++) {
        #pragma GCC unroll 8
        for (int j = 0; j < 8; j++) { float v = x[p*8+j]; fss = std::fma(v, v, fss); }
    }
    for (int c = (npack<<3); c < cols; c++) { float v = x[c]; fss = std::fma(v, v, fss); }
    float finv = 1.f/std::sqrt(fss/cols + eps);
    double err = 0; for (int c = 0; c < cols; c++) err = std::max(err, std::abs((double)(x[c]*finv*wt[c]) - ref[c]));
    return err;
}

// ---------------------------------------------------------------------------
// 5c. add_rmsnorm2 sequencing (PR #44): the fused residual+norm kernel must keep
//     its exact numeric sequencing under vectorization — SS accumulates on the
//     fp32 sum (x+residual), the sum is round-tripped through bf16 into out_sum,
//     and the norm pass re-reads that bf16 sum. This models that scalar vs 8-wide
//     grouped sequencing produce the same normalized output (bf16-exact), so the
//     vectorized rewrite is byte-faithful, not just close.
// ---------------------------------------------------------------------------
static double test_add_rmsnorm2_seq(int cols) {
    vector<float> x(cols), r(cols), wt(cols);
    for (auto& v : x)  v = frand();
    for (auto& v : r)  v = frand();
    for (auto& v : wt) v = frand();
    const float eps = 1e-6f;

    // Scalar reference sequencing (original kernel).
    vector<float> sum_bf(cols);
    float ss_s = 0.f;
    for (int c = 0; c < cols; c++) {
        float v = x[c] + r[c];          // fp32 sum
        sum_bf[c] = to_bf16(v);          // out_sum stored as bf16
        ss_s = std::fma(v, v, ss_s);     // SS on the fp32 sum
    }
    float inv_s = 1.f/std::sqrt(ss_s/cols + eps);
    vector<float> norm_s(cols);
    for (int c = 0; c < cols; c++) norm_s[c] = to_bf16(sum_bf[c] * inv_s * wt[c]);

    // 8-wide grouped sequencing (PR kernel): same per-element ops, packed.
    const int npack = cols >> 3;
    vector<float> sum_bf2(cols);
    float ss_v = 0.f;
    for (int p = 0; p < npack; p++)
        for (int j = 0; j < 8; j++) {
            float v = x[p*8+j] + r[p*8+j];
            sum_bf2[p*8+j] = to_bf16(v);
            ss_v = std::fma(v, v, ss_v);
        }
    float inv_v = 1.f/std::sqrt(ss_v/cols + eps);
    vector<float> norm_v(cols);
    for (int p = 0; p < npack; p++)
        for (int j = 0; j < 8; j++)
            norm_v[p*8+j] = to_bf16(sum_bf2[p*8+j] * inv_v * wt[p*8+j]);

    double err = 0;
    for (int c = 0; c < cols; c++) {
        err = std::max(err, std::abs((double)sum_bf2[c] - sum_bf[c]));
        err = std::max(err, std::abs((double)norm_v[c]  - norm_s[c]));
    }
    return err;  // expect 0: scalar and 8-wide grouping are bit-identical here
}

// ---------------------------------------------------------------------------
// argmax two-pass (decode): the multi-block scan + final reduce must return the
// SAME index as a serial argmax, including the smallest-index tie-break.
// ---------------------------------------------------------------------------
static double test_argmax_twopass(int vocab, int nblocks, bool ties) {
    vector<float> L(vocab);
    for (auto& v : L) v = frand();
    if (ties) { for (auto& v : L) v = 0.f; for (int i : {1000 % vocab, 5, 77777 % vocab, 250}) L[i] = 7.f; }

    auto merge = [](float& bv, int& bi, float ov, int oi) {
        if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; }
    };
    // serial ground truth (smallest index on ties)
    float gv = -1e30f; int gi = 0;
    for (int v = 0; v < vocab; v++) merge(gv, gi, L[v], v);

    // pass 1: nblocks grid-stride partials, each block 256 threads
    const int BT = 256;
    vector<float> pv(nblocks); vector<int> pi(nblocks);
    for (int b = 0; b < nblocks; b++) {
        float bbv = -1e30f; int bbi = 0;
        vector<float> tv(BT, -1e30f); vector<int> ti(BT, 0);
        for (int t = 0; t < BT; t++)
            for (int v = b * BT + t; v < vocab; v += BT * nblocks) merge(tv[t], ti[t], L[v], v);
        for (int t = 0; t < BT; t++) merge(bbv, bbi, tv[t], ti[t]);
        pv[b] = bbv; pi[b] = bbi;
    }
    // pass 2: reduce partials
    float rv = -1e30f; int ri = 0;
    for (int b = 0; b < nblocks; b++) merge(rv, ri, pv[b], pi[b]);
    return (double)std::abs(ri - gi);   // expect 0: same index as serial argmax
}

int main() {
    printf("sparkinfer kernel algorithm correctness (CPU reference)\n");
    check("attention hd128 kv1",   test_attention(128, 1),    1e-4);
    check("attention hd128 kv333", test_attention(128, 333),  1e-4);
    check("attention hd256 kv1024",test_attention(256, 1024), 2e-4);
    check("attention hd512 kv777", test_attention(512, 777),  2e-4);
    check("router E256 k8",        test_router(256, 8),       1e-6);
    check("router E128 k8",        test_router(128, 8),       1e-6);
    check("swiglu H2048 F512",     test_swiglu(2048, 512),    1e-3);
    check("swiglu H512 F1536",     test_swiglu(512, 1536),    1e-3);
    check("gemm 64x96x128",        test_gemm(64, 96, 128),    1e-3);
    check("gemm 17x33x49",         test_gemm(17, 33, 49),     1e-3);
    check("rmsnorm cols2048",      test_rmsnorm(2048),        1e-4);
    check("rmsnorm vec8 cols2048", test_rmsnorm_vec8(2048),   1e-4);
    check("rmsnorm vec8 cols128",  test_rmsnorm_vec8(128),    1e-4);
    check("add_rmsnorm2 seq c2048",test_add_rmsnorm2_seq(2048),1e-9);
    check("add_rmsnorm2 seq c1536",test_add_rmsnorm2_seq(1536),1e-9);
    check("argmax 2pass qwen vocab",test_argmax_twopass(151936, 512, false), 0.0);
    check("argmax 2pass gemma vocab",test_argmax_twopass(262144, 512, false), 0.0);
    check("argmax 2pass tie-break",  test_argmax_twopass(151936, 512, true),  0.0);
    printf("%s (%d failures)\n", g_fail ? "FAILED" : "ALL PASSED", g_fail);
    return g_fail ? 1 : 0;
}
