"""robust_bench spec: relu dtype coverage on H100 — ours vs torch.

Upstream kernels-community/relu is fp32-only (TORCH_CHECK), so fp16/bf16/int8
are ours-vs-torch only. Working sets on H100 (50 MB L2):
  4096^2: fp16/bf16 67 MB (DRAM, just past L2), int8 33.5 MB (L2-resident)
  8192^2: fp16/bf16 268 MB, int8 134 MB (all DRAM-bound)
"""
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
from robust_bench import Case, run

torch.manual_seed(0)
# 2-D shapes to match the fp32 spec's methodology (ours/torch are
# shape-agnostic grid-stride kernels, but keep all suites consistent).
SIZES = {"4096": (4096, 4096), "8192": (8192, 8192)}
DTYPES = {"fp16": torch.float16, "bf16": torch.bfloat16, "int8": torch.int8}


def make(shape, dt):
    if dt == torch.int8:
        return torch.randint(-128, 128, shape, dtype=torch.int8, device="cuda")
    return torch.randn(*shape, device="cuda", dtype=dt)


def correctness():
    for shape in list(SIZES.values()) + [(4097,), (333,), (1,)]:
        for dt in DTYPES.values():
            x = make(shape, dt)
            assert torch.equal(ours.relu(x), F.relu(x)), f"{dt} @ {shape}"


cases = []
for tag, shape in SIZES.items():
    for dname, dt in DTYPES.items():
        x = make(shape, dt)
        out = torch.empty_like(x)
        cases.append(Case(f"torch_{dname}_{tag}",
                          lambda x=x, out=out: torch.clamp(x, min=0, out=out)))
        cases.append(Case(f"ours_{dname}_{tag}",
                          lambda x=x, out=out: ours.relu(x, out=out)))

run(cases, correctness_fn=correctness)