# Copyright 2026 FlagOS Contributors / GuanghuLab
"""Benchmark gelu_and_mul v2 vs FlagGems baseline.

Usage:
    python -m bench.bench_gelu_and_mul --shape 4096,4096 --dtype fp16
    python -m bench.bench_gelu_and_mul --shape 8192,8192 --dtype bf16
"""

import argparse
import time
import torch

from flag_gems.fused.gelu_and_mul import gelu_and_mul as baseline_fn
from flag_gems_local.fused.gelu_and_mul_v2 import gelu_and_mul as v2_fn


def _to_ms(t):
    return t * 1000.0


def bench(fn, x, y, iters=100, warmup=20):
    # warmup
    for _ in range(warmup):
        out = fn(x, y)
    torch.cuda.synchronize()

    # measure
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    for _ in range(iters):
        out = fn(x, y)
    end.record()
    torch.cuda.synchronize()
    return start.elapsed_time(end) / iters


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--shape", type=str, default="4096,4096")
    parser.add_argument("--dtype", type=str, default="fp16",
                        choices=["fp16", "bf16", "fp32"])
    parser.add_argument("--iters", type=int, default=100)
    args = parser.parse_args()

    dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[args.dtype]
    shape = tuple(int(s) for s in args.shape.split(","))

    print(f"Benchmark gelu_and_mul: shape={shape} dtype={args.dtype} iters={args.iters}")
    print("=" * 72)

    x = torch.randn(shape, dtype=dtype, device="cuda")
    y = torch.randn(shape, dtype=dtype, device="cuda")

    # Correctness check
    out_b = baseline_fn(x, y)
    out_v = v2_fn(x, y)
    abs_diff = (out_b - out_v).abs().max().item()
    rel_diff = (out_b - out_v).abs().div(out_b.abs().clamp_min(1e-6)).max().item()
    print(f"Correctness: max abs_diff={abs_diff:.2e}, max rel_diff={rel_diff:.2e}")
    assert abs_diff < 1e-2, "v2 diverges from baseline"
    print()

    t_base = bench(baseline_fn, x, y, iters=args.iters)
    t_v2 = bench(v2_fn, x, y, iters=args.iters)
    speedup = t_base / t_v2
    print(f"Baseline : {_to_ms(t_base):.4f} ms/iter")
    print(f"v2 (ours): {_to_ms(t_v2):.4f} ms/iter")
    print(f"Speedup  : {speedup:.3f}x")
    print()

    # Find best autotune config
    print("Best v2 autotune config:")
    for k, v in v2_fn.__self__.forward.__func__.__code__.co_consts:
        if k == "BLOCK_SIZE":
            print(f"  BLOCK_SIZE = {v}")


if __name__ == "__main__":
    main()