ywkang
821bbf26a2
Adds the user-orchestrated async GEMM variants used in paper §3.4 to contrast with the composite (load_ref) path: - matmul-async : naive (tl.load full A+B, one tl.dot) - matmul-async-chunked : depth-inf prefetch (all B-tiles queued) - matmul-async-chunked-db : depth-2 double-buffer (TCM-bounded) plus scripts/paper/paper_plot_gemm_async_vs_composite.py (4-way per-PE TFLOPS sweep + figure) and the regenerated outputs under 1H_milestone_output/gemm/. Sweep regenerated under the ADR-0064 D8 single-op cost model (commit2d8271c). At K=3072 the composite-vs-async-tiled gap narrows from the pre-D8 ~6.3x to ~2.8x (composite 7.18 / async-tiled 2.54 TFLOPS): D8 makes the async-tiled kernel's ~191 single-op dispatches 5x cheaper, so its dispatch overhead drops to ~1.5us. Qualitative order persists: composite > async-full (3.91) > async-tiled (2.54). test_bench_registry.py: register matmul-async / -chunked / -chunked-db (also picks up the earlier milestone-gqa-headline -> milestone-1h-gqa rename already present in the tree). --- Remaining work (resume here if interrupted) --- - Paper §3.4 (03-gemm.tex sec:gemm-vs-async): finish naive->async-full / chunked->async-tiled rename AND reframe "FIXED_DMA=8 for DMA descriptors" to single-op(8) vs composite(40). Figure already copied to docs/report/1H-codesign-paper/figures/gemm_composite_vs_async_tflops.png. - Paper §3.4 K=3072 para + mechanism #3: update dispatch breakdown to new model (191 single-op * 8c ~= 1.5us, was 4.6us) and headline gap ~6.3x -> ~2.8x. - Rebuild build/main.pdf (tectonic) + verify via pdftotext. - Then commit Group 3 (paper + diagrams + PDF). - Table 2 / §2 prose / ADR-0064 D8 already done in2d8271c. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
102 lines
2.9 KiB
Python
102 lines
2.9 KiB
Python
"""Tests for kernbench.benches.registry — @bench decorator + resolve/list."""
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from kernbench.benches import registry
|
|
|
|
|
|
EXPECTED_NAMES = [
|
|
"ccl-allreduce",
|
|
"gemm-single-pe",
|
|
"gpt3-qkv",
|
|
"ipcq-allreduce",
|
|
"matmul-async",
|
|
"matmul-async-chunked",
|
|
"matmul-async-chunked-db",
|
|
"matmul-composite",
|
|
"milestone-1h-ccl",
|
|
"milestone-1h-gemm",
|
|
"milestone-1h-gqa",
|
|
"qkv-gemm",
|
|
"qkv-gemm-multi-pe",
|
|
"va-offset-verify",
|
|
]
|
|
|
|
|
|
def test_registry_lists_all_benches():
|
|
specs = registry.list_all()
|
|
names = [s.name for s in specs]
|
|
assert names == EXPECTED_NAMES
|
|
|
|
|
|
def test_registry_indices_are_1_based_sorted_by_name():
|
|
specs = registry.list_all()
|
|
assert [s.index for s in specs] == list(range(1, len(EXPECTED_NAMES) + 1))
|
|
assert sorted(s.name for s in specs) == [s.name for s in specs]
|
|
|
|
|
|
def test_resolve_by_name_returns_spec():
|
|
spec = registry.resolve("gemm-single-pe")
|
|
assert spec.name == "gemm-single-pe"
|
|
assert callable(spec.run)
|
|
assert spec.description.strip()
|
|
|
|
|
|
def test_resolve_by_index_string_matches_list_order():
|
|
specs = registry.list_all()
|
|
third = specs[2]
|
|
resolved = registry.resolve(str(third.index))
|
|
assert resolved is third
|
|
|
|
|
|
def test_resolve_unknown_name_raises():
|
|
with pytest.raises(ValueError, match="kernbench list"):
|
|
registry.resolve("does-not-exist")
|
|
|
|
|
|
def test_resolve_unknown_index_raises():
|
|
with pytest.raises(ValueError, match="kernbench list"):
|
|
registry.resolve("99")
|
|
|
|
|
|
def test_resolve_empty_identifier_raises():
|
|
with pytest.raises(ValueError):
|
|
registry.resolve("")
|
|
|
|
|
|
def test_bench_decorator_rejects_invalid_name():
|
|
with pytest.raises(ValueError, match="kebab-case"):
|
|
registry.bench(name="Invalid_Name", description="x")
|
|
|
|
|
|
def test_bench_decorator_rejects_empty_description():
|
|
with pytest.raises(ValueError, match="non-empty"):
|
|
registry.bench(name="ok-name", description=" ")
|
|
|
|
|
|
def test_audit_raises_on_missing_decorator():
|
|
with pytest.raises(RuntimeError, match="missing @bench decorator"):
|
|
registry._audit_modules(
|
|
imported=["kernbench.benches.fake_no_dec", "kernbench.benches.real"],
|
|
registered={"kernbench.benches.real"},
|
|
)
|
|
|
|
|
|
def test_audit_passes_when_all_registered():
|
|
registry._audit_modules(
|
|
imported=["kernbench.benches.a", "kernbench.benches.b"],
|
|
registered={"kernbench.benches.a", "kernbench.benches.b"},
|
|
)
|
|
|
|
|
|
def test_duplicate_name_at_finalize_fails(monkeypatch):
|
|
"""_finalize() rejects two pending entries with the same name."""
|
|
monkeypatch.setattr(registry, "_PENDING", [
|
|
("dup", "d1", lambda: None),
|
|
("dup", "d2", lambda: None),
|
|
])
|
|
monkeypatch.setattr(registry, "_REGISTRY", {})
|
|
with pytest.raises(RuntimeError, match="duplicate bench name"):
|
|
registry._finalize()
|