ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

benches: package as kernbench.benches, add @bench registry + list subcommand

Browse Source 
Move benches/ -> src/kernbench/benches/ and src/kernbench/cli/probe.py ->
src/kernbench/probes/probe.py. Each bench self-registers via
@bench(name=..., description=...); kernbench list enumerates benches
with auto-assigned indices, --bench accepts kebab-case name or numeric
index. Audit at package-import time fails if any non-underscore module
forgets the decorator. ADR-0010 (EN + KO) updated to reflect the new
resolver path, list subcommand, and probes package separation.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook Kang
2026-05-20 14:42:10 -07:00
parent 168b0c89f0
commit 049e3d8bb3
28 changed files with 398 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/probes/__init__.py
+5
View File
@@ -0,0 +1,5 @@
"""kernbench.probes: latency/BW diagnostic utilities (not benchmarks).
See ADR-0010 D4. Probe is a developer tool for verifying the latency/BW
model; it bypasses the bench registry.
"""
src/kernbench/probes/probe.py
+451
View File
@@ -0,0 +1,451 @@
"""kernbench probe: latency and BW verification utility.
Runs predefined traffic patterns through the simulation engine and reports
latency, effective bandwidth, bottleneck bandwidth, and utilization for each
case. Validates monotonicity invariants across hop counts and access types.
"""
from __future__ import annotations
from pathlib import Path
from kernbench.policy.address.phyaddr import PhysAddr
from kernbench.policy.routing.router import AddressResolver, PathRouter
from kernbench.runtime_api.kernel import MemoryReadMsg, MemoryWriteMsg, PeDmaMsg
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import load_topology
from kernbench.topology.types import TopologyGraph
# -- Helpers ----------------------------------------------------------
def _hbm_pa(sip: int, cube: int, pe_id: int, spec: dict) -> int:
mm = spec["cube"]["memory_map"]
slice_bytes = mm["hbm_total_gb_per_cube"] * (1 << 30) // mm["hbm_slices_per_cube"]
pa = PhysAddr.pe_hbm_addr(
sip_id=sip, die_id=cube, pe_id=pe_id,
pe_local_hbm_offset=0x1000, slice_size_bytes=slice_bytes,
)
return pa.encode()
def _build_edge_map(graph: TopologyGraph) -> dict[tuple[str, str], object]:
return {(e.src, e.dst): e for e in graph.edges}
def _formula_breakdown(
path: list[str], nbytes: int, edge_map: dict, graph: TopologyGraph,
) -> tuple[float, float, float, float]:
"""Return (wire_ns, overhead_ns, drain_ns, formula_ns) for a path."""
ns_per_mm = graph.spec.get("system", {}).get("ns_per_mm", 0.01)
wire_ns = 0.0
for i in range(len(path) - 1):
e = edge_map.get((path[i], path[i + 1]))
if e:
wire_ns += e.distance_mm * ns_per_mm
overhead_ns = 0.0
for nid in path:
node = graph.nodes.get(nid)
if node:
overhead_ns += float(node.attrs.get("overhead_ns", 0.0))
bws = [e.bw_gbs for i in range(len(path) - 1)
if (e := edge_map.get((path[i], path[i + 1]))) and e.bw_gbs]
drain_ns = nbytes / min(bws) if bws else 0.0
return wire_ns, overhead_ns, drain_ns, wire_ns + overhead_ns + drain_ns
def _hop_timestamps(
path: list[str], nbytes: int, edge_map: dict, graph: TopologyGraph,
) -> list[tuple[str, float, str]]:
"""Return per-hop timestamps: [(node_short, cumulative_ns, annotation), ...].
Annotations mark bottleneck edges and significant overhead nodes.
"""
ns_per_mm = graph.spec.get("system", {}).get("ns_per_mm", 0.01)
# Find bottleneck BW for annotation
bws = [e.bw_gbs for i in range(len(path) - 1)
if (e := edge_map.get((path[i], path[i + 1]))) and e.bw_gbs]
bn_bw = min(bws) if bws else None
cumulative = 0.0
result: list[tuple[str, float, str]] = []
result.append((_short_name(path[0]), 0.0, ""))
for i in range(len(path) - 1):
e = edge_map.get((path[i], path[i + 1]))
ann = ""
if e:
cumulative += e.distance_mm * ns_per_mm
if bn_bw is not None and e.bw_gbs and e.bw_gbs == bn_bw:
ann = f"<BN:{e.bw_gbs:.0f}GB/s>"
node = graph.nodes.get(path[i + 1])
if node:
ovhd = float(node.attrs.get("overhead_ns", 0.0))
cumulative += ovhd
if ovhd > 0 and not ann:
ann = f"+{ovhd:.1f}ns"
result.append((_short_name(path[i + 1]), cumulative, ann))
# Add drain at terminal
if bn_bw and nbytes > 0:
cumulative += nbytes / bn_bw
result[-1] = (result[-1][0], cumulative, result[-1][2] + f" drain:{nbytes/bn_bw:.1f}ns")
return result
def _bottleneck_bw(path: list[str], edge_map: dict) -> float | None:
"""Per-request bottleneck: single request uses one connection."""
bws: list[float] = []
for i in range(len(path) - 1):
e = edge_map.get((path[i], path[i + 1]))
if e and e.bw_gbs:
bws.append(e.bw_gbs)
return min(bws) if bws else None
def _fmt_bw(bw: float | None) -> str:
return f"{bw:.1f}" if bw is not None else "-"
def _fmt_util(eff: float, bn: float | None) -> str:
if bn is None or bn <= 0:
return "-"
return f"{eff / bn * 100:.1f}%"
def _short_name(node_id: str) -> str:
"""Shorten node id: keep last 2 segments to avoid ambiguity (router.pe0 vs pe0)."""
parts = node_id.split(".")
return ".".join(parts[-2:]) if len(parts) >= 2 else node_id
def _short_path(path: list[str]) -> str:
return " -> ".join(_short_name(n) for n in path)
def _print_hop_trace(timestamps: list[tuple[str, float, str]], indent: str = " ") -> None:
"""Print per-hop timestamp trace."""
for node, t_ns, ann in timestamps:
ann_str = f" {ann}" if ann else ""
print(f"{indent}{t_ns:>8.2f}ns {node}{ann_str}")
SWEEP_SIZES = [4096, 16384, 65536, 262144, 1048576]
SWEEP_LABELS = ["4KB", "16KB", "64KB", "256KB", "1MB"]
def _sweep_util(ovhd_ns: float, wire_ns: float, bn_bw: float | None, sizes: list[int] = SWEEP_SIZES) -> list[float]:
"""Compute utilization % for each data size using formula model."""
if bn_bw is None or bn_bw <= 0:
return [0.0] * len(sizes)
result = []
for nb in sizes:
drain = nb / bn_bw
total = ovhd_ns + wire_ns + drain
eff = nb / total if total > 0 else 0.0
result.append(eff / bn_bw * 100)
return result
def _print_sweep_table(case_names: list[str], sweep_data: list[list[float]]) -> None:
"""Print compact BW saturation table."""
hdr = f" {'Case':<26}" + "".join(f" {l:>7}" for l in SWEEP_LABELS)
print(f"\n BW Saturation (Util% by data size):")
print(hdr)
print(" " + "-" * (26 + 8 * len(SWEEP_LABELS)))
for name, utils in zip(case_names, sweep_data):
cols = "".join(f" {u:>6.1f}%" for u in utils)
print(f" {name:<26}{cols}")
# -- Probe runner -----------------------------------------------------
def run_probe(topology_path: str, case_filter: str | None = None) -> int:
path = Path(topology_path).expanduser().resolve()
graph = load_topology(path)
edge_map = _build_edge_map(graph)
spec = graph.spec
resolver = AddressResolver(graph)
router = PathRouter(graph)
nbytes = 32768
show_all = case_filter is None or case_filter == "all"
# === Collect H2D results ===
h2d_cases = [
("h2d-1hop", 0, 1),
("h2d-2hop", 4, 2),
("h2d-3hop", 8, 3),
("h2d-4hop", 12, 4),
]
h2d_results: list[tuple[str, int, float, float, float | None, float, float, float, float, float]] = []
h2d_route_data: list[tuple[str, list[str], list[str], list[str], list[str]]] = []
for name, cube, hops in h2d_cases:
if not show_all and case_filter != name:
continue
engine = GraphEngine(graph)
pa = _hbm_pa(sip=0, cube=cube, pe_id=0, spec=spec)
msg = MemoryWriteMsg(
correlation_id="probe", request_id=name,
dst_sip=0, dst_cube=cube, dst_pe=0,
dst_pa=pa, nbytes=nbytes, pattern="zero",
)
h = engine.submit(msg)
engine.wait(h)
_, trace = engine.get_completion(h)
total_ns = trace["total_ns"]
eff_bw = nbytes / total_ns if total_ns > 0 else 0.0
pa_obj = PhysAddr.decode(pa)
dst_node = resolver.resolve(pa_obj)
pcie_ep = resolver.find_pcie_ep(0)
io_cpu = resolver.find_io_cpu(0)
m_cpu = resolver.find_m_cpu(0, cube)
leg1 = router.find_node_path(pcie_ep, io_cpu)
leg2 = router.find_node_path(io_cpu, m_cpu)
leg3 = router.find_mcpu_dma_path(m_cpu, dst_node)
full_path = leg1 + leg2[1:] + leg3[1:]
bn_bw = _bottleneck_bw(full_path, edge_map)
fwd_path = leg1 + leg2[1:] + leg3[1:]
wire, ovhd, drain, formula = _formula_breakdown(fwd_path, nbytes, edge_map, graph)
ovhd_pct = ovhd / total_ns * 100 if total_ns > 0 else 0
drain_pct = drain / total_ns * 100 if total_ns > 0 else 0
h2d_results.append((name, hops, total_ns, eff_bw, bn_bw, ovhd, drain, wire, ovhd_pct, drain_pct))
h2d_route_data.append((name, leg1, leg2, leg3, fwd_path))
# === Collect D2H Read results ===
d2h_cases = [
("d2h-1hop", 0, 1),
("d2h-2hop", 4, 2),
("d2h-3hop", 8, 3),
("d2h-4hop", 12, 4),
]
d2h_results: list[tuple[str, int, float, float, float | None, float, float, float, float, float]] = []
d2h_route_data: list[tuple[str, list[str], list[str], list[str], list[str]]] = []
for name, cube, hops in d2h_cases:
if not show_all and case_filter != name:
continue
engine = GraphEngine(graph)
pa = _hbm_pa(sip=0, cube=cube, pe_id=0, spec=spec)
msg = MemoryReadMsg(
correlation_id="probe", request_id=name,
src_sip=0, src_cube=cube, src_pe=0,
src_pa=pa, nbytes=nbytes,
)
h = engine.submit(msg)
engine.wait(h)
_, trace = engine.get_completion(h)
total_ns = trace["total_ns"]
eff_bw = nbytes / total_ns if total_ns > 0 else 0.0
pa_obj = PhysAddr.decode(pa)
dst_node = resolver.resolve(pa_obj)
pcie_ep = resolver.find_pcie_ep(0)
fwd_path = router.find_memory_path(pcie_ep, dst_node)
rev_path = list(reversed(fwd_path))
bn_bw = _bottleneck_bw(fwd_path, edge_map)
wire, ovhd, drain, formula = _formula_breakdown(fwd_path, nbytes, edge_map, graph)
ovhd_pct = ovhd / total_ns * 100 if total_ns > 0 else 0
drain_pct = drain / total_ns * 100 if total_ns > 0 else 0
d2h_results.append((name, hops, total_ns, eff_bw, bn_bw, ovhd, drain, wire, ovhd_pct, drain_pct))
d2h_route_data.append((name, fwd_path, rev_path, [], fwd_path))
# === Collect PE DMA results ===
pe_cases = [
("pe-local-hbm", 0, 0, 0, 0, 0),
("pe-same-half-hbm", 0, 0, 0, 0, 1),
("pe-cross-half-hbm", 0, 0, 0, 0, 4),
("pe-cross-cube-hbm-best", 0, 0, 0, 1, 0), # adjacent cube
("pe-cross-cube-hbm-worst", 0, 0, 0, 15, 0), # diagonal far cube
]
pe_results: list[tuple[str, float, float, float | None, float, float, float, float, float]] = []
pe_route_data: list[tuple[str, list[str], str]] = []
for name, sip, src_cube, src_pe, dst_cube, dst_pe in pe_cases:
if not show_all and case_filter != name:
continue
engine = GraphEngine(graph)
dst_pa = _hbm_pa(sip=sip, cube=dst_cube, pe_id=dst_pe, spec=spec)
msg = PeDmaMsg(
correlation_id="probe", request_id=name,
src_sip=sip, src_cube=src_cube, src_pe=src_pe,
dst_pa=dst_pa, nbytes=nbytes,
)
h = engine.submit(msg)
engine.wait(h)
_, trace = engine.get_completion(h)
total_ns = trace["total_ns"]
eff_bw = nbytes / total_ns if total_ns > 0 else 0.0
pe_ref = f"sip{sip}.cube{src_cube}.pe{src_pe}"
pa_obj = PhysAddr.decode(dst_pa)
dst_node = resolver.resolve(pa_obj)
dma_path = router.find_path(pe_ref, dst_node)
bn_bw = _bottleneck_bw(dma_path, edge_map)
wire, ovhd, drain, formula = _formula_breakdown(dma_path, nbytes, edge_map, graph)
ovhd_pct = ovhd / total_ns * 100 if total_ns > 0 else 0
drain_pct = drain / total_ns * 100 if total_ns > 0 else 0
target_str = f"c{src_cube}.pe{src_pe}->c{dst_cube}.slice{dst_pe}"
pe_results.append((name, total_ns, eff_bw, bn_bw, ovhd, drain, wire, ovhd_pct, drain_pct))
pe_route_data.append((name, dma_path, target_str))
# ================================================================
# OUTPUT: Summary tables first, then route details
# ================================================================
# --- H2D Summary Table ---
print()
print(f"=== H2D Write Latency (IO->HBM, data={nbytes}B) ===")
print(f" {'Case':<14} {'Target':<16} {'Hops':>4} {'Actual':>8}"
f" {'Ovhd':>6} {'Drain':>6} {'Wire':>5} {'Ovhd%':>6} {'Drain%':>7}"
f" {'Eff.BW':>8} {'BN.BW':>8} {'Util%':>6}")
print(" " + "-" * 115)
for i, (name, hops, total_ns, eff_bw, bn_bw, ovhd, drain, wire, ovhd_pct, drain_pct) in enumerate(h2d_results):
cube = h2d_cases[i][1] if i < len(h2d_cases) else 0
print(f" {name:<14} cube{cube}.pe0{'':<8} {hops:>4} {total_ns:>8.2f}"
f" {ovhd:>6.1f} {drain:>6.1f} {wire:>5.2f} {ovhd_pct:>5.1f}% {drain_pct:>5.1f}%"
f" {eff_bw:>8.2f} {_fmt_bw(bn_bw):>8} {_fmt_util(eff_bw, bn_bw):>6}")
if len(h2d_results) >= 2:
lats = [r[2] for r in h2d_results]
mono = all(lats[i] < lats[i + 1] for i in range(len(lats) - 1))
sym = "[v]" if mono else "[x]"
print(f" {sym} Monotonic increase: {'PASS' if mono else 'FAIL'}")
if h2d_results:
h2d_sweep = [_sweep_util(r[5], r[7], r[4]) for r in h2d_results]
_print_sweep_table([r[0] for r in h2d_results], h2d_sweep)
# --- D2H Summary Table ---
print()
print(f"=== D2H Read Latency (HBM->IO, data={nbytes}B) ===")
print(f" {'Case':<14} {'Source':<16} {'Hops':>4} {'Actual':>8}"
f" {'Ovhd':>6} {'Drain':>6} {'Wire':>5} {'Ovhd%':>6} {'Drain%':>7}"
f" {'Eff.BW':>8} {'BN.BW':>8} {'Util%':>6}")
print(" " + "-" * 115)
for i, (name, hops, total_ns, eff_bw, bn_bw, ovhd, drain, wire, ovhd_pct, drain_pct) in enumerate(d2h_results):
cube = d2h_cases[i][1] if i < len(d2h_cases) else 0
print(f" {name:<14} cube{cube}.pe0{'':<8} {hops:>4} {total_ns:>8.2f}"
f" {ovhd:>6.1f} {drain:>6.1f} {wire:>5.2f} {ovhd_pct:>5.1f}% {drain_pct:>5.1f}%"
f" {eff_bw:>8.2f} {_fmt_bw(bn_bw):>8} {_fmt_util(eff_bw, bn_bw):>6}")
if len(d2h_results) >= 2:
lats = [r[2] for r in d2h_results]
mono = all(lats[i] < lats[i + 1] for i in range(len(lats) - 1))
sym = "[v]" if mono else "[x]"
print(f" {sym} Monotonic increase: {'PASS' if mono else 'FAIL'}")
if d2h_results:
# D2H fixed cost = actual_total - drain (includes fwd+rev overhead)
d2h_sweep = [_sweep_util(r[2] - r[6], 0.0, r[4]) for r in d2h_results]
_print_sweep_table([r[0] for r in d2h_results], d2h_sweep)
# H2D vs D2H comparison
if h2d_results and d2h_results and len(h2d_results) == len(d2h_results):
all_gte = all(d2h_results[i][2] >= h2d_results[i][2] for i in range(len(h2d_results)))
sym = "[v]" if all_gte else "[x]"
print(f" {sym} D2H >= H2D (reverse data path): {'PASS' if all_gte else 'FAIL'}")
# --- PE DMA Summary Table ---
print()
print(f"=== PE DMA Latency (pe_dma -> router -> HBM, data={nbytes}B) ===")
print(f" {'Case':<26} {'Target':<28} {'Actual':>8}"
f" {'Ovhd':>6} {'Drain':>6} {'Wire':>5} {'Ovhd%':>6} {'Drain%':>7}"
f" {'Eff.BW':>8} {'BN.BW':>8} {'Util%':>6}")
print(" " + "-" * 124)
for name, total_ns, eff_bw, bn_bw, ovhd, drain, wire, ovhd_pct, drain_pct in pe_results:
target_str = [t for n, _, t in pe_route_data if n == name]
t_str = target_str[0] if target_str else ""
print(f" {name:<26} {t_str:<28} {total_ns:>8.2f}"
f" {ovhd:>6.1f} {drain:>6.1f} {wire:>5.2f} {ovhd_pct:>5.1f}% {drain_pct:>5.1f}%"
f" {eff_bw:>8.2f} {_fmt_bw(bn_bw):>8} {_fmt_util(eff_bw, bn_bw):>6}")
if len(pe_results) >= 2:
local = [r for r in pe_results if "local" in r[0]]
remote = [r for r in pe_results if "local" not in r[0]]
if local and remote:
print(f" * Local BN: {_fmt_bw(local[0][3])} GB/s, "
f"Remote BN: {_fmt_bw(remote[0][3])} GB/s")
best = [r for r in pe_results if "best" in r[0]]
worst = [r for r in pe_results if "worst" in r[0]]
if best and worst:
sym = "[v]" if best[0][1] < worst[0][1] else "[x]"
print(f" {sym} Cross-cube best < worst: {'PASS' if best[0][1] < worst[0][1] else 'FAIL'}"
f" ({best[0][1]:.2f}ns < {worst[0][1]:.2f}ns)")
if pe_results:
pe_sweep = [_sweep_util(r[4], r[6], r[3]) for r in pe_results]
_print_sweep_table([r[0] for r in pe_results], pe_sweep)
# ================================================================
# ROUTE DETAILS (grouped below all tables)
# ================================================================
print()
print("=" * 60)
print(" ROUTE DETAILS (per-hop timestamps)")
print("=" * 60)
# --- H2D Routes ---
if h2d_route_data:
print()
print(" --- H2D Write Routes ---")
for name, leg1, leg2, leg3, fwd_path in h2d_route_data:
timestamps = _hop_timestamps(fwd_path, nbytes, edge_map, graph)
print(f"\n [{name}]")
print(f" Leg1: {_short_path(leg1)}")
print(f" Leg2: {_short_path(leg2)}")
print(f" Leg3: {_short_path(leg3)}")
print(f" Per-hop trace:")
_print_hop_trace(timestamps, indent=" ")
# --- D2H Routes ---
if d2h_route_data:
print()
print(" --- D2H Read Routes ---")
for name, fwd_path, rev_path, _, _ in d2h_route_data:
timestamps_fwd = _hop_timestamps(fwd_path, 0, edge_map, graph)
timestamps_rev = _hop_timestamps(rev_path, nbytes, edge_map, graph)
print(f"\n [{name}]")
print(f" Fwd (cmd): {_short_path(fwd_path)}")
print(f" Rev (data): {_short_path(rev_path)}")
print(f" Forward cmd trace (no data):")
_print_hop_trace(timestamps_fwd, indent=" ")
print(f" Reverse data trace:")
_print_hop_trace(timestamps_rev, indent=" ")
# --- PE DMA Routes ---
if pe_route_data:
print()
print(" --- PE DMA Routes ---")
for name, dma_path, target_str in pe_route_data:
timestamps = _hop_timestamps(dma_path, nbytes, edge_map, graph)
print(f"\n [{name}] {target_str}")
print(f" Path: {_short_path(dma_path)}")
print(f" Per-hop trace:")
_print_hop_trace(timestamps, indent=" ")
print()
return 0
def cmd_probe(args) -> int:
return run_probe(args.topology, getattr(args, "case", "all"))