metakernel v1: device dossier, throttle-rejected bench, ULP compare, fuzz, sweep, stamps
0c6aadc verified | import math | |
| import pytest | |
| import torch | |
| import kernels | |
| mk = kernels.get_kernel("phanerozoic/metakernel", version=1, | |
| trust_remote_code=True) | |
| requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), | |
| reason="CUDA required") | |
| _dossier = None | |
| def dossier(): | |
| global _dossier | |
| if _dossier is None: | |
| _dossier = mk.probe_device(quick=True) | |
| return _dossier | |
| def test_clock_measurement_sane(): | |
| ghz = mk.measure_clock() | |
| assert 0.3 < ghz < 4.5, ghz | |
| def test_dossier_bandwidth_hierarchy(): | |
| """Measured L2 bandwidth must exceed the best HBM number; streaming | |
| patterns must land in plausible ranges; random gather must cost far | |
| more than streaming.""" | |
| d = dossier() | |
| hbm = d["hbm_bw"]["value"] | |
| best_stream = max(hbm["triad_4B"], hbm["triad_8B"], hbm["triad_16B"], | |
| hbm["read_only"], hbm["write_only"]) | |
| assert 50 < best_stream < 20000, hbm | |
| assert hbm["triad_16B"] >= hbm["triad_4B"] * 0.8, hbm | |
| assert hbm["read_only"] > 50 and hbm["write_only"] > 50, hbm | |
| assert hbm["random_gather_16B"] < best_stream, hbm | |
| assert d["l2_bw"]["value"] > best_stream, (d["l2_bw"]["value"], | |
| best_stream) | |
| def test_dossier_latency_hierarchy(): | |
| """Dependent-load latency must order smem < L2 < HBM with the L1 | |
| point between cache levels.""" | |
| lat = dossier()["latency"]["value"] | |
| assert lat["smem_cycles"] < lat["l2_cycles"] < lat["hbm_cycles"], lat | |
| assert lat["l1_cycles"] < lat["l2_cycles"], lat | |
| assert 10 < lat["l1_cycles"] < 300, lat | |
| assert 5 < lat["smem_cycles"] < 150, lat | |
| assert 50 < lat["hbm_cycles"] < 3000, lat | |
| def test_dossier_mma_rates(): | |
| """Tensor cores beat the FMA pipe; int8 is not slower than fp16; fp8 | |
| support matches the architecture.""" | |
| d = dossier() | |
| mma = d["mma"]["value"] | |
| fma = d["fma_f32"]["value"] | |
| assert isinstance(mma["fp16"], float) and mma["fp16"] > fma, (mma, fma) | |
| assert isinstance(mma["bf16"], float) and mma["bf16"] > fma | |
| assert isinstance(mma["tf32"], float) | |
| assert isinstance(mma["int8"], float) and mma["int8"] > 0.9 * mma["fp16"] | |
| major = int(d["device"]["sm_arch"][2]) | |
| minor = int(d["device"]["sm_arch"][3:]) | |
| if major * 10 + minor >= 89: | |
| assert isinstance(mma["fp8_e4m3"], float), mma | |
| else: | |
| assert mma["fp8_e4m3"] == "unsupported", mma | |
| def test_dossier_atomics_contention_curve(): | |
| """Throughput must not fall as the number of distinct addresses grows, | |
| and the endpoints must differ by a large factor.""" | |
| at = dossier()["atomics_f32"]["value"] | |
| curve = at["curve_slots"] | |
| slots = sorted(int(k) for k in curve) | |
| vals = [curve[str(s)] for s in slots] | |
| for lo, hi in zip(vals, vals[1:]): | |
| assert hi >= lo * 0.7, curve # monotone up to measurement noise | |
| assert vals[-1] > vals[0] * 4, curve | |
| assert at["shared_same_addr"] > 0, at | |
| def test_dossier_fp64_pipe(): | |
| d = dossier() | |
| assert 0.02 < d["fma_f64"]["value"] < d["fma_f32"]["value"] * 0.6, \ | |
| (d["fma_f64"]["value"], d["fma_f32"]["value"]) | |
| def test_dossier_occupancy_knees_timed(): | |
| """The timed counterpart to the API table: a dependent-chain workload | |
| must lose throughput when dynamic smem forces occupancy down.""" | |
| knees = dossier()["occupancy_knees_fma_tflops"]["value"] | |
| healthy = knees["t256"]["0KB"] | |
| starved = knees["t128"]["48KB"] # 2 blocks x 128 threads per SM | |
| assert starved < healthy * 0.8, knees | |
| assert knees["t1024"]["48KB"] >= starved, knees | |
| def test_dossier_launch_and_barrier(): | |
| d = dossier() | |
| lo = d["launch_overhead"]["value"] | |
| assert 0.2 < lo["queued_us"] < 500, lo | |
| assert lo["round_trip_us"] > lo["queued_us"] * 0.5, lo | |
| gb = d["grid_barrier"]["value"] | |
| assert gb == "unsupported" or 0.05 < gb < 500, gb | |
| def test_dossier_occupancy_monotone(): | |
| occ = dossier()["occupancy_blocks_per_sm"]["value"] | |
| assert occ["t128"]["0KB"] >= occ["t1024"]["0KB"], occ | |
| assert occ["t256"]["48KB"] <= occ["t256"]["0KB"], occ | |
| assert occ["t256"]["0KB"] >= 1, occ | |
| def test_bench_reports_and_roofline(): | |
| d = dossier() | |
| n = 4096 | |
| x = torch.randn(n, n, device="cuda", dtype=torch.float16) | |
| y = torch.randn(n, n, device="cuda", dtype=torch.float16) | |
| flops = 2 * n ** 3 | |
| byts = 3 * n * n * 2 | |
| rep = mk.bench(lambda: x @ y, iters=48, warmup=8, bytes=byts, | |
| flops=flops, dossier=d, dtype="fp16") | |
| assert rep["median_ms"] > 0 | |
| assert rep["samples"] == 48 | |
| assert rep["rejected_throttled"] <= rep["samples"] | |
| assert 1 < rep["achieved_tflops"] < 5000, rep | |
| assert rep["roofline"]["bound"] == "compute", rep["roofline"] | |
| big = torch.randn(64 << 20, device="cuda") | |
| dst = torch.empty_like(big) | |
| rep2 = mk.bench(lambda: dst.copy_(big), iters=32, warmup=4, | |
| bytes=2 * big.numel() * 4, flops=big.numel(), | |
| dossier=d) | |
| assert rep2["roofline"]["bound"] == "memory", rep2["roofline"] | |
| assert rep2["achieved_gbs"] > 20, rep2 | |
| def test_bench_graph_mode_strips_dispatch(): | |
| """Graph replay removes Python and dispatch overhead: for a train of | |
| tiny kernels the captured timing must not exceed the eager timing.""" | |
| x = torch.randn(1024, device="cuda") | |
| def tiny(): | |
| for _ in range(16): | |
| x.mul_(1.0000001) | |
| eager = mk.bench(tiny, iters=32, warmup=8) | |
| graphed = mk.bench(tiny, iters=32, warmup=8, graph=True) | |
| assert graphed["graph_captured"] and not eager["graph_captured"] | |
| assert graphed["median_ms"] <= eager["median_ms"] * 1.10, \ | |
| (graphed["median_ms"], eager["median_ms"]) | |
| def test_compare_multi_output_and_tolerance_route(): | |
| def gen(seed): | |
| g = torch.Generator(device="cuda").manual_seed(seed) | |
| return (torch.randn(64, 64, device="cuda", generator=g),) | |
| ref = lambda x: (x * 2.0, x.sum(dim=1)) # noqa: E731 | |
| same = mk.compare(ref, ref, gen, trials=2) | |
| assert same["pass"] and same["worst_trial"]["max_ulp"] == 0 | |
| shifted = lambda x: (x * 2.0 + 1e-6, x.sum(dim=1)) # noqa: E731 | |
| ulp_fail = mk.compare(ref, shifted, gen, trials=2) | |
| assert not ulp_fail["pass"] | |
| tol_pass = mk.compare(ref, shifted, gen, trials=2, atol=1e-4) | |
| assert tol_pass["pass"], tol_pass["criterion"] | |
| assert tol_pass["max_abs"] <= 2e-6, tol_pass["max_abs"] | |
| def test_compare_builtin_matmul_k_shuffle(): | |
| """The built-in K-permutation helper: a matmul judged against its own | |
| accumulation-order band.""" | |
| def gen(seed): | |
| g = torch.Generator(device="cuda").manual_seed(seed) | |
| a = torch.randn(96, 4096, device="cuda", generator=g, | |
| dtype=torch.float16) | |
| b = torch.randn(4096, 64, device="cuda", generator=g, | |
| dtype=torch.float16) | |
| return (a, b) | |
| ref = lambda a, b: a @ b # noqa: E731 | |
| rep = mk.compare(ref, ref, gen, trials=2, | |
| order_shuffle=mk.order_shuffles.matmul_k(), | |
| order_trials=4) | |
| assert rep["pass"], rep["criterion"] | |
| bad = lambda a, b: (a @ b) * 1.01 # noqa: E731 | |
| rep2 = mk.compare(ref, bad, gen, trials=2, | |
| order_shuffle=mk.order_shuffles.matmul_k(), | |
| order_trials=4) | |
| assert not rep2["pass"], rep2["criterion"] | |
| def test_ulp_diff_unit_cases(): | |
| a = torch.tensor([1.0, -0.0, float("nan"), 2.0], device="cuda") | |
| b = torch.tensor([torch.nextafter(torch.tensor(1.0), | |
| torch.tensor(2.0)).item(), | |
| 0.0, float("nan"), 2.0], device="cuda") | |
| d = mk.ulp_diff(a, b) | |
| assert d[0].item() == 1 | |
| assert d[1].item() == 0 # -0.0 and +0.0 are the same value | |
| assert d[2].item() == 0 # both-NaN agree | |
| assert d[3].item() == 0 | |
| c = torch.tensor([float("nan")], device="cuda") | |
| e = torch.tensor([1.0], device="cuda") | |
| assert mk.ulp_diff(c, e)[0].item() == 1 << 62 # one-sided NaN | |
| h = torch.tensor([1.0], dtype=torch.float16, device="cuda") | |
| h2 = torch.nextafter(h, torch.tensor([2.0], dtype=torch.float16, | |
| device="cuda")) | |
| assert mk.ulp_diff(h, h2)[0].item() == 1 | |
| def test_compare_pass_and_fail(): | |
| def gen(seed): | |
| g = torch.Generator(device="cuda").manual_seed(seed) | |
| return (torch.randn(256, 512, device="cuda", generator=g),) | |
| ref = lambda x: torch.relu(x) * 2.0 # noqa: E731 | |
| same = mk.compare(ref, ref, gen, trials=3) | |
| assert same["pass"] and same["worst_trial"]["max_ulp"] == 0 | |
| close = mk.compare(ref, lambda x: (torch.relu(x) * 2.0), gen, trials=3) | |
| assert close["pass"] | |
| broken = mk.compare(ref, lambda x: torch.relu(x) * 2.0 + 1e-3, gen, | |
| trials=3) | |
| assert not broken["pass"] | |
| assert broken["worst_trial"]["max_ulp"] > 100 | |
| assert broken["worst_trial"]["first_divergent_index"] >= 0 | |
| assert len(broken["ulp_hist_log2"]) > 0 | |
| def test_compare_reduction_order_band(): | |
| """A long fp32 sum judged against its own accumulation-order spread: | |
| the shuffled reference bounds the legitimate band, and a candidate | |
| that IS a reordered sum passes while a corrupted one fails.""" | |
| def gen(seed): | |
| g = torch.Generator(device="cuda").manual_seed(seed) | |
| return (torch.randn(8, 200_000, device="cuda", generator=g) | |
| * 100.0,) | |
| def shuffle(args, seed): | |
| g = torch.Generator(device="cuda").manual_seed(seed) | |
| p = torch.randperm(args[0].shape[1], device="cuda", generator=g) | |
| return (args[0][:, p],) | |
| ref = lambda x: x.sum(dim=1) # noqa: E731 | |
| cand = lambda x: x.flip(1).sum(dim=1) # noqa: E731 | |
| rep = mk.compare(ref, cand, gen, trials=3, order_shuffle=shuffle, | |
| order_trials=6) | |
| assert rep["pass"], rep["worst_trial"] | |
| bad = lambda x: x.sum(dim=1) * (1 + 1e-4) # noqa: E731 | |
| rep2 = mk.compare(ref, bad, gen, trials=3, order_shuffle=shuffle, | |
| order_trials=6) | |
| assert not rep2["pass"], rep2["worst_trial"] | |
| def test_fuzz_shapes_envelope(): | |
| def make(shape, dtype=torch.float32): | |
| n = math.prod(shape) | |
| return torch.randn(*shape, device="cuda", dtype=dtype) if n \ | |
| else torch.zeros(*shape, device="cuda", dtype=dtype) | |
| ref = lambda x: torch.relu(x.contiguous()) # noqa: E731 | |
| good = mk.fuzz_shapes(torch.relu, ref, make, | |
| variants=("contig", "transposed", "strided_rows", | |
| "offset_slice")) | |
| assert good["failed"] == 0, good | |
| assert not good["context_lost"] | |
| # an op that silently assumes contiguity must be caught | |
| cheat = lambda x: torch.relu( # noqa: E731 | |
| torch.as_strided(x, x.shape, [x.shape[-1], 1] if x.dim() == 2 | |
| else [1])) | |
| bad = mk.fuzz_shapes(cheat, ref, make, | |
| variants=("contig", "transposed", "strided_rows", | |
| "offset_slice")) | |
| assert bad["failed"] > 0, bad | |
| def test_fuzz_multi_arg_dtypes_overflow(): | |
| """Binary op across dtypes with the overflow_expand variant: past | |
| 2**31 logical elements the case either runs correctly (large card) or | |
| reports skipped_vram (small card); silent absence is not allowed.""" | |
| def make(shape, dtype=torch.float16): | |
| n = math.prod(shape) | |
| f = torch.randn if n else lambda *s, **k: torch.zeros(*s, **k) | |
| return (f(*shape, device="cuda", dtype=dtype), | |
| f(*shape, device="cuda", dtype=dtype)) | |
| op = torch.add | |
| ref = lambda a, b: a.contiguous() + b.contiguous() # noqa: E731 | |
| rep = mk.fuzz_shapes(op, ref, make, shapes=[(1, 17), (128, 128)], | |
| dtypes=(torch.float16, torch.float32)) | |
| assert rep["failed"] == 0, [c for c in rep["cases"] | |
| if c["status"] not in ("ok", "skipped_vram", | |
| "skipped_empty")] | |
| ov = [c for c in rep["cases"] if c["variant"] == "overflow_expand"] | |
| assert len(ov) == 4 | |
| assert all(c["status"] in ("ok", "skipped_vram") for c in ov), ov | |
| ran = [c for c in ov if c["status"] == "ok"] | |
| for c in ran: | |
| assert c["logical_elems"] > (1 << 31), c | |
| def test_sweep_orders_configurations(): | |
| x = torch.randn(1 << 22, device="cuda") | |
| def factory(chunk): | |
| def run(): | |
| for piece in x.split(chunk): | |
| piece.mul_(1.0000001) | |
| return run | |
| rows = mk.sweep(factory, {"chunk": [1 << 14, 1 << 20, 1 << 22]}, | |
| iters=12, warmup=3) | |
| assert len(rows) == 3 | |
| assert rows[0]["median_ms"] <= rows[-1]["median_ms"] | |
| assert all("chunk" in r for r in rows) | |
| def test_sweep_check_gate_and_budget(): | |
| x = torch.randn(1 << 18, device="cuda") | |
| def factory(scale): | |
| return lambda: x.mul(scale) | |
| rows = mk.sweep(factory, {"scale": [1.0, 2.0]}, iters=6, warmup=2, | |
| check=lambda fn: bool(torch.isfinite(fn()).all())) | |
| assert all(r["check"] is True for r in rows) | |
| rows2 = mk.sweep(factory, {"scale": [1.0, 2.0, 3.0, 4.0]}, iters=6, | |
| warmup=2, budget_s=0.0) | |
| assert any(r.get("status") == "not_run" for r in rows2) | |
| assert len(rows2) == 4 # skipped points are visible, not dropped | |
| def test_stamps_per_block(): | |
| st = torch.zeros(3, 3, dtype=torch.int64) | |
| st[0] = torch.tensor([100, 300, 600]) | |
| st[1] = torch.tensor([100, 500, 800]) | |
| st[2] = torch.tensor([0, 0, 0]) # a block that never stamped | |
| rep = mk.read_stamps(st.cuda(), labels=["load", "compute"], | |
| clock_ghz=1.0) | |
| assert rep["blocks_reporting"] == 2 and rep["blocks_total"] == 3 | |
| load = rep["phases"][0] | |
| assert load["ticks_min"] == 200 and load["ticks_max"] == 400 | |
| assert load["ticks_median"] in (200, 300, 400) | |
| def test_stamps_roundtrip(): | |
| st = mk.alloc_stamps(4) | |
| base = 1_000_000 | |
| ticks = [base, base + 2_000_000, base + 3_000_000, base + 3_500_000] | |
| st.copy_(torch.tensor(ticks, dtype=torch.int64)) | |
| rep = mk.read_stamps(st, labels=["load", "compute", "store"], | |
| clock_ghz=2.0) | |
| assert [p["phase"] for p in rep["phases"]] == ["load", "compute", | |
| "store"] | |
| assert rep["phases"][0]["ticks"] == 2_000_000 | |
| assert abs(rep["phases"][0]["us"] - 1000.0) < 1e-6 | |
| assert rep["phases"][2]["ticks"] == 500_000 | |
| def test_chase_repeatability(): | |
| """The instrument itself must be stable: back-to-back shared-memory | |
| chases agree within 20%.""" | |
| ring = mk.ops # ensure ops import path is alive | |
| import metakernel as _m | |
| r = _m._sattolo_ring(4096, "cuda") | |
| out = torch.zeros(2, dtype=torch.int64, device="cuda") | |
| mk.ops.mk_chase_shared(r, 1 << 14, out) | |
| torch.cuda.synchronize() | |
| t1 = int(out[0].item()) | |
| mk.ops.mk_chase_shared(r, 1 << 14, out) | |
| torch.cuda.synchronize() | |
| t2 = int(out[0].item()) | |
| assert abs(t1 - t2) / max(t1, t2) < 0.2, (t1, t2) | |
| # the ring is a single cycle: 2^14 hops through 4096 slots ends where | |
| # modular arithmetic says it must, proving every hop was taken | |
| assert 0 <= int(out[1].item()) < 4096 | |