tests.cpp

// slop tests
#include <cstdint>
#include <cstdio>
#include <filesystem>
#include <fstream>
#include <limits>
#include <string>
#include <variant>
#include <vector>

#include "config.hpp"
#include "isa.hpp"
#include "cost.hpp"
#include "array.hpp"
#include "machine.hpp"
#include "sim.hpp"
#include "harness.hpp"
#include "selfcheck.hpp"
#include "lower.hpp"
#include "golden.hpp"

// shared by every group
struct Tally {
    int passed = 0;
    int failed = 0;
    void check(bool ok, const char* what) {
        (ok ? passed : failed)++;
        std::printf("  [%s] %s\n", ok ? "PASS" : "FAIL", what);
    }
};
 
inline void functional_tests(Tally& t) {
    std::printf("=== functional ===\n");

    // variuos shapes
    struct S { std::size_t R, C, K; };
    for (S s : {S{8,8,8}, S{8,8,1}, S{8,8,64}, S{4,6,10}, S{16,4,32}, S{1,1,5}})
        for (Dataflow df : {Dataflow::OutputStationary, Dataflow::WeightStationary}) {
            const CaseResult r = run_case({s.R, s.C, s.K, df});
            t.check(r.sim_out == r.golden, "sim == golden");
        }

    // input vals
    for (InputKind k : {InputKind::Zeros, InputKind::Ones, InputKind::Identity,
                        InputKind::Int8Max, InputKind::Int8Min, InputKind::Random}) {
        const CaseResult r = run_case({8, 8, 16, Dataflow::OutputStationary, false, k});
        t.check(r.sim_out == r.golden, "corner == golden");
    }

    const CaseResult acc = run_case({8, 8, 8, Dataflow::OutputStationary, true});
    t.check(acc.sim_out == acc.golden, "accumulate == golden");

    t.check(run_case({8, 8, 8, Dataflow::OutputStationary}).sim_out ==
            run_case({8, 8, 8, Dataflow::WeightStationary}).sim_out, "OS == WS");

    // check for overflow
    Machine m; m.cfg = Config{1, 1};
    allocate(m, 64);
    m.mem.data[0] = 127; m.mem.data[1] = 127;
    m.accbanks.bank[2][0] = std::numeric_limits<Acc>::max() - 1;
    run(m, {Config_{Dataflow::OutputStationary},
            Load{0, {0,0}, Queue::Dma, 1, 1}, Load{1, {1,0}, Queue::Dma, 1, 1},
            Matmul{{0,0}, {1,0}, {2,0}, 1, true}});
    t.check(m.accbanks.bank[2][0] == std::numeric_limits<Acc>::max(), "saturation clamps");
}
inline void timing_tests(Tally& t) {
    std::printf("=== timing ===\n");

    const std::size_t Ks[] = {1, 4, 8, 32};
    for (Dataflow df : {Dataflow::OutputStationary, Dataflow::WeightStationary})
        for (std::size_t K : Ks) {
            Config c; c.dataflow = df;
            const ArrayRun r = (df == Dataflow::OutputStationary)
                                   ? step_array_os(c, K) : step_array_ws(c, K);
            t.check(r.cycles == cost(Matmul{{0,0},{0,0},{0,0},K}, c), "latency == cost");
            t.check(r.pe_cycles == c.rows * c.cols * K, "pe_cycles == R*C*K");
        }

    // serial stream
    Config c;
    const std::vector<Inst> prog = {
        Config_{c.dataflow},
        Load{0,  {0,0}, Queue::Dma, 1, 64},
        Load{64, {1,0}, Queue::Dma, 1, 64},
        Matmul{{0,0}, {1,0}, {2,0}, 8},
        Store{{2,0}, 128, Queue::Dma, 1, 256},
    };
    Machine m; m.cfg = c; allocate(m, 1 << 20);
    const Counters k = run(m, prog);

    std::uint32_t sum = 0, useful = 0;   // useful == sum of mm(K) 
    for (const Inst& in : prog) {
        sum += cost(in, c);
        if (const Matmul* mm = std::get_if<Matmul>(&in)) useful += mm->k;
    }
    t.check(k.total_cycles <= sum, "total <= sum(cost)");
    t.check(useful + k.stall_fill_drain + k.stall_memory == k.total_cycles, "waterfall balances");
    Config c22{2, 2};
    t.check(step_array_os(c22, 2).cycles == 4, "2x2 K=2 -> 4");
    t.check(step_array_os(c, 8).cycles == 22, "8x8 K=8 -> 22");
}

inline std::vector<Acc> run_sim_tile(std::size_t R, std::size_t C, std::size_t K,
                                     const std::vector<Elem>& A, const std::vector<Elem>& B) {
    Config cfg; cfg.rows = R; cfg.cols = C;
    Machine m; m.cfg = cfg; allocate(m, 1 << 20);
    for (std::size_t i = 0; i < A.size(); ++i) m.mem.data[i] = A[i];
    for (std::size_t i = 0; i < B.size(); ++i) m.mem.data[A.size() + i] = B[i];
    run(m, {Config_{cfg.dataflow},
            Load{0, {0,0}, Queue::Dma, 1, A.size()},
            Load{A.size(), {1,0}, Queue::Dma, 1, B.size()},
            Matmul{{0,0}, {1,0}, {2,0}, K}});
    std::vector<Acc> out(R * C);
    for (std::size_t i = 0; i < out.size(); ++i) out[i] = m.accbanks.bank[2][i];
    return out;
}

inline void numpy_crosscheck(Tally& t) {
    std::printf("=== fixture cross-check ===\n");
    const std::filesystem::path dir = "fixtures";
    if (!std::filesystem::exists(dir)) {
        std::printf("  [SKIP] no fixtures/ (run: python3 fixtures/gen_fixtures.py)\n");
        return;
    }
    for (const auto& e : std::filesystem::directory_iterator(dir)) {
        if (e.path().extension() != ".txt") continue;
        std::ifstream f(e.path());
        std::size_t R, C, K;
        f >> R >> C >> K;
        std::vector<Elem> A(R * K), B(K * C);
        std::vector<Acc>  want(R * C);
        for (Elem& x : A) { int v; f >> v; x = static_cast<Elem>(v); }  // >> on int8 reads a char
        for (Elem& x : B) { int v; f >> v; x = static_cast<Elem>(v); }
        for (Acc& x : want) f >> x;
        t.check(run_sim_tile(R, C, K, A, B) == want, e.path().filename().string().c_str());
    }
}

struct GemmRun { std::vector<Acc> out; Counters ctr; };

// lower an (m,n,k) gemm onto the array, run it, return the int32 out + counters
inline GemmRun run_gemm_full(std::size_t M, std::size_t N, std::size_t K,
                             const std::vector<Elem>& A, const std::vector<Elem>& B,
                             Dataflow df, Sched sched) {
    Config cfg; cfg.dataflow = df;               // default 8x8 array
    Machine m; m.cfg = cfg; allocate(m, 1 << 20);
    m.out.assign(M * N, Acc{0});
    for (std::size_t i = 0; i < A.size(); ++i) m.mem.data[i] = A[i];
    for (std::size_t i = 0; i < B.size(); ++i) m.mem.data[A.size() + i] = B[i];
    GemmRun r;
    r.ctr = run(m, lower(GemmSpec{M, N, K, 0, A.size()}, cfg, sched));
    r.out.assign(m.out.begin(), m.out.begin() + M * N);
    return r;
}

inline std::vector<Acc> run_gemm(std::size_t M, std::size_t N, std::size_t K,
                                 const std::vector<Elem>& A, const std::vector<Elem>& B,
                                 Dataflow df) {
    return run_gemm_full(M, N, K, A, B, df, Sched::Naive).out;
}

inline void lowering_tests(Tally& t) {
    std::printf("=== lowering ===\n");
    struct S { std::size_t M, N, K; };
    // exact tiles, partial r/c/k tiles, split-k, tall/wide, degenerate
    for (S s : {S{8,8,8}, S{8,8,16}, S{16,16,8}, S{10,12,20}, S{5,3,17}, S{1,1,9}, S{20,7,4}})
        for (Dataflow df : {Dataflow::OutputStationary, Dataflow::WeightStationary}) {
            std::vector<Elem> A(s.M * s.K), B(s.K * s.N);
            std::uint64_t x = s.M * 131 + s.N * 17 + s.K + 1;
            for (Elem& e : A) { x = x*6364136223846793005ULL + 1; e = static_cast<Elem>((x>>56)%7 - 3); }
            for (Elem& e : B) { x = x*6364136223846793005ULL + 1; e = static_cast<Elem>((x>>56)%7 - 3); }
            const bool ok = run_gemm(s.M, s.N, s.K, A, B, df) ==
                            reference_gemm_mnk(s.M, s.N, s.K, A.data(), B.data());
            t.check(ok, "lowered gemm == golden");
        }
}

// serial cost split into the two queues; pipelined total can't beat the busier one
inline void queue_bounds(std::size_t M, std::size_t N, std::size_t K,
                         const Config& cfg, std::uint64_t& compute_b, std::uint64_t& dma_b) {
    compute_b = dma_b = 0;
    for (const Inst& in : lower(GemmSpec{M, N, K, 0, M * K}, cfg)) {
        if (std::holds_alternative<Matmul>(in)) compute_b += cost(in, cfg);
        else if (!std::holds_alternative<Config_>(in) && !std::holds_alternative<Fence>(in))
            dma_b += cost(in, cfg);
    }
}

// naive vs pipelined lowering, same kernel: overlap must help and stay correct
inline void scheduling_tests(Tally& t) {
    std::printf("=== scheduling ===\n");
    struct S { std::size_t M, N, K; };
    for (S s : {S{8,8,64}, S{16,16,32}, S{8,8,24}, S{10,12,20}}) {
        std::vector<Elem> A(s.M * s.K), B(s.K * s.N);
        std::uint64_t x = s.M * 131 + s.N * 17 + s.K + 7;
        for (Elem& e : A) { x = x*6364136223846793005ULL + 1; e = static_cast<Elem>((x>>56)%7 - 3); }
        for (Elem& e : B) { x = x*6364136223846793005ULL + 1; e = static_cast<Elem>((x>>56)%7 - 3); }

        const auto want = reference_gemm_mnk(s.M, s.N, s.K, A.data(), B.data());
        const GemmRun nv = run_gemm_full(s.M, s.N, s.K, A, B, Dataflow::OutputStationary, Sched::Naive);
        const GemmRun pp = run_gemm_full(s.M, s.N, s.K, A, B, Dataflow::OutputStationary, Sched::Pipelined);

        Config cfg;
        std::uint64_t cb, db;
        queue_bounds(s.M, s.N, s.K, cfg, cb, db);

        t.check(pp.out == want, "pipelined == golden");
        t.check(pp.ctr.total_cycles <= nv.ctr.total_cycles, "pipelined not slower");
        t.check(pp.ctr.total_cycles >= std::max(cb, db), "pipelined respects queue bound");
    }
}

int main() {
    Tally t;

    functional_tests(t);
    timing_tests(t);
    numpy_crosscheck(t);
    lowering_tests(t);
    scheduling_tests(t);

    Config os; os.dataflow = Dataflow::OutputStationary;
    Config ws; ws.dataflow = Dataflow::WeightStationary;
    t.check(run_self_checks(os), "self-checks (OS)");
    t.check(run_self_checks(ws), "self-checks (WS)");

    std::printf("=== summary: %d passed, %d failed ===\n", t.passed, t.failed);
    return t.failed == 0 ? 0 : 1;
}