"""robust_bench spec: small-working-set dtype check for the Hopper override.

1024^2 in all four dtypes exercises the ws<16MB u2/st.cs branch (fp32 8.4MB,
fp16/bf16 4.2MB, int8 2.1MB); int8 4096^2 (33.5MB) confirms the reverted
33MB band returns to baseline. ours vs torch, interleaved 1:1.
"""
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)


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)


DTYPES = {"fp32": torch.float32, "fp16": torch.float16,
          "bf16": torch.bfloat16, "int8": torch.int8}


def correctness():
    for shape in [(1024, 1024), (4096, 4096), (4097,), (333,), (1,)]:
        for dt in DTYPES.values():
            x = make(shape, dt)
            assert torch.equal(ours.relu(x), F.relu(x)), f"{dt} @ {shape}"
    x = torch.full((4096,), float("nan"), device="cuda")
    assert (ours.relu(x) == 0).all(), "NaN -> 0 expected"


cases = []
for dname, dt in DTYPES.items():
    x = make((1024, 1024), dt)
    out = torch.empty_like(x)
    cases.append(Case(f"torch_{dname}_1024",
                      lambda x=x, out=out: torch.clamp(x, min=0, out=out)))
    cases.append(Case(f"ours_{dname}_1024",
                      lambda x=x, out=out: ours.relu(x, out=out)))
xi = make((4096, 4096), torch.int8)
oi = torch.empty_like(xi)
cases.append(Case("torch_int8_4096",
                  lambda x=xi, out=oi: torch.clamp(x, min=0, out=out)))
cases.append(Case("ours_int8_4096",
                  lambda x=xi, out=oi: ours.relu(x, out=out)))

run(cases, correctness_fn=correctness)