// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// GEMM throughput benchmark. Times each kernel over the shapes a transformer
// step actually issues (fwd X·Wᵀ, dX = dY·W, dW = dYᵀ·X) plus square shapes
// for a clean TFLOPs number. GPU time comes from command-buffer
// GPUStartTime/GPUEndTime, so it excludes CPU encode overhead.
#include <Foundation/Foundation.hpp>
#include <Metal/Metal.hpp>

#include "core/device.h"
#include "core/tensor.h"
#include "ops/metal/metal_ops.h"

#include <cstdio>
#include <random>
#include <vector>

using forge::metal::MatmulKernel;

namespace {

struct Case {
    int64_t M, K, N;
    bool ta, tb;
    const char* label;
};

double bench(const Case& c, MatmulKernel kernel, int iters) {
    forge::Tensor a = c.ta ? forge::Tensor::empty({c.K, c.M})
                           : forge::Tensor::empty({c.M, c.K});
    forge::Tensor b = c.tb ? forge::Tensor::empty({c.N, c.K})
                           : forge::Tensor::empty({c.K, c.N});
    forge::Tensor out = forge::Tensor::empty({c.M, c.N});
    std::mt19937 rng(7);
    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
    for (int64_t i = 0; i < a.numel(); ++i) a.data<float>()[i] = dist(rng);
    for (int64_t i = 0; i < b.numel(); ++i) b.data<float>()[i] = dist(rng);

    // warmup (also builds the pipeline)
    forge::metal::matmul(a, b, out, c.ta, c.tb, false, kernel);
    forge::metal::sync();

    const double t0 = forge::metal::Stream::get().gpu_seconds();
    for (int i = 0; i < iters; ++i)
        forge::metal::matmul(a, b, out, c.ta, c.tb, false, kernel);
    forge::metal::sync();
    const double elapsed = forge::metal::Stream::get().gpu_seconds() - t0;

    const double flops = 2.0 * double(c.M) * double(c.N) * double(c.K) * iters;
    return flops / elapsed / 1e12; // TFLOP/s
}

} // namespace

int main() {
    NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init();
    std::printf("device: %s\n\n", forge::Device::get().name().c_str());

    // gpt-25m-ish shapes: batch*seq = 64*1024 rows, d_model 512, d_ff 1408
    const std::vector<Case> cases = {
        {4096, 4096, 4096, false, false, "square 4096"},
        {2048, 2048, 2048, false, false, "square 2048"},
        {1024, 1024, 1024, false, false, "square 1024"},
        {65536, 512, 1408, false, true, "fwd  mlp  X·W1ᵀ"},
        {65536, 1408, 512, false, true, "fwd  mlp  H·W2ᵀ"},
        {65536, 512, 512, false, true, "fwd  attn X·Wqᵀ"},
        {65536, 512, 1408, false, false, "bwd  dX = dY·W"},
        {512, 65536, 1408, true, false, "bwd  dW = dYᵀ·X"},
        {65536, 512, 4096, false, true, "lm head (vocab 4096)"},
    };

    std::printf("%-24s %10s %10s %10s\n", "case", "naive", "tiled", "simd");
    for (const Case& c : cases) {
        const double gflop = 2.0 * double(c.M) * double(c.N) * double(c.K) / 1e9;
        const int iters = gflop > 50.0 ? 3 : 10;
        const double n = bench(c, MatmulKernel::Naive, iters);
        const double t = bench(c, MatmulKernel::Tiled, iters);
        const double s = bench(c, MatmulKernel::Simdgroup, iters);
        std::printf("%-24s %9.2fT %9.2fT %9.2fT\n", c.label, n, t, s);
        std::fflush(stdout);
    }

    pool->drain();
    return 0;
}