// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Does simdgroup_matrix actually run faster with f16/bf16 operands than with
// f32 on this hardware? The literature disagrees: metal-benchmarks measures
// f16 and f32 FMA at the SAME rate on Apple GPUs (the win coming from
// registers/bandwidth), while Apple's M3 material claims up to 2x ALU for
// family 9 via FP16/FP32/INT co-issue. Since mixed precision is a large
// amount of work to plumb through a training framework, measure before
// committing.
//
// All three variants use the identical 64x64x16 tiling from matmul_simd and
// an f32 accumulator (required for training); only the staged-tile operand
// type and the MMA fragment type differ.
#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 <string>

using namespace forge;

namespace {

struct Params { uint32_t M, N, K; };

double run(const char* kernel, DType dt, int64_t M, int64_t N, int64_t K, int iters) {
    Tensor a = Tensor::empty({M, K}, dt);
    Tensor b = Tensor::empty({K, N}, dt);
    Tensor c = Tensor::empty({M, N}, DType::F32);
    std::mt19937 rng(3);
    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
    for (int64_t i = 0; i < a.numel(); ++i) a.set_item(i, dist(rng));
    for (int64_t i = 0; i < b.numel(); ++i) b.set_item(i, dist(rng));

    Device& dev = Device::get();
    MTL::ComputePipelineState* pso = dev.pipeline(kernel);
    Params p{uint32_t(M), uint32_t(N), uint32_t(K)};

    auto encode = [&]() {
        MTL::ComputeCommandEncoder* enc = metal::Stream::get().encoder();
        enc->setComputePipelineState(pso);
        enc->setBuffer(a.buffer(), a.buffer_offset(), 0);
        enc->setBuffer(b.buffer(), b.buffer_offset(), 1);
        enc->setBuffer(c.buffer(), c.buffer_offset(), 2);
        enc->setBytes(&p, sizeof(p), 3);
        enc->dispatchThreadgroups(
            MTL::Size(NS::UInteger(N) / 64, NS::UInteger(M) / 64, 1),
            MTL::Size(128, 1, 1));
    };

    encode();
    metal::sync(); // warm up, build pipeline
    const double t0 = metal::Stream::get().gpu_seconds();
    for (int i = 0; i < iters; ++i) encode();
    metal::sync();
    const double dt_s = metal::Stream::get().gpu_seconds() - t0;
    return 2.0 * double(M) * double(N) * double(K) * iters / dt_s / 1e12;
}

double run_mpp(const char* kernel, DType dt, int64_t M, int64_t N, int64_t K,
               int iters) {
    Tensor a = Tensor::empty({M, K}, dt);
    Tensor b = Tensor::empty({K, N}, dt);
    Tensor c = Tensor::empty({M, N}, DType::F32);
    std::mt19937 rng(3);
    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
    for (int64_t i = 0; i < a.numel(); ++i) a.set_item(i, dist(rng));
    for (int64_t i = 0; i < b.numel(); ++i) b.set_item(i, dist(rng));

    Device& dev = Device::get();
    MTL::ComputePipelineState* pso = dev.pipeline(kernel);
    Params p{uint32_t(M), uint32_t(N), uint32_t(K)};

    auto encode = [&]() {
        MTL::ComputeCommandEncoder* enc = metal::Stream::get().encoder();
        enc->setComputePipelineState(pso);
        enc->setBuffer(a.buffer(), a.buffer_offset(), 0);
        enc->setBuffer(b.buffer(), b.buffer_offset(), 1);
        enc->setBuffer(c.buffer(), c.buffer_offset(), 2);
        enc->setBytes(&p, sizeof(p), 3);
        enc->dispatchThreadgroups(
            MTL::Size(NS::UInteger(N) / 32, NS::UInteger(M) / 64, 1),
            MTL::Size(128, 1, 1));
    };

    encode();
    metal::sync();
    const double t0 = metal::Stream::get().gpu_seconds();
    for (int i = 0; i < iters; ++i) encode();
    metal::sync();
    const double dt_s = metal::Stream::get().gpu_seconds() - t0;
    return 2.0 * double(M) * double(N) * double(K) * iters / dt_s / 1e12;
}

} // namespace

int main() {
    NS::AutoreleasePool* pool = NS::AutoreleasePool::alloc()->init();
    std::printf("device: %s\n\n", Device::get().name().c_str());
    std::printf("simdgroup_matrix throughput by operand precision "
                "(f32 accumulator throughout)\n");
    std::printf("%-16s %10s %10s %10s\n", "shape", "f32", "f16", "bf16");

    struct Shape { int64_t M, N, K; };
    const Shape shapes[] = {{2048, 2048, 2048}, {4096, 4096, 4096},
                            {65536, 512, 1408}};
    for (const Shape& s : shapes) {
        const double gflop = 2.0 * double(s.M) * double(s.N) * double(s.K) / 1e9;
        const int iters = gflop > 50.0 ? 3 : 10;
        const double f32 = run("gemm_f32", DType::F32, s.M, s.N, s.K, iters);
        const double f16 = run("gemm_f16", DType::F16, s.M, s.N, s.K, iters);
        const double bf16 = run("gemm_bf16", DType::BF16, s.M, s.N, s.K, iters);
        char tag[32];
        std::snprintf(tag, sizeof(tag), "%lldx%lldx%lld", (long long)s.M,
                      (long long)s.N, (long long)s.K);
        std::printf("%-16s %9.2fT %9.2fT %9.2fT\n", tag, f32, f16, bf16);
        std::fflush(stdout);
    }

    // Metal Performance Primitives cooperative-tensor matmul2d — the Metal 4
    // path that targets M5 neural accelerators. 64x32 tile, 4 simdgroups.
    std::printf("\nMPP matmul2d (cooperative tensors, Metal 4)\n");
    std::printf("%-16s %10s %10s\n", "shape", "f32", "f16");
    for (const Shape& s : shapes) {
        if (s.M % 64 != 0 || s.N % 32 != 0) continue;
        const double gflop = 2.0 * double(s.M) * double(s.N) * double(s.K) / 1e9;
        const int iters = gflop > 50.0 ? 3 : 10;
        const double f32 = run_mpp("matmul_mpp_f32", DType::F32, s.M, s.N, s.K, iters);
        const double f16 = run_mpp("matmul_mpp_f16", DType::F16, s.M, s.N, s.K, iters);
        char tag[32];
        std::snprintf(tag, sizeof(tag), "%lldx%lldx%lld", (long long)s.M,
                      (long long)s.N, (long long)s.K);
        std::printf("%-16s %9.2fT %9.2fT\n", tag, f32, f16);
        std::fflush(stdout);
    }

    pool->drain();
    return 0;
}