File size: 1,514 Bytes
a6e5c43 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | """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,
)
|