"""Roofline floor for a task on a GPU: t_min = max(FLOPs / peak compute, bytes / peak bandwidth).

FLOPs come from torch.utils.flop_counter on the reference (matmul/conv/attention only, so
elementwise work counts as 0) and bytes are the compulsory traffic: inputs + parameters read once,
outputs written once. Both are lower bounds, so t_min is a floor no correct kernel can beat,
except tasks that legitimately skip input data (slicing), which KernelBench rarely has.
"""

from typing import Dict, Optional

# Dense datasheet peaks: (FP32 SIMT, TF32 tensor, BF16/FP16 tensor) in TFLOP/s, HBM in TB/s.
PEAKS = {
    "H100 PCIe": (51.2, 378.0, 756.0, 2.0),
    "H100": (66.9, 494.7, 989.4, 3.35),  # SXM, "NVIDIA H100 80GB HBM3"
    "H200": (66.9, 494.7, 989.4, 4.8),
    "A100-SXM4-40GB": (19.5, 156.0, 312.0, 1.555),
    "A100": (19.5, 156.0, 312.0, 2.039),
    "L40S": (91.6, 183.0, 362.0, 0.864),
    "L4": (30.3, 60.0, 121.0, 0.3),
}


def peaks(gpu: Optional[str]):
    for key in sorted(PEAKS, key=len, reverse=True):  # "L40S" before "L4", "H100 PCIe" before "H100"
        if gpu and key in gpu:
            return PEAKS[key]
    return None


def roofline(flops, nbytes, dtype, gpu, t_ref_ms, t_new_ms=None) -> Optional[Dict]:
    p = peaks(gpu)
    if p is None or not nbytes or flops is None or not t_ref_ms:
        return None
    _, tf32, half, bw = p
    # fp32 tasks: TF32 tensor cores are the roof, since tl.dot uses TF32 by default and passes atol=rtol=1e-2
    compute = {"torch.float32": tf32, "torch.float16": half, "torch.bfloat16": half}.get(dtype)
    if compute is None:
        return None
    t_compute = flops / (compute * 1e12) * 1e3
    t_memory = nbytes / (bw * 1e12) * 1e3
    t_min = max(t_compute, t_memory)
    r = {
        "arithmetic_intensity": flops / nbytes,
        "bound": "compute" if t_compute > t_memory else "memory",
        "t_min_ms": t_min,
        "ref_roofline_frac": t_min / t_ref_ms,
    }
    if t_new_ms:
        r["roofline_frac"] = t_min / t_new_ms
    return r


def describe(r: Dict, gpu: str) -> str:
    s = (f"Roofline on {gpu}: {r['bound']}-bound (arithmetic intensity {r['arithmetic_intensity']:.3g} FLOP/byte), "
         f"floor {r['t_min_ms']:.4f} ms; PyTorch reaches {r['ref_roofline_frac']:.0%} of it")
    if "roofline_frac" in r:
        s += f", yours {r['roofline_frac']:.0%}"
    return s + "."