bench.cpp

#include <cstdio>
#include <cstdlib>
#include <vector>

#include "config.hpp"
#include "machine.hpp"
#include "sim.hpp"
#include "lower.hpp"
#include "golden.hpp"

// deterministic ops so a cycle delta is applicable to the swept axis
static void fill_ab(std::vector<Elem>& A, std::vector<Elem>& B) {
    for (std::size_t i = 0; i < A.size(); ++i) A[i] = static_cast<Elem>((i % 5) - 2);
    for (std::size_t i = 0; i < B.size(); ++i) B[i] = static_cast<Elem>((i % 3) - 1);
}

struct Run { Counters ctr; bool golden_ok; };

// one gemm through lower + run on a given array/dataflow/sched
static Run run_one(const Config& cfg, std::size_t M, std::size_t N, std::size_t K, Sched sched) {
    std::vector<Elem> A(M * K), B(K * N);
    fill_ab(A, B);

    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];

    Run r;
    r.ctr = run(m, lower(GemmSpec{M, N, K, 0, A.size()}, cfg, sched));
    const std::vector<Acc> want = reference_gemm_mnk(M, N, K, A.data(), B.data());
    r.golden_ok = std::vector<Acc>(m.out.begin(), m.out.begin() + M * N) == want;
    return r;
}

// scalar cpu doing the same arithmetic, one mac per cycle
static std::uint64_t cpu_cycles(std::size_t M, std::size_t N, std::size_t K) {
    return static_cast<std::uint64_t>(M) * N * K;
}

static const char* df_name(Dataflow d) {
    return d == Dataflow::OutputStationary ? "OS" : "WS";
}

// one csv row: sweep axes then the whole Counters then the two derived speedups
// pipe holds naive/pipelined for the sched sweep, 0 elsewhere (means n/a)
static void row(const char* kind, const Config& cfg, Sched sched,
                std::size_t M, std::size_t N, std::size_t K,
                const Counters& c, float pipe) {
    const std::uint64_t cpu = cpu_cycles(M, N, K);
    const float vs_cpu = c.total_cycles ? static_cast<float>(cpu) / c.total_cycles : 0.0f;

    // util = busy_frac * eff_busy: how much of the time the array runs, times how
    // efficient it is while running (fill/drain is the loss inside eff_busy)
    const float busy = c.total_cycles ? static_cast<float>(c.compute_cycles) / c.total_cycles : 0.0f;
    const float eff  = c.compute_cycles ? c.utilization * c.total_cycles / c.compute_cycles : 0.0f;
    const float fd   = c.total_cycles ? static_cast<float>(c.stall_fill_drain) / c.total_cycles : 0.0f;

    std::printf("%s,%zu,%zu,%s,%zu,%s,%zu,%zu,%zu,%u,%u,%u,%.4f,%.4f,%u,%u,%llu,%.3f,%.3f,%.4f,%.4f,%.4f\n",
                kind, cfg.rows, cfg.cols, df_name(cfg.dataflow), KT,
                sched == Sched::Naive ? "naive" : "pipelined", M, N, K,
                c.total_cycles, c.compute_cycles, c.dma_cycles,
                c.utilization, c.bandwidth, c.stall_memory, c.stall_fill_drain,
                static_cast<unsigned long long>(cpu), vs_cpu, pipe,
                busy, eff, fd);
}

int main() {
    std::printf("kind,R,C,dataflow,KT,sched,M,N,K,"
                "total_cycles,compute_cycles,dma_cycles,utilization,bandwidth,"
                "stall_memory,stall_fill_drain,cpu_cycles,speedup_vs_cpu,speedup_pipeline,"
                "busy_frac,eff_busy,fill_drain_frac\n");

    bool all_golden = true;
    auto emit = [&](const char* kind, const Config& cfg, Sched sched,
                    std::size_t M, std::size_t N, std::size_t K, float pipe) {
        const Run r = run_one(cfg, M, N, K, sched);
        all_golden = all_golden && r.golden_ok;
        row(kind, cfg, sched, M, N, K, r.ctr, pipe);
    };

    // array-multiple shapes so utilization's full-array numerator is honest
    // sweep 1: utilization vs array size, same gemm, one axis (R=C) moving
    for (std::size_t s : {std::size_t{4}, std::size_t{8}, std::size_t{16}, std::size_t{32}}) {
        Config cfg; cfg.rows = s; cfg.cols = s;
        emit("array_size", cfg, Sched::Naive, 32, 32, 64, 0.0f);
    }

    // sweep 2: OS vs WS, same gemm, dataflow is the only axis
    for (Dataflow df : {Dataflow::OutputStationary, Dataflow::WeightStationary}) {
        Config cfg; cfg.dataflow = df;
        emit("dataflow", cfg, Sched::Naive, 32, 32, 64, 0.0f);
    }

    // sweep 3: naive vs pipelined, K-heavy shapes where overlap pays
    struct Shape { std::size_t M, N, K; };
    for (Shape sh : {Shape{8, 8, 64}, Shape{16, 16, 64}, Shape{8, 8, 128}, Shape{32, 32, 64}}) {
        Config cfg;
        const Run nv = run_one(cfg, sh.M, sh.N, sh.K, Sched::Naive);
        const Run pp = run_one(cfg, sh.M, sh.N, sh.K, Sched::Pipelined);
        all_golden = all_golden && nv.golden_ok && pp.golden_ok;
        const float pipe = pp.ctr.total_cycles
                               ? static_cast<float>(nv.ctr.total_cycles) / pp.ctr.total_cycles
                               : 0.0f;
        row("sched", cfg, Sched::Naive,     sh.M, sh.N, sh.K, nv.ctr, pipe);
        row("sched", cfg, Sched::Pipelined, sh.M, sh.N, sh.K, pp.ctr, pipe);
    }

    // sweep 4: accelerator vs cpu across shapes, best (pipelined) lowering
    for (Shape sh : {Shape{8, 8, 8}, Shape{8, 8, 64}, Shape{16, 16, 64},
                     Shape{32, 32, 64}, Shape{32, 32, 128}}) {
        Config cfg;
        emit("vs_cpu", cfg, Sched::Pipelined, sh.M, sh.N, sh.K, 0.0f);
    }

    if (!all_golden) {
        std::fprintf(stderr, "WARNING: a bench run did not match golden\n");
        return 1;
    }
    return 0;
}