File size: 1,751 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: relu dtype coverage on H100 — ours vs torch.
Upstream kernels-community/relu is fp32-only (TORCH_CHECK), so fp16/bf16/int8
are ours-vs-torch only. Working sets on H100 (50 MB L2):
4096^2: fp16/bf16 67 MB (DRAM, just past L2), int8 33.5 MB (L2-resident)
8192^2: fp16/bf16 268 MB, int8 134 MB (all DRAM-bound)
"""
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)
# 2-D shapes to match the fp32 spec's methodology (ours/torch are
# shape-agnostic grid-stride kernels, but keep all suites consistent).
SIZES = {"4096": (4096, 4096), "8192": (8192, 8192)}
DTYPES = {"fp16": torch.float16, "bf16": torch.bfloat16, "int8": torch.int8}
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)
def correctness():
for shape in list(SIZES.values()) + [(4097,), (333,), (1,)]:
for dt in DTYPES.values():
x = make(shape, dt)
assert torch.equal(ours.relu(x), F.relu(x)), f"{dt} @ {shape}"
cases = []
for tag, shape in SIZES.items():
for dname, dt in DTYPES.items():
x = make(shape, dt)
out = torch.empty_like(x)
cases.append(Case(f"torch_{dname}_{tag}",
lambda x=x, out=out: torch.clamp(x, min=0, out=out)))
cases.append(Case(f"ours_{dname}_{tag}",
lambda x=x, out=out: ours.relu(x, out=out)))
run(cases, correctness_fn=correctness)
|