222 lines
7.3 KiB
Python
222 lines
7.3 KiB
Python
"""Tests for H2D writes and PE DMA probe latency invariants.
|
|
|
|
H2D tests use MemoryWriteMsg (pcie_ep → io_cpu → m_cpu → hbm_ctrl → response).
|
|
PE DMA tests use PeDmaMsg (direct pe_dma → xbar → hbm_ctrl injection).
|
|
"""
|
|
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 MemoryWriteMsg, PeDmaMsg
|
|
from kernbench.sim_engine.engine import GraphEngine
|
|
from kernbench.topology.builder import load_topology
|
|
|
|
TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml"
|
|
|
|
|
|
def _engine():
|
|
return GraphEngine(load_topology(TOPOLOGY_PATH))
|
|
|
|
|
|
def _hbm_pa(sip: int = 0, cube: int = 0, pe_id: int = 0) -> int:
|
|
slice_bytes = 48 * (1 << 30) // 8
|
|
pa = PhysAddr.pe_hbm_addr(
|
|
rack_id=0, sip_id=sip, cube_id=cube, pe_id=pe_id,
|
|
pe_local_hbm_offset=0x1000, slice_size_bytes=slice_bytes,
|
|
)
|
|
return pa.encode()
|
|
|
|
|
|
def _h2d_latency(dst_cube: int, dst_pe: int = 0) -> float:
|
|
engine = _engine()
|
|
msg = MemoryWriteMsg(
|
|
correlation_id="probe", request_id=f"h2d-c{dst_cube}-p{dst_pe}",
|
|
dst_sip=0, dst_cube=dst_cube, dst_pe=dst_pe,
|
|
dst_pa=_hbm_pa(sip=0, cube=dst_cube, pe_id=dst_pe), nbytes=4096,
|
|
pattern="zero", target_pe=dst_pe,
|
|
)
|
|
h = engine.submit(msg)
|
|
engine.wait(h)
|
|
_, trace = engine.get_completion(h)
|
|
return trace["total_ns"]
|
|
|
|
|
|
# ── 1. Single-PE write completes ──────────────────────────────────
|
|
|
|
|
|
def test_single_pe_write_completes():
|
|
"""MemoryWriteMsg(target_pe=0) must complete with ok=True, latency > 0."""
|
|
engine = _engine()
|
|
msg = MemoryWriteMsg(
|
|
correlation_id="probe", request_id="pe-local",
|
|
dst_sip=0, dst_cube=0, dst_pe=0,
|
|
dst_pa=_hbm_pa(sip=0, cube=0, pe_id=0), nbytes=4096,
|
|
pattern="zero", target_pe=0,
|
|
)
|
|
h = engine.submit(msg)
|
|
engine.wait(h)
|
|
comp, trace = engine.get_completion(h)
|
|
assert comp.ok is True
|
|
assert trace["total_ns"] > 0
|
|
|
|
|
|
# ── 2. Cross-cube write positive latency ─────────────────────────
|
|
|
|
|
|
def test_cross_cube_write_positive():
|
|
"""Cross-cube MemoryWriteMsg(target_pe=0) must complete with latency > 0."""
|
|
lat = _h2d_latency(dst_cube=1, dst_pe=0)
|
|
assert lat > 0
|
|
|
|
|
|
# ── 3. H2D latency monotonicity ──────────────────────────────────
|
|
|
|
|
|
def test_h2d_latency_monotonic():
|
|
"""1hop < 2hop < 3hop < 4hop."""
|
|
cubes = [0, 4, 8, 12]
|
|
latencies: list[tuple[int, float]] = []
|
|
for cube in cubes:
|
|
lat = _h2d_latency(dst_cube=cube, dst_pe=0)
|
|
latencies.append((cube, lat))
|
|
|
|
for i in range(len(latencies) - 1):
|
|
assert latencies[i][1] < latencies[i + 1][1], (
|
|
f"cube{latencies[i][0]}({latencies[i][1]:.2f}) "
|
|
f"must < cube{latencies[i + 1][0]}({latencies[i + 1][1]:.2f})"
|
|
)
|
|
|
|
|
|
# ── 4. Single-PE write deterministic ─────────────────────────────
|
|
|
|
|
|
def test_single_pe_write_deterministic():
|
|
"""Same MemoryWriteMsg on two engines must produce identical latency."""
|
|
msg = MemoryWriteMsg(
|
|
correlation_id="probe", request_id="det",
|
|
dst_sip=0, dst_cube=0, dst_pe=0,
|
|
dst_pa=_hbm_pa(sip=0, cube=0, pe_id=0), nbytes=4096,
|
|
pattern="zero", target_pe=0,
|
|
)
|
|
e1 = _engine()
|
|
h1 = e1.submit(msg)
|
|
e1.wait(h1)
|
|
_, t1 = e1.get_completion(h1)
|
|
|
|
e2 = _engine()
|
|
h2 = e2.submit(msg)
|
|
e2.wait(h2)
|
|
_, t2 = e2.get_completion(h2)
|
|
|
|
assert t1["total_ns"] == t2["total_ns"]
|
|
|
|
|
|
# ── 5. Cut-through (wormhole) wire model invariants ──────────────
|
|
|
|
|
|
def test_h2d_local_cube_cut_through():
|
|
"""H2D to local cube with cut-through should be < 50ns for 4096B.
|
|
|
|
Full command path: pcie_ep → io_cpu → ucie → noc → m_cpu
|
|
DMA: m_cpu → noc → xbar → hbm_ctrl (drain once at terminal)
|
|
Plus response path back.
|
|
With store-and-forward each hop would serialize; cut-through keeps it low.
|
|
"""
|
|
lat = _h2d_latency(dst_cube=0, dst_pe=0)
|
|
assert lat < 65.0, f"Local H2D {lat:.2f}ns; cut-through expects < 65ns"
|
|
|
|
|
|
def test_h2d_remote_cube_cut_through():
|
|
"""H2D to 1-hop remote cube: cut-through drain dominates, not per-hop serialization.
|
|
|
|
With store-and-forward, each hop would serialize 4096B, total >> 100ns.
|
|
With cut-through, drain happens once at bottleneck.
|
|
"""
|
|
lat = _h2d_latency(dst_cube=4, dst_pe=0)
|
|
assert lat < 80.0, f"Remote H2D {lat:.2f}ns; cut-through expects < 80ns"
|
|
|
|
|
|
# ── 6. PE DMA: direct injection tests ─────────────────────────
|
|
|
|
|
|
def _graph():
|
|
return load_topology(TOPOLOGY_PATH)
|
|
|
|
|
|
def _pe_dma_latency(src_cube: int, src_pe: int, dst_pe: int) -> float:
|
|
engine = _engine()
|
|
msg = PeDmaMsg(
|
|
correlation_id="probe", request_id=f"dma-c{src_cube}-p{src_pe}-s{dst_pe}",
|
|
src_sip=0, src_cube=src_cube, src_pe=src_pe,
|
|
dst_pa=_hbm_pa(sip=0, cube=src_cube, pe_id=dst_pe), nbytes=4096,
|
|
)
|
|
h = engine.submit(msg)
|
|
engine.wait(h)
|
|
_, trace = engine.get_completion(h)
|
|
return trace["total_ns"]
|
|
|
|
|
|
def _pe_dma_bottleneck(src_cube: int, src_pe: int, dst_pe: int) -> float | None:
|
|
graph = _graph()
|
|
edge_map = {(e.src, e.dst): e for e in graph.edges}
|
|
resolver = AddressResolver(graph)
|
|
router = PathRouter(graph)
|
|
pa = _hbm_pa(sip=0, cube=src_cube, pe_id=dst_pe)
|
|
pa_obj = PhysAddr.decode(pa)
|
|
dst_node = resolver.resolve(pa_obj)
|
|
pe_ref = f"sip0.cube{src_cube}.pe{src_pe}"
|
|
path = router.find_path(pe_ref, dst_node)
|
|
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 test_pe_dma_local_completes():
|
|
"""PeDmaMsg to local slice0 must complete with ok=True, latency > 0."""
|
|
engine = _engine()
|
|
msg = PeDmaMsg(
|
|
correlation_id="probe", request_id="dma-local",
|
|
src_sip=0, src_cube=0, src_pe=0,
|
|
dst_pa=_hbm_pa(sip=0, cube=0, pe_id=0), nbytes=4096,
|
|
)
|
|
h = engine.submit(msg)
|
|
engine.wait(h)
|
|
comp, trace = engine.get_completion(h)
|
|
assert comp.ok is True
|
|
assert trace["total_ns"] > 0
|
|
|
|
|
|
def test_pe_dma_local_bottleneck_256():
|
|
"""PE DMA pe0→slice0 (local): bottleneck = 256 GB/s (direct xbar→hbm)."""
|
|
bn = _pe_dma_bottleneck(src_cube=0, src_pe=0, dst_pe=0)
|
|
assert bn == 256.0, f"Local PE DMA bottleneck {bn}, expected 256.0"
|
|
|
|
|
|
def test_pe_dma_chain_bottleneck_128():
|
|
"""PE DMA pe0→slice1 (xbar chain): bottleneck = 128 GB/s."""
|
|
bn = _pe_dma_bottleneck(src_cube=0, src_pe=0, dst_pe=1)
|
|
assert bn == 128.0, f"Chain PE DMA bottleneck {bn}, expected 128.0"
|
|
|
|
|
|
def test_pe_dma_deterministic():
|
|
"""Same PeDmaMsg on two engines must produce identical latency."""
|
|
msg = PeDmaMsg(
|
|
correlation_id="probe", request_id="det",
|
|
src_sip=0, src_cube=0, src_pe=0,
|
|
dst_pa=_hbm_pa(sip=0, cube=0, pe_id=0), nbytes=4096,
|
|
)
|
|
e1 = _engine()
|
|
h1 = e1.submit(msg)
|
|
e1.wait(h1)
|
|
_, t1 = e1.get_completion(h1)
|
|
|
|
e2 = _engine()
|
|
h2 = e2.submit(msg)
|
|
e2.wait(h2)
|
|
_, t2 = e2.get_completion(h2)
|
|
|
|
assert t1["total_ns"] == t2["total_ns"]
|