"""Benchmark silu_and_mul_masked v2 vs FlagGems baseline.

Run: python bench/bench_silu_and_mul_masked.py --shape 4096,4096 --dtype fp16
"""

import argparse

import torch


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

    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    for _ in range(iters):
        out = fn(*args)
    end.record()
    torch.cuda.synchronize()
    return start.elapsed_time(end) / iters


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--shape", default="4096,4096")
    parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16", "fp32"])
    parser.add_argument("--iters", type=int, default=100)
    parser.add_argument("--with-mask", action="store_true", help="enable mask")
    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"Bench silu_and_mul_masked: shape={shape} dtype={args.dtype} mask={args.with_mask}")
    print("=" * 72)

    x = torch.randn(shape, dtype=dtype, device="cuda")
    y = torch.randn(shape, dtype=dtype, device="cuda")
    mask = torch.randn(shape, dtype=dtype, device="cuda") if args.with_mask else None

    from flag_gems.fused.silu_and_mul import silu_and_mul as baseline_fn
    from flag_gems_local.fused.silu_and_mul_masked_v2 import silu_and_mul_masked as v2_fn

    # Reference (using PyTorch native ops)
    def ref_fn(x, y, mask=None):
        x_in = x * mask if mask is not None else x
        return torch.nn.functional.silu(x_in) * y

    out_r = ref_fn(x, y, mask)
    if mask is None:
        out_b = baseline_fn(x, y)
    else:
        # baseline doesn't support mask; compare v2 to reference
        out_b = None
    out_v = v2_fn(x, y, mask)

    if out_b is not None:
        abs_diff = (out_b - out_v).abs().max().item()
        print(f"v2 vs baseline: abs_diff={abs_diff:.2e}")
    abs_diff_r = (out_r - out_v).abs().max().item()
    print(f"v2 vs reference: abs_diff={abs_diff_r:.2e}")
    assert abs_diff_r < 1e-2, "v2 diverges from reference"
    print()

    # Benchmark
    if mask is None:
        t_base = bench(baseline_fn, x, y, iters=args.iters)
        print(f"Baseline : {t_base:.4f} ms/iter")
    t_v2 = bench(v2_fn, x, y, mask, iters=args.iters)
    print(f"v2 (ours): {t_v2:.4f} ms/iter")
    if mask is None:
        speedup = t_base / t_v2
        print(f"Speedup  : {speedup:.3f}x")
    else:
        print("Baseline doesn't support mask, no speedup comparison.")


if __name__ == "__main__":
    main()