"""robust_bench spec: old-vs-new relu at the sizes the H100 retune targeted.

2048^2 fp32 = 33.6 MB rw (deep-L2 regime), 2560^2 fp32 = 52.4 MB rw (the L2
boundary where the sweep showed +20%). Old and new .so load side by side
(distinct torch.ops namespaces), interleaved 1:1 in one process.
"""
import importlib.util
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 new_relu
from robust_bench import Case, run
from kernels import get_kernel

OLD = BENCH / "old_bundle"
spec = importlib.util.spec_from_file_location(
    "relu_old", OLD / "__init__.py", submodule_search_locations=[str(OLD)])
old_relu = importlib.util.module_from_spec(spec)
sys.modules["relu_old"] = old_relu
spec.loader.exec_module(old_relu)

hub = get_kernel("kernels-community/relu", version=1)

torch.manual_seed(0)
SIZES = {"2048": 2048, "2560": 2560}


def correctness():
    for s in SIZES.values():
        x = torch.randn(s, s, device="cuda")
        r_new = new_relu.relu(x)
        assert torch.equal(r_new, F.relu(x)), f"new != F.relu @ {s}"
        assert torch.equal(r_new, old_relu.relu(x)), f"new != old @ {s}"
        assert torch.equal(r_new, hub.relu(x)), f"new != hub @ {s}"


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)))
    cases.append(Case(f"hub_{tag}", lambda x=x: hub.relu(x)))
    cases.append(Case(f"old_{tag}",
                      lambda x=x, out=out: old_relu.relu(x, out=out)))
    cases.append(Case(f"new_{tag}",
                      lambda x=x, out=out: new_relu.relu(x, out=out)))

run(cases, correctness_fn=correctness)