# SPDX-License-Identifier: CC-BY-NC-4.0
# Copyright (c) 2026 Jakub Ćwirlej
"""FFN block micro-benchmark (CUDA): dense SwiGLU vs eugen_sparse (selection + kernel), forward and
forward+backward, bf16, optionally under torch.compile.

    python speed_check.py --tokens 65536 --d-model 512 --d-ff 4096 --group-size 32 \
        --active-groups 8 --samples-per-group 8 [--compile]

This times one FFN block on random inputs. It is not a prediction of end-to-end training speed:
measure your own training loop. Random inputs spread tokens evenly over the groups; a trained model
is more skewed (popular groups), which the load-balancing bias counteracts.
"""
import argparse
import statistics
import time

import torch
import torch.nn.functional as F

import eugen_sparse as es


def timer(fn, warmup=5, iters=20):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    ts = []
    for _ in range(iters):
        t0 = time.perf_counter()
        fn()
        torch.cuda.synchronize()
        ts.append((time.perf_counter() - t0) * 1000)
    return statistics.median(ts)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--tokens', type=int, default=65536)
    ap.add_argument('--d-model', type=int, default=512)
    ap.add_argument('--d-ff', type=int, default=4096)
    ap.add_argument('--group-size', type=int, default=32)
    ap.add_argument('--active-groups', type=int, default=8)
    ap.add_argument('--samples-per-group', type=int, default=8)
    ap.add_argument('--compile', action='store_true')
    a = ap.parse_args()
    assert torch.cuda.is_available(), 'CUDA required'
    G, kg = a.group_size, a.active_groups
    if not es.kernel_supported(True, a.d_model, a.d_ff, G, kg):
        raise SystemExit('shape not supported by the kernel (see kernel_supported)')
    dt = torch.bfloat16
    torch.manual_seed(1)
    x = torch.randn(a.tokens, a.d_model, device='cuda', dtype=dt)
    gw, uw = (torch.randn(a.d_ff, a.d_model, device='cuda', dtype=dt) * a.d_model ** -0.5 for _ in range(2))
    dw = torch.randn(a.d_model, a.d_ff, device='cuda', dtype=dt) * a.d_ff ** -0.5
    dy = torch.randn_like(x)
    sub, spg = es.make_subsample_index(a.d_ff, G, a.samples_per_group)
    sub = sub.cuda()

    def dense(x_, g_, u_, d_):
        return F.linear(F.silu(F.linear(x_, g_)) * F.linear(x_, u_), d_)

    def sparse(x_, g_, u_, d_):
        idx = es.grouped_topk_select(x_, g_, G, kg, sub, spg)
        return es.sparse_swiglu_ffn(x_, g_, u_, d_, idx, G)

    print(f'tokens={a.tokens} d_model={a.d_model} d_ff={a.d_ff} group_size={G} active_groups={kg} '
          f'(density {kg * G / a.d_ff:.2%}) compile={a.compile} | {torch.cuda.get_device_name(0)} | '
          f'torch {torch.__version__}')
    print('plan:', es.plan(a.tokens, a.d_model, a.d_ff, G, kg))
    res = {}
    for name, fn in (('dense', dense), ('sparse', sparse)):
        if a.compile:
            fn = torch.compile(fn, fullgraph=True, dynamic=False)
        t_f = timer(lambda: fn(x, gw, uw, dw))
        leaves = tuple(t.clone().requires_grad_() for t in (x, gw, uw, dw))

        def fb(fn=fn, l=leaves):
            torch.autograd.grad(fn(*l), l, dy)
        t_fb = timer(fb)
        res[name] = (t_f, t_fb)
        print(f'  {name:6s} fwd {t_f:8.2f} ms   fwd+bwd {t_fb:8.2f} ms')
    print(f'  sparse vs dense: fwd x{res["dense"][0] / res["sparse"][0]:.2f}   '
          f'fwd+bwd x{res["dense"][1] / res["sparse"][1]:.2f}  (>1 = sparse faster)')


if __name__ == '__main__':
    main()