File size: 1,884 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 53 54 55 56 57 58 | """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)
|