"""robust_bench spec: fp32 relu three-way on H100 — ours vs upstream Hub vs torch.

Sizes chosen for H100 (50 MB L2, HBM3):
  1024^2 fp32:  8 MB working set  -> L2-resident
  4096^2 fp32: 134 MB             -> DRAM-bound
  8192^2 fp32: 536 MB             -> DRAM-bound, far beyond cache
"""
import inspect
import sys
from pathlib import Path

import torch
import torch.nn.functional as F

BENCH = Path.home() / "relu-bench"
sys.path.insert(0, str(BENCH))
sys.path.insert(0, str(BENCH / "torch212-cxx11-cu126-x86_64-linux"))

import relu as ours  # our optimized bundle (sm_90 SASS verified present)
from robust_bench import Case, run
from kernels import get_kernel

hub = get_kernel("kernels-community/relu", version=1)
HUB_HAS_OUT = "out" in inspect.signature(hub.relu).parameters

torch.manual_seed(0)
# 2-D shapes: upstream hub's launch heuristic is shape-dependent
# (grid = numel/size(-1)); a flat 1-D tensor collapses it to grid=1.
# 2-D matches how the 4090 comparison was run (1024^2 / 4096^2 matrices).
SIZES = {"1024": 1024, "4096": 4096, "8192": 8192}


def correctness():
    for tag, s in SIZES.items():
        x = torch.randn(s, s, device="cuda")
        r_ours = ours.relu(x)
        assert torch.equal(r_ours, F.relu(x)), f"ours != F.relu @ {tag}"
        assert torch.equal(r_ours, hub.relu(x)), f"ours != hub @ {tag}"
    # tail/odd sizes
    for n in (4097, 333, 1):
        x = torch.randn(n, device="cuda")
        assert torch.equal(ours.relu(x), F.relu(x)), f"tail {n}"
    # NaN semantics: ours and hub both map NaN -> 0 (upstream semantic)
    x = torch.full((4096,), float("nan"), device="cuda")
    assert torch.equal(ours.relu(x), hub.relu(x)), "NaN semantic mismatch"
    assert (ours.relu(x) == 0).all(), "NaN -> 0 expected"


cases = []
for tag, s in SIZES.items():
    x = torch.randn(s, s, device="cuda")
    out = torch.empty_like(x)
    cases.append(Case(f"torch_{tag}",
                      lambda x=x, out=out: torch.clamp(x, min=0, out=out)))
    if HUB_HAS_OUT:
        cases.append(Case(f"hub_{tag}",
                          lambda x=x, out=out: hub.relu(x, out=out)))
    else:
        cases.append(Case(f"hub_{tag}", lambda x=x: hub.relu(x)))
    cases.append(Case(f"ours_{tag}",
                      lambda x=x, out=out: ours.relu(x, out=out)))

run(cases, correctness_fn=correctness)