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

Add probe CLI improvements, D2H read, UCIe/HBM tuning, BW sweep

Browse Source 
- Probe CLI: restructured output (tables first, routes below), per-hop
  timestamps, split cross-cube into best/worst cases, D2H read section
- UCIe overhead: 1ns -> 8ns per port (16ns per crossing) to fix
  cross-cube-best < cross-half latency inversion
- HBM efficiency: added efficiency=0.8 factor to hbm_ctrl, reducing
  effective BW from 256 to 204.8 GB/s
- Multi-size BW sweep: saturation tables (4KB-1MB) for all probe cases
- Probe default data size: 4KB -> 32KB for more realistic measurements
- IOChiplet NOC + D2H topology and tests
- NOC mesh, xbar, BW occupancy components and tests
- Cube mesh visualization diagram

278 tests pass.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Yangwook
2026-03-19 01:16:18 -07:00
parent 6f43807900
commit d75da439c6
24 changed files with 3456 additions and 501 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/topology/builder.py
+337 -201
View File
@@ -5,11 +5,13 @@ TopologyGraph with nodes, edges, and representative view projections.
"""
from __future__ import annotations
import math
from pathlib import Path
from typing import Any
import yaml
from .mesh_gen import ensure_mesh_file
from .types import Edge, Node, TopologyGraph, TopologyHandle, ViewGraph
@@ -42,6 +44,10 @@ def load_topology(path: Path) -> TopologyGraph:
"""Load topology spec from file and compile into a topology graph."""
spec = _read_spec(path)
_validate_spec(spec)
# Generate cube_mesh.yaml alongside the topology file
mesh_path = path.parent / "cube_mesh.yaml"
mesh_data = ensure_mesh_file(spec["cube"], mesh_path)
spec["_mesh"] = mesh_data
return _compile_graph(spec)
@@ -110,7 +116,7 @@ def _compile_graph(spec: dict) -> TopologyGraph:
cid = row * mesh_w + col
cp = f"{sp}.cube{cid}"
origin = (col * stride_x, row * stride_y)
_instantiate_cube(nodes, edges, cp, cube_spec, origin)
_instantiate_cube(nodes, edges, cp, cube_spec, origin, spec["_mesh"])
# Inter-cube UCIe mesh
_add_inter_cube_edges(edges, sp, mesh_w, mesh_h, sip_spec)
@@ -148,9 +154,9 @@ def _cube_local_positions(cube_w: float, cube_h: float) -> dict[str, tuple[float
"ucie-W": (uw, cy),
"ucie-E": (cube_w - uw, cy),
"m_cpu": (cube_w - 2.5, cy - 1.5),
"xbar.top": (cx, 3.5), # Y reference for top-half xbar.pe nodes
"xbar_top": (cx, 3.5),
"hbm_ctrl": (cx - 2.0, cy),
"xbar.bottom": (cx, cube_h - 3.5), # Y reference for bottom-half xbar.pe nodes
"xbar_bot": (cx, cube_h - 3.5),
"bridge.left": (2.5, cy + 2.0),
"bridge.right": (cube_w - 2.5, cy + 2.0),
"noc": (cx + 2.0, cy),
@@ -195,10 +201,11 @@ def _instantiate_io_chiplets(
mesh_h: int,
seam: float,
) -> None:
"""Add IO chiplet nodes and internal pcie_ep → io_cpu edges."""
"""Add IO chiplet nodes: pcie_ep, io_cpu, io_noc, io_ucie PHYs, conn nodes."""
io_spec = sip_spec["iochiplet"]
comp = io_spec["components"]
links = io_spec["links"]
ucie_cfg = io_spec.get("ucie", {})
mesh_total_w = mesh_w * cube_w + (mesh_w - 1) * seam
mesh_total_h = mesh_h * cube_h + (mesh_h - 1) * seam
@@ -208,9 +215,9 @@ def _instantiate_io_chiplets(
side = inst["place"]["side"]
cx = mesh_total_w / 2
if side == "N":
pcie_y, cpu_y = -5.0, -3.0
pcie_y, cpu_y, noc_y = -5.0, -3.0, -4.0
else:
pcie_y, cpu_y = mesh_total_h + 5.0, mesh_total_h + 3.0
pcie_y, cpu_y, noc_y = mesh_total_h + 5.0, mesh_total_h + 3.0, mesh_total_h + 4.0
# pcie_ep
ep = comp["pcie_ep"]
@@ -228,13 +235,114 @@ def _instantiate_io_chiplets(
attrs=cpu["attrs"], pos_mm=(cx, cpu_y), label="IO CPU",
)
# Internal edge
# io_noc (central switch inside IOChiplet)
noc = comp["io_noc"]
noc_id = f"{prefix}.noc"
nodes[noc_id] = Node(
id=noc_id, kind=noc["kind"], impl=noc["impl"],
attrs=noc["attrs"], pos_mm=(cx, noc_y), label="IO NOC",
)
# pcie_ep ↔ io_noc (bidirectional)
edges.append(Edge(
src=ep_id, dst=cpu_id,
distance_mm=links["pcie_ep_to_io_cpu_mm"],
bw_gbs=links["pcie_ep_to_io_cpu_bw_gbs"],
src=ep_id, dst=noc_id,
distance_mm=links["pcie_ep_to_noc_mm"],
bw_gbs=links["pcie_ep_to_noc_bw_gbs"],
kind="io_internal",
))
edges.append(Edge(
src=noc_id, dst=ep_id,
distance_mm=links["pcie_ep_to_noc_mm"],
bw_gbs=links["pcie_ep_to_noc_bw_gbs"],
kind="io_internal",
))
# io_cpu ↔ io_noc (bidirectional)
edges.append(Edge(
src=cpu_id, dst=noc_id,
distance_mm=links["io_cpu_to_noc_mm"],
bw_gbs=links["io_cpu_to_noc_bw_gbs"],
kind="io_internal",
))
edges.append(Edge(
src=noc_id, dst=cpu_id,
distance_mm=links["io_cpu_to_noc_mm"],
bw_gbs=links["io_cpu_to_noc_bw_gbs"],
kind="io_internal",
))
# io_ucie PHY nodes + conn nodes per PHY
io_ucie_ns = float(ucie_cfg.get("overhead_ns", 1.0))
io_n_conn = int(ucie_cfg.get("n_connections", 4))
io_conn_bw = float(ucie_cfg.get("per_connection_bw_gbs", 128.0))
io_noc_to_ucie_mm = float(ucie_cfg.get("noc_to_ucie_mm", 0.5))
for phy in inst["ucie"]["phys"]:
phy_id = f"{prefix}.ucie-{phy}"
nodes[phy_id] = Node(
id=phy_id, kind="io_ucie", impl="ucie_v1",
attrs={"overhead_ns": io_ucie_ns},
pos_mm=(cx, noc_y), label=f"IO UCIe-{phy}",
)
for ci in range(io_n_conn):
conn_id = f"{phy_id}.conn{ci}"
nodes[conn_id] = Node(
id=conn_id, kind="io_ucie_conn", impl="ucie_v1",
attrs={"overhead_ns": 0.0},
pos_mm=(cx, noc_y), label=f"IO UCIe-{phy} C{ci}",
)
# io_noc ↔ conn (per-connection BW)
edges.append(Edge(
src=noc_id, dst=conn_id,
distance_mm=io_noc_to_ucie_mm,
bw_gbs=io_conn_bw,
kind="io_noc_to_conn",
))
edges.append(Edge(
src=conn_id, dst=noc_id,
distance_mm=io_noc_to_ucie_mm,
bw_gbs=io_conn_bw,
kind="conn_to_io_noc",
))
# conn ↔ io_ucie (internal, no BW limit)
edges.append(Edge(
src=conn_id, dst=phy_id,
distance_mm=0.0, kind="io_ucie_internal",
))
edges.append(Edge(
src=phy_id, dst=conn_id,
distance_mm=0.0, kind="io_ucie_internal",
))
# ── PE-to-router distance ─────────────────────────────────────────
def _compute_pe_noc_distances(
mesh_data: dict,
corner_pos: dict[str, list[tuple[float, float]]],
corners: list[str],
pe_per_corner: int,
) -> dict[int, float]:
"""Compute per-PE Euclidean distance from physical position to assigned router."""
distances: dict[int, float] = {}
routers = mesh_data["routers"]
pe_idx = 0
for corner in corners:
for ci in range(pe_per_corner):
pe_cx, pe_cy = corner_pos[corner][ci]
target = f"pe{pe_idx}.dma"
for _rkey, rval in routers.items():
if rval is not None and target in rval.get("attach", []):
rx, ry = rval["pos_mm"]
dist = math.sqrt((pe_cx - rx) ** 2 + (pe_cy - ry) ** 2)
distances[pe_idx] = round(dist, 2)
break
else:
distances[pe_idx] = 0.0
pe_idx += 1
return distances
# ── Instantiation: cube + PEs ───────────────────────────────────────
@@ -246,18 +354,26 @@ def _instantiate_cube(
cp: str,
cube: dict,
origin: tuple[float, float],
mesh_data: dict,
) -> None:
"""Add all cube-internal nodes and edges, including PE instances."""
"""Add all cube-internal nodes and edges, including PE instances.
Topology: PE_DMA → NOC → xbar_top/bot → HBM_CTRL.
No per-PE xbar nodes; position-aware XBAR top/bottom replaces chaining.
"""
cube_w = cube["geometry"]["cube_mm"]["w"]
cube_h = cube["geometry"]["cube_mm"]["h"]
ox, oy = origin
local_pos = _cube_local_positions(cube_w, cube_h)
clinks = cube["links"]
n_slices = cube["memory_map"]["hbm_slices_per_cube"]
half = n_slices // 2
# ── UCIe ports ──
ucie_ns = cube["ucie"]["overhead_ns"]
for port in cube["ucie"]["ports"]:
# ── UCIe ports + connection nodes ──
ucie_cfg = cube["ucie"]
ucie_ns = ucie_cfg["overhead_ns"]
ucie_n_conn = ucie_cfg.get("n_connections", 1)
for port in ucie_cfg["ports"]:
pid = f"{cp}.ucie-{port}"
lx, ly = local_pos[f"ucie-{port}"]
nodes[pid] = Node(
@@ -265,6 +381,14 @@ def _instantiate_cube(
attrs={"overhead_ns": ucie_ns}, pos_mm=(ox + lx, oy + ly),
label=f"UCIe-{port}",
)
for ci in range(ucie_n_conn):
conn_id = f"{cp}.ucie-{port}.conn{ci}"
nodes[conn_id] = Node(
id=conn_id, kind="ucie_conn", impl="ucie_v1",
attrs={"overhead_ns": 0.0},
pos_mm=(ox + lx, oy + ly),
label=f"UCIe-{port} C{ci}",
)
# ── Named components: noc, m_cpu, sram ──
for name in ("noc", "m_cpu", "sram"):
@@ -277,7 +401,19 @@ def _instantiate_cube(
label=name.upper().replace("_", " "),
)
# ── HBM controller slices (one per PE) ──
# ── xbar_top and xbar_bot (position-aware XBAR) ──
xbar_spec = cube["components"]["xbar"]
for xbar_name, xbar_cfg in [("xbar_top", xbar_spec["top"]),
("xbar_bot", xbar_spec["bottom"])]:
nid = f"{cp}.{xbar_name}"
lx, ly = local_pos[xbar_name]
nodes[nid] = Node(
id=nid, kind=xbar_cfg["kind"], impl=xbar_cfg["impl"],
attrs=xbar_cfg["attrs"], pos_mm=(ox + lx, oy + ly),
label=xbar_name.upper().replace("_", " "),
)
# ── HBM controller slices ──
hbm_spec = cube["components"]["hbm_ctrl"]
hbm_lx, hbm_ly = local_pos["hbm_ctrl"]
for sl in range(n_slices):
@@ -289,7 +425,7 @@ def _instantiate_cube(
)
# ── Bridges ──
for br in cube["components"]["xbar"]["bridges"]:
for br in xbar_spec["bridges"]:
bname = br["id"]
nid = f"{cp}.bridge.{bname}"
lx, ly = local_pos[f"bridge.{bname}"]
@@ -299,34 +435,22 @@ def _instantiate_cube(
label=f"Bridge {bname.upper()}",
)
# ── PE instances + per-PE xbar entry nodes ──
# ── PE instances (no per-PE xbar nodes) ──
corners = cube["pe_layout"]["corners"]
pe_per_corner = cube["pe_layout"]["pe_per_corner"]
corner_pos = _corner_pe_positions(cube_w, cube_h)
pe_tmpl = cube["pe_template"]
pe_links = pe_tmpl["links"]
xbar_pe_spec = cube["components"]["xbar"]["pe"]
xbar_top_y = local_pos["xbar.top"][1]
xbar_bot_y = local_pos["xbar.bottom"][1]
pe_noc_distances = _compute_pe_noc_distances(
mesh_data, corner_pos, corners, pe_per_corner,
)
pe_idx = 0
for corner in corners:
is_top = corner in ("NW", "NE")
xbar_y = xbar_top_y if is_top else xbar_bot_y
mm_key = "pe_to_xbar_row_n_mm" if is_top else "pe_to_xbar_row_s_mm"
for ci in range(pe_per_corner):
pp = f"{cp}.pe{pe_idx}"
pe_cx, pe_cy = corner_pos[corner][ci]
# Per-PE xbar entry node
xbar_nid = f"{cp}.xbar.pe{pe_idx}"
nodes[xbar_nid] = Node(
id=xbar_nid, kind=xbar_pe_spec["kind"], impl=xbar_pe_spec["impl"],
attrs=xbar_pe_spec["attrs"], pos_mm=(ox + pe_cx, oy + xbar_y),
label=f"XBAR PE{pe_idx}",
)
# PE template components
for comp_name, comp_spec in pe_tmpl["components"].items():
cid = f"{pp}.{comp_name}"
@@ -341,18 +465,10 @@ def _instantiate_cube(
# PE-internal edges
_add_pe_internal_edges(edges, pp, pe_links)
# PE_DMA → xbar.pe_i (HBM data path)
edges.append(Edge(
src=f"{pp}.pe_dma", dst=xbar_nid,
distance_mm=clinks[mm_key],
bw_gbs=clinks["pe_to_xbar_bw_gbs"],
kind="pe_to_xbar",
))
# PE_DMA → noc (non-HBM data path: SRAM, inter-cube, etc.)
# PE_DMA → noc (distance auto-computed from PE physical position)
edges.append(Edge(
src=f"{pp}.pe_dma", dst=f"{cp}.noc",
distance_mm=clinks["pe_dma_to_noc_mm"],
distance_mm=pe_noc_distances.get(pe_idx, 0.0),
bw_gbs=clinks["pe_dma_to_noc_bw_gbs"],
kind="pe_to_noc",
))
@@ -366,97 +482,96 @@ def _instantiate_cube(
pe_idx += 1
# ── Cube fabric edges ──
# xbar.pe_i ↔ hbm_ctrl.slice_i (local Y-path, bidirectional for response)
for i in range(n_slices):
# ── xbar_top/bot → HBM slices ──
hbm_eff = float(hbm_spec.get("attrs", {}).get("efficiency", 1.0))
hbm_bw = clinks["xbar_to_hbm_bw_gbs"] * hbm_eff
for i in range(half):
edges.append(Edge(
src=f"{cp}.xbar.pe{i}", dst=f"{cp}.hbm_ctrl.slice{i}",
src=f"{cp}.xbar_top", dst=f"{cp}.hbm_ctrl.slice{i}",
distance_mm=clinks["xbar_to_hbm_mm"],
bw_gbs=clinks["xbar_to_hbm_bw_gbs"],
bw_gbs=hbm_bw,
kind="xbar_to_hbm",
))
edges.append(Edge(
src=f"{cp}.hbm_ctrl.slice{i}", dst=f"{cp}.xbar.pe{i}",
src=f"{cp}.hbm_ctrl.slice{i}", dst=f"{cp}.xbar_top",
distance_mm=clinks["xbar_to_hbm_mm"],
bw_gbs=clinks["xbar_to_hbm_bw_gbs"],
bw_gbs=hbm_bw,
kind="hbm_to_xbar",
))
for i in range(half, n_slices):
edges.append(Edge(
src=f"{cp}.xbar_bot", dst=f"{cp}.hbm_ctrl.slice{i}",
distance_mm=clinks["xbar_to_hbm_mm"],
bw_gbs=hbm_bw,
kind="xbar_to_hbm",
))
edges.append(Edge(
src=f"{cp}.hbm_ctrl.slice{i}", dst=f"{cp}.xbar_bot",
distance_mm=clinks["xbar_to_hbm_mm"],
bw_gbs=hbm_bw,
kind="hbm_to_xbar",
))
# xbar chain: pe0↔pe1↔pe2↔pe3 (top), pe4↔pe5↔pe6↔pe7 (bottom)
half = n_slices // 2
for half_start in (0, half):
for i in range(half_start, half_start + half - 1):
intra = ((i - half_start) % pe_per_corner) != (pe_per_corner - 1)
x_dist = clinks["xbar_chain_intra_corner_mm"] if intra else clinks["xbar_chain_inter_corner_mm"]
for a, b in [(i, i + 1), (i + 1, i)]:
edges.append(Edge(
src=f"{cp}.xbar.pe{a}", dst=f"{cp}.xbar.pe{b}",
distance_mm=x_dist,
bw_gbs=clinks["xbar_x_bw_gbs"],
kind="xbar_chain",
))
# ── NOC ↔ xbar_top/bot ──
# xbar_top: primary (low routing weight), xbar_bot: secondary (high routing weight
# steers Dijkstra through xbar_top→bridge→xbar_bot for cross-half access)
noc_xbar_bw = clinks.get("noc_to_xbar_bw_gbs", 256.0)
noc_xbar_mm = clinks.get("noc_to_xbar_mm", 0.0)
for xbar_name, rw in [("xbar_top", None), ("xbar_bot", 100.0)]:
edges.append(Edge(
src=f"{cp}.noc", dst=f"{cp}.{xbar_name}",
distance_mm=noc_xbar_mm, bw_gbs=noc_xbar_bw,
routing_weight_mm=rw, kind="noc_to_xbar",
))
edges.append(Edge(
src=f"{cp}.{xbar_name}", dst=f"{cp}.noc",
distance_mm=noc_xbar_mm, bw_gbs=noc_xbar_bw,
routing_weight_mm=rw, kind="xbar_to_noc",
))
# bridge connections: pe0↔bridge.left↔pe4, pe3↔bridge.right↔pe7
for bname, pe_top, pe_bot in [("left", 0, half), ("right", half - 1, n_slices - 1)]:
# ── Bridge connections: xbar_top ↔ bridge ↔ xbar_bot ──
bridge_mm = clinks.get("xbar_to_bridge_mm", 3.0)
bridge_bw = clinks.get("xbar_to_bridge_bw_gbs", 128.0)
for bname in ("left", "right"):
br_node = f"{cp}.bridge.{bname}"
for pe_i, br_mm_key in [(pe_top, "xbar_row_n_to_bridge_mm"),
(pe_bot, "xbar_row_s_to_bridge_mm")]:
xbar_node = f"{cp}.xbar.pe{pe_i}"
for xbar_name in ("xbar_top", "xbar_bot"):
edges.append(Edge(
src=xbar_node, dst=br_node,
distance_mm=clinks[br_mm_key],
bw_gbs=clinks["xbar_to_bridge_bw_gbs"],
src=f"{cp}.{xbar_name}", dst=br_node,
distance_mm=bridge_mm, bw_gbs=bridge_bw,
kind="xbar_to_bridge",
))
edges.append(Edge(
src=br_node, dst=xbar_node,
distance_mm=clinks[br_mm_key],
bw_gbs=clinks["xbar_to_bridge_bw_gbs"],
src=br_node, dst=f"{cp}.{xbar_name}",
distance_mm=bridge_mm, bw_gbs=bridge_bw,
kind="bridge_to_xbar",
))
# ucie ↔ noc (UCIe-NOC boundary; per_connection_bw_gbs = 128 GB/s, n_connections = 4)
_noc_ucie = clinks["noc_to_ucie"]
for port in cube["ucie"]["ports"]:
edges.append(Edge(
src=f"{cp}.ucie-{port}", dst=f"{cp}.noc",
distance_mm=0.0,
bw_gbs=_noc_ucie["per_connection_bw_gbs"],
n_connections=_noc_ucie["n_connections"],
kind="ucie_to_noc",
))
# ── UCIe ↔ conn ↔ NOC ──
ucie_conn_bw = ucie_cfg.get("per_connection_bw_gbs", 128.0)
for port in ucie_cfg["ports"]:
ucie_id = f"{cp}.ucie-{port}"
for ci in range(ucie_n_conn):
conn_id = f"{cp}.ucie-{port}.conn{ci}"
edges.append(Edge(
src=ucie_id, dst=conn_id,
distance_mm=0.0, kind="ucie_internal",
))
edges.append(Edge(
src=conn_id, dst=ucie_id,
distance_mm=0.0, kind="ucie_internal",
))
edges.append(Edge(
src=conn_id, dst=f"{cp}.noc",
distance_mm=0.0, bw_gbs=ucie_conn_bw,
kind="ucie_conn_to_noc",
))
edges.append(Edge(
src=f"{cp}.noc", dst=conn_id,
distance_mm=0.0, bw_gbs=ucie_conn_bw,
kind="noc_to_ucie_conn",
))
for port in cube["ucie"]["ports"]:
edges.append(Edge(
src=f"{cp}.noc", dst=f"{cp}.ucie-{port}",
distance_mm=0.0,
bw_gbs=_noc_ucie["per_connection_bw_gbs"],
n_connections=_noc_ucie["n_connections"],
kind="noc_to_ucie",
))
# noc ↔ xbar.pe{i}: wire delay is 0 (NOC traversal latency computed by TwoDMeshNocComponent);
# routing_weight_mm=50.0 steers PE DMA Dijkstra away from this path (prefer direct pe_dma→xbar)
_noc_xbar = clinks.get("noc_to_xbar", {})
_noc_xbar_bw = _noc_xbar.get("per_connection_bw_gbs")
for i in range(n_slices):
edges.append(Edge(
src=f"{cp}.noc", dst=f"{cp}.xbar.pe{i}",
distance_mm=0.0,
bw_gbs=_noc_xbar_bw,
routing_weight_mm=50.0,
kind="noc_to_xbar",
))
edges.append(Edge(
src=f"{cp}.xbar.pe{i}", dst=f"{cp}.noc",
distance_mm=0.0,
bw_gbs=_noc_xbar_bw,
routing_weight_mm=50.0,
kind="xbar_to_noc",
))
# m_cpu ↔ noc (command dispatch, both directions)
# ── m_cpu ↔ noc (command dispatch) ──
edges.append(Edge(
src=f"{cp}.m_cpu", dst=f"{cp}.noc",
distance_mm=clinks["m_cpu_to_noc_mm"],
@@ -468,7 +583,7 @@ def _instantiate_cube(
kind="command",
))
# noc ↔ sram (shared SRAM access; per_connection_bw_gbs = 128 GB/s, n_connections = 4)
# ── noc ↔ sram ──
_noc_sram = clinks["noc_to_sram"]
edges.append(Edge(
src=f"{cp}.noc", dst=f"{cp}.sram",
@@ -550,28 +665,27 @@ def _add_inter_cube_edges(
def _add_io_to_cube_edges(
edges: list[Edge], sp: str, sip_spec: dict, mesh_w: int,
) -> None:
"""Add IO chiplet io_cpu ↔ cube UCIe edges (bidirectional for response)."""
io_links = sip_spec["iochiplet"]["links"]
io_to_ucie_mm = io_links["io_cpu_to_ucie_mm"]
io_to_ucie_bw = io_links["io_cpu_to_ucie_bw_gbs"]
"""Add IO chiplet io_ucie ↔ cube UCIe edges (bidirectional)."""
for inst in sip_spec["iochiplet"]["instances"]:
iid = inst["id"]
io_cpu_id = f"{sp}.{iid}.io_cpu"
phy_bw = float(inst["ucie"]["phy_bw_gbs"])
for port in inst["cube_ports"]:
cube_col, cube_row = port["cube"]["xy"]
cube_id = cube_row * mesh_w + cube_col
cube_side = port["cube_side"]
ucie_id = f"{sp}.cube{cube_id}.ucie-{cube_side}"
phy = port["phy"]
io_ucie_id = f"{sp}.{iid}.ucie-{phy}"
cube_ucie_id = f"{sp}.cube{cube_id}.ucie-{cube_side}"
edges.append(Edge(
src=io_cpu_id, dst=ucie_id,
distance_mm=io_to_ucie_mm + port["distance_mm"],
bw_gbs=io_to_ucie_bw,
src=io_ucie_id, dst=cube_ucie_id,
distance_mm=port["distance_mm"],
bw_gbs=phy_bw,
kind="io_to_cube",
))
edges.append(Edge(
src=ucie_id, dst=io_cpu_id,
distance_mm=io_to_ucie_mm + port["distance_mm"],
bw_gbs=io_to_ucie_bw,
src=cube_ucie_id, dst=io_ucie_id,
distance_mm=port["distance_mm"],
bw_gbs=phy_bw,
kind="cube_to_io",
))
@@ -704,11 +818,13 @@ def _build_sip_view(spec: dict) -> ViewGraph:
))
# IO chiplets
io_links = sip_spec["iochiplet"]["links"]
io_ucie_cfg = sip_spec["iochiplet"].get("ucie", {})
io_noc_to_ucie_mm = float(io_ucie_cfg.get("noc_to_ucie_mm", 0.5))
for inst in sip_spec["iochiplet"]["instances"]:
iid = inst["id"]
side = inst["place"]["side"]
iy = 2.0 if side == "N" else canvas_h - 2.0
phy_bw = float(inst["ucie"]["phy_bw_gbs"])
nodes[iid] = Node(
id=iid, kind="iochiplet", impl="",
attrs={}, pos_mm=(mesh_total_w / 2, iy), label=f"IO {iid}",
@@ -718,8 +834,8 @@ def _build_sip_view(spec: dict) -> ViewGraph:
cube_id = cube_row * mesh_w + cube_col
view_edges.append(Edge(
src=iid, dst=f"cube{cube_id}",
distance_mm=io_links["io_cpu_to_ucie_mm"] + port["distance_mm"],
bw_gbs=io_links["io_cpu_to_ucie_bw_gbs"],
distance_mm=io_noc_to_ucie_mm + port["distance_mm"],
bw_gbs=phy_bw,
kind="io_to_cube",
))
@@ -737,31 +853,52 @@ def _build_cube_view(spec: dict) -> ViewGraph:
local_pos = _cube_local_positions(cube_w, cube_h)
clinks = cube["links"]
n_slices = cube["memory_map"]["hbm_slices_per_cube"]
half = n_slices // 2
nodes: dict[str, Node] = {}
view_edges: list[Edge] = []
# UCIe ports
for port in cube["ucie"]["ports"]:
# UCIe ports + connection nodes
ucie_cfg = cube["ucie"]
ucie_n_conn = ucie_cfg.get("n_connections", 1)
for port in ucie_cfg["ports"]:
pid = f"ucie-{port}"
lx, ly = local_pos[pid]
nodes[pid] = Node(
id=pid, kind="ucie_port", impl="ucie_v1",
attrs={}, pos_mm=(lx, ly), label=f"UCIe-{port}",
)
for ci in range(ucie_n_conn):
conn_id = f"ucie-{port}.conn{ci}"
nodes[conn_id] = Node(
id=conn_id, kind="ucie_conn", impl="ucie_v1",
attrs={"overhead_ns": 0.0}, pos_mm=(lx, ly),
label=f"UCIe-{port} C{ci}",
)
# Named components (hbm_ctrl as single representative node in view)
for name in ("noc", "m_cpu", "hbm_ctrl", "sram"):
c = cube["components"][name]
lx, ly = local_pos[name]
lx, ly = local_pos.get(name, local_pos.get("hbm_ctrl"))
nodes[name] = Node(
id=name, kind=c["kind"], impl=c["impl"],
attrs=c["attrs"], pos_mm=(lx, ly),
label=name.upper().replace("_", " "),
)
# xbar_top, xbar_bot
xbar_spec = cube["components"]["xbar"]
for xbar_name, xbar_cfg in [("xbar_top", xbar_spec["top"]),
("xbar_bot", xbar_spec["bottom"])]:
lx, ly = local_pos[xbar_name]
nodes[xbar_name] = Node(
id=xbar_name, kind=xbar_cfg["kind"], impl=xbar_cfg["impl"],
attrs=xbar_cfg["attrs"], pos_mm=(lx, ly),
label=xbar_name.upper().replace("_", " "),
)
# Bridges
for br in cube["components"]["xbar"]["bridges"]:
for br in xbar_spec["bridges"]:
bname = br["id"]
bid = f"bridge.{bname}"
lx, ly = local_pos[bid]
@@ -771,46 +908,29 @@ def _build_cube_view(spec: dict) -> ViewGraph:
label=f"Bridge {bname.upper()}",
)
# PEs as opaque blocks + per-PE xbar entry nodes
# PEs as opaque blocks (no per-PE xbar nodes)
corners = cube["pe_layout"]["corners"]
pe_per_corner = cube["pe_layout"]["pe_per_corner"]
corner_pos = _corner_pe_positions(cube_w, cube_h)
xbar_pe_spec = cube["components"]["xbar"]["pe"]
xbar_top_y = local_pos["xbar.top"][1]
xbar_bot_y = local_pos["xbar.bottom"][1]
mesh_data = spec.get("_mesh", {})
pe_noc_distances = _compute_pe_noc_distances(
mesh_data, corner_pos, corners, pe_per_corner,
) if mesh_data else {}
pe_idx = 0
for corner in corners:
is_top = corner in ("NW", "NE")
xbar_y = xbar_top_y if is_top else xbar_bot_y
mm_key = "pe_to_xbar_row_n_mm" if is_top else "pe_to_xbar_row_s_mm"
for ci in range(pe_per_corner):
pid = f"pe{pe_idx}"
xbar_id = f"xbar.pe{pe_idx}"
px, py = corner_pos[corner][ci]
nodes[pid] = Node(
id=pid, kind="pe", impl="",
attrs={"corner": corner}, pos_mm=(px, py),
label=f"PE{pe_idx}",
)
nodes[xbar_id] = Node(
id=xbar_id, kind=xbar_pe_spec["kind"], impl=xbar_pe_spec["impl"],
attrs=xbar_pe_spec["attrs"], pos_mm=(px, xbar_y),
label=f"XBAR PE{pe_idx}",
)
# PE → xbar.pe_i (HBM data path)
view_edges.append(Edge(
src=pid, dst=xbar_id,
distance_mm=clinks[mm_key],
bw_gbs=clinks["pe_to_xbar_bw_gbs"],
kind="pe_to_xbar",
))
# PE → noc (non-HBM data path)
# PE → noc (distance auto-computed from PE physical position)
view_edges.append(Edge(
src=pid, dst="noc",
distance_mm=clinks["pe_dma_to_noc_mm"],
distance_mm=pe_noc_distances.get(pe_idx, 0.0),
bw_gbs=clinks["pe_dma_to_noc_bw_gbs"],
kind="pe_to_noc",
))
@@ -822,60 +942,76 @@ def _build_cube_view(spec: dict) -> ViewGraph:
))
pe_idx += 1
# Cube fabric edges
# xbar.pe_i → hbm_ctrl (single representative node in view)
for i in range(n_slices):
# xbar_top/bot → hbm_ctrl
view_edges.append(Edge(
src="xbar_top", dst="hbm_ctrl",
distance_mm=clinks["xbar_to_hbm_mm"],
bw_gbs=clinks["xbar_to_hbm_bw_gbs"],
kind="xbar_to_hbm",
))
view_edges.append(Edge(
src="xbar_bot", dst="hbm_ctrl",
distance_mm=clinks["xbar_to_hbm_mm"],
bw_gbs=clinks["xbar_to_hbm_bw_gbs"],
kind="xbar_to_hbm",
))
# noc ↔ xbar_top/bot
noc_xbar_bw = clinks.get("noc_to_xbar_bw_gbs", 256.0)
noc_xbar_mm = clinks.get("noc_to_xbar_mm", 0.0)
for xbar_name in ("xbar_top", "xbar_bot"):
view_edges.append(Edge(
src=f"xbar.pe{i}", dst="hbm_ctrl",
distance_mm=clinks["xbar_to_hbm_mm"],
bw_gbs=clinks["xbar_to_hbm_bw_gbs"],
kind="xbar_to_hbm",
src="noc", dst=xbar_name,
distance_mm=noc_xbar_mm, bw_gbs=noc_xbar_bw,
kind="noc_to_xbar",
))
view_edges.append(Edge(
src=xbar_name, dst="noc",
distance_mm=noc_xbar_mm, bw_gbs=noc_xbar_bw,
kind="xbar_to_noc",
))
# xbar chain
half = n_slices // 2
for half_start in (0, half):
for i in range(half_start, half_start + half - 1):
intra = ((i - half_start) % pe_per_corner) != (pe_per_corner - 1)
x_dist = clinks["xbar_chain_intra_corner_mm"] if intra else clinks["xbar_chain_inter_corner_mm"]
for a, b in [(i, i + 1), (i + 1, i)]:
view_edges.append(Edge(
src=f"xbar.pe{a}", dst=f"xbar.pe{b}",
distance_mm=x_dist,
bw_gbs=clinks["xbar_x_bw_gbs"],
kind="xbar_chain",
))
# bridge connections
for bname, pe_top, pe_bot in [("left", 0, half), ("right", half - 1, n_slices - 1)]:
# bridge connections: xbar_top ↔ bridge ↔ xbar_bot
bridge_mm = clinks.get("xbar_to_bridge_mm", 3.0)
bridge_bw = clinks.get("xbar_to_bridge_bw_gbs", 128.0)
for bname in ("left", "right"):
br_id = f"bridge.{bname}"
for pe_i, br_mm_key in [(pe_top, "xbar_row_n_to_bridge_mm"),
(pe_bot, "xbar_row_s_to_bridge_mm")]:
xbar_id = f"xbar.pe{pe_i}"
for xbar_name in ("xbar_top", "xbar_bot"):
view_edges.append(Edge(
src=xbar_id, dst=br_id,
distance_mm=clinks[br_mm_key],
bw_gbs=clinks["xbar_to_bridge_bw_gbs"],
src=xbar_name, dst=br_id,
distance_mm=bridge_mm, bw_gbs=bridge_bw,
kind="xbar_to_bridge",
))
view_edges.append(Edge(
src=br_id, dst=xbar_id,
distance_mm=clinks[br_mm_key],
bw_gbs=clinks["xbar_to_bridge_bw_gbs"],
src=br_id, dst=xbar_name,
distance_mm=bridge_mm, bw_gbs=bridge_bw,
kind="bridge_to_xbar",
))
_noc_ucie_v = clinks["noc_to_ucie"]
for port in cube["ucie"]["ports"]:
view_edges.append(Edge(
src="noc", dst=f"ucie-{port}",
distance_mm=0.0,
bw_gbs=_noc_ucie_v["per_connection_bw_gbs"],
n_connections=_noc_ucie_v["n_connections"],
kind="noc_to_ucie",
))
ucie_conn_bw_v = ucie_cfg.get("per_connection_bw_gbs", 128.0)
for port in ucie_cfg["ports"]:
for ci in range(ucie_n_conn):
conn_id = f"ucie-{port}.conn{ci}"
view_edges.append(Edge(
src="noc", dst=conn_id,
distance_mm=0.0, bw_gbs=ucie_conn_bw_v,
kind="noc_to_ucie_conn",
))
view_edges.append(Edge(
src=conn_id, dst=f"ucie-{port}",
distance_mm=0.0, kind="ucie_internal",
))
view_edges.append(Edge(
src=f"ucie-{port}", dst=conn_id,
distance_mm=0.0, kind="ucie_internal",
))
view_edges.append(Edge(
src=conn_id, dst="noc",
distance_mm=0.0, bw_gbs=ucie_conn_bw_v,
kind="ucie_conn_to_noc",
))
# m_cpu ↔ noc (command dispatch, both directions)
# m_cpu ↔ noc
view_edges.append(Edge(
src="m_cpu", dst="noc",
distance_mm=clinks["m_cpu_to_noc_mm"],
@@ -887,7 +1023,7 @@ def _build_cube_view(spec: dict) -> ViewGraph:
kind="command",
))
# noc ↔ sram (shared SRAM access, bidirectional)
# noc ↔ sram
_noc_sram_v = clinks["noc_to_sram"]
view_edges.append(Edge(
src="noc", dst="sram",
src/kernbench/topology/mesh_gen.py
+284
View File
@@ -0,0 +1,284 @@
"""Auto-layout mesh generation for CUBE NOC router mesh.
Generates cube_mesh.yaml describing the internal router grid, PE/UCIe/XBAR
attachments, and HBM exclusion zone. The file is cached with a source_hash
so it is only regenerated when relevant topology parameters change.
Algorithm (final, per Phase 1 design iteration):
cols = physical_cols (PE x-positions + relay cols for max_spacing)
rows_per_half = ceil(n_connections / 2)
total_rows = rows_per_half * 2 + 2 (+ 2 HBM rows)
PEs: 1 PE per row when rows available, corners at fixed positions
Hot path: min_connections = max(n_connections, 2)
"""
from __future__ import annotations
import hashlib
import math
from pathlib import Path
from typing import Any
import yaml
# ── Public API ────────────────────────────────────────────────────────
def ensure_mesh_file(cube_spec: dict, mesh_path: Path) -> dict:
"""Generate cube_mesh.yaml if needed, return parsed mesh dict."""
source_hash = _compute_source_hash(cube_spec)
if mesh_path.exists():
existing = yaml.safe_load(mesh_path.read_text(encoding="utf-8"))
if existing and existing.get("source_hash") == source_hash:
return existing
mesh = _generate_mesh(cube_spec, source_hash)
mesh_path.write_text(
yaml.dump(mesh, default_flow_style=False, sort_keys=False),
encoding="utf-8",
)
return mesh
# ── Hash ──────────────────────────────────────────────────────────────
def _compute_source_hash(cube_spec: dict) -> str:
"""Hash relevant topology params that determine mesh layout."""
relevant = {
"geometry": cube_spec["geometry"],
"pe_layout": cube_spec["pe_layout"],
"ucie_n_connections": cube_spec["ucie"]["n_connections"],
}
raw = yaml.dump(relevant, sort_keys=True)
return hashlib.sha256(raw.encode()).hexdigest()[:16]
# ── Layout helpers ────────────────────────────────────────────────────
def _corner_pe_positions(
cube_w: float, cube_h: float
) -> dict[str, list[tuple[float, float]]]:
"""PE center positions per corner, relative to cube origin."""
return {
"NW": [(1.5, 1.5), (4.5, 1.5)],
"NE": [(cube_w - 4.5, 1.5), (cube_w - 1.5, 1.5)],
"SW": [(1.5, cube_h - 1.5), (4.5, cube_h - 1.5)],
"SE": [(cube_w - 4.5, cube_h - 1.5), (cube_w - 1.5, cube_h - 1.5)],
}
def _compute_col_positions(cube_w: float, pe_positions: dict) -> list[float]:
"""Compute X positions for grid columns based on PE positions + relay spacing."""
xs: set[float] = set()
for positions in pe_positions.values():
for x, _y in positions:
xs.add(x)
sorted_xs = sorted(xs)
# Insert relay columns for gaps > max_spacing (3mm)
max_spacing = 3.0
result: list[float] = []
for i, x in enumerate(sorted_xs):
if i > 0:
gap = x - result[-1]
while gap > max_spacing + 0.01:
mid = result[-1] + max_spacing
if mid < x - 0.5:
result.append(round(mid, 1))
gap = x - result[-1]
else:
break
result.append(x)
return result
def _compute_row_positions(
cube_h: float, n_connections: int, pe_positions: dict
) -> tuple[list[float], int]:
"""Compute Y positions for grid rows.
Returns (y_positions, rows_per_half).
Layout: [top PE rows] [HBM row top] [HBM row bot] [bottom PE rows]
"""
n_conn = max(n_connections, 2) # hot path minimum
rows_per_half = math.ceil(n_conn / 2)
# Top half: evenly spaced from top PE y to just above HBM zone
top_pe_y = 1.5
hbm_top_y = cube_h / 2 - 1.5 # ~5.5 for h=14
hbm_bot_y = cube_h / 2 + 1.5 # ~8.5 for h=14
bot_pe_y = cube_h - 1.5
top_rows: list[float] = []
if rows_per_half == 1:
top_rows = [top_pe_y]
else:
step = (hbm_top_y - top_pe_y) / (rows_per_half - 1) if rows_per_half > 1 else 0
for i in range(rows_per_half):
top_rows.append(round(top_pe_y + i * step, 1))
# HBM rows
hbm_rows = [round(hbm_top_y, 1), round(hbm_bot_y, 1)]
# Bottom half: mirror of top
bot_rows: list[float] = []
if rows_per_half == 1:
bot_rows = [bot_pe_y]
else:
step = (bot_pe_y - hbm_bot_y) / (rows_per_half - 1) if rows_per_half > 1 else 0
for i in range(rows_per_half):
bot_rows.append(round(hbm_bot_y + i * step, 1))
return top_rows + hbm_rows + bot_rows, rows_per_half
# ── Mesh generation ──────────────────────────────────────────────────
def _generate_mesh(cube_spec: dict, source_hash: str) -> dict:
geom = cube_spec["geometry"]
cube_w = geom["cube_mm"]["w"]
cube_h = geom["cube_mm"]["h"]
pe_layout = cube_spec["pe_layout"]
corners = pe_layout["corners"]
pe_per_corner = pe_layout["pe_per_corner"]
n_connections = cube_spec["ucie"]["n_connections"]
pe_positions = _corner_pe_positions(cube_w, cube_h)
col_xs = _compute_col_positions(cube_w, pe_positions)
row_ys, rows_per_half = _compute_row_positions(
cube_h, n_connections, pe_positions
)
n_rows = len(row_ys)
n_cols = len(col_xs)
# HBM exclusion zone: center rows, center cols
hbm_row_start = rows_per_half # first HBM row index
hbm_row_end = rows_per_half + 1 # last HBM row index (inclusive)
hbm_col_start = n_cols // 2 - 1 # center-left col
hbm_col_end = n_cols // 2 # center-right col
# Build routers dict
routers: dict[str, Any] = {}
for r in range(n_rows):
for c in range(n_cols):
key = f"r{r}c{c}"
if (hbm_row_start <= r <= hbm_row_end
and hbm_col_start <= c <= hbm_col_end):
routers[key] = None # HBM excluded
else:
routers[key] = {
"pos_mm": [col_xs[c], row_ys[r]],
"attach": [],
}
# PE assignment: map each PE to a router based on corner and position.
# All PEs in the same corner share one row. Corner order determines row:
# Top half: NW → row 0, NE → row 1
# Bottom half: SW → row 4, SE → row 5 (for rows_per_half=2)
pe_idx = 0
top_pe_routers: list[str] = []
bot_pe_routers: list[str] = []
top_corners = [c for c in corners if c in ("NW", "NE")]
bot_corners = [c for c in corners if c in ("SW", "SE")]
for corner in corners:
is_top = corner in ("NW", "NE")
if is_top:
corner_idx = top_corners.index(corner)
row = corner_idx if corner_idx < rows_per_half else rows_per_half - 1
else:
corner_idx = bot_corners.index(corner)
bot_start = hbm_row_end + 1
row = bot_start + corner_idx if (bot_start + corner_idx) < n_rows else n_rows - 1
for ci in range(pe_per_corner):
pe_x, _pe_y = pe_positions[corner][ci]
col = min(range(n_cols), key=lambda c: abs(col_xs[c] - pe_x))
key = f"r{row}c{col}"
router = routers[key]
if router is not None:
router["attach"].append(f"pe{pe_idx}.dma")
router["attach"].append(f"pe{pe_idx}.cpu")
if is_top:
top_pe_routers.append(key)
else:
bot_pe_routers.append(key)
pe_idx += 1
# M_CPU and SRAM attachments (HBM row, leftmost available)
mcpu_key = f"r{hbm_row_start}c0"
if routers.get(mcpu_key) is not None:
routers[mcpu_key]["attach"].append("m_cpu")
sram_key = f"r{hbm_row_end}c0"
if routers.get(sram_key) is not None:
routers[sram_key]["attach"].append("sram")
# UCIe PE rows: top-half rows + bottom-half rows (1 per PE row)
ucie_pe_rows = []
for r in range(rows_per_half):
ucie_pe_rows.append(r)
for r in range(rows_per_half):
ucie_pe_rows.append(hbm_row_end + 1 + r)
# UCIe-E distribution: 1 per PE row, rightmost column
for i, row in enumerate(ucie_pe_rows):
key = f"r{row}c{n_cols - 1}"
router = routers.get(key)
if router is not None:
router["attach"].append(f"ucie_e.c{i}")
# UCIe-W distribution: 1 per PE row, leftmost column (mirror of E)
for i, row in enumerate(ucie_pe_rows):
key = f"r{row}c0"
router = routers.get(key)
if router is not None:
router["attach"].append(f"ucie_w.c{i}")
# UCIe PE columns: left-half + right-half PE columns (for N/S distribution)
pe_xs = set()
for positions in pe_positions.values():
for x, _y in positions:
pe_xs.add(x)
left_pe_cols = sorted(c for c in range(n_cols)
if col_xs[c] in pe_xs and c < hbm_col_start)
right_pe_cols = sorted(c for c in range(n_cols)
if col_xs[c] in pe_xs and c > hbm_col_end)
n_ucie = len(ucie_pe_rows)
half_n = n_ucie // 2
ucie_pe_cols = left_pe_cols[:half_n] + right_pe_cols[:n_ucie - half_n]
# UCIe-N distribution: PE columns on top row (row 0)
for i, col in enumerate(ucie_pe_cols):
key = f"r0c{col}"
router = routers.get(key)
if router is not None:
router["attach"].append(f"ucie_n.c{i}")
# UCIe-S distribution: PE columns on bottom row (row n_rows-1)
for i, col in enumerate(ucie_pe_cols):
key = f"r{n_rows - 1}c{col}"
router = routers.get(key)
if router is not None:
router["attach"].append(f"ucie_s.c{i}")
return {
"source_hash": source_hash,
"mesh": {
"rows": n_rows,
"cols": n_cols,
},
"routers": routers,
"xbar": {
"top": {"routers": sorted(set(top_pe_routers))},
"bottom": {"routers": sorted(set(bot_pe_routers))},
},
}