bench/spec_pair_h100.py

"""robust_bench spec: symmetric pairwise old-vs-new relu A/B, one size per run.

Two cases only, own buffers each, strict alternation: old's predecessor is
always new and vice versa, so streaming-hint L2 aftermath cancels between
them (the 4-case spec biased whoever ran after the evict-first kernel).
Size selected via SPEC_SIZE env (matrix side, e.g. 1024/2048/2560).
"""
import os
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

import importlib.util
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)

S = int(os.environ["SPEC_SIZE"])
torch.manual_seed(0)

x_old = torch.randn(S, S, device="cuda")
x_new = x_old.clone()
out_old = torch.empty_like(x_old)
out_new = torch.empty_like(x_new)


def correctness():
    r_new = new_relu.relu(x_new)
    assert torch.equal(r_new, F.relu(x_new)), "new != F.relu"
    assert torch.equal(r_new, old_relu.relu(x_old)), "new != old"


run(
    [
        Case(f"old_{S}", lambda: old_relu.relu(x_old, out=out_old)),
        Case(f"new_{S}", lambda: new_relu.relu(x_new, out=out_new)),
    ],
    correctness_fn=correctness,
)