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

commit - release 1

Browse Source
This commit is contained in:
Yangwook
2026-03-18 11:47:48 -07:00
commit 6f43807900
109 changed files with 14909 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/topology/__init__.py
View File
src/kernbench/topology/builder.py
+965
View File
@@ -0,0 +1,965 @@
# kernbench/topology/builder.py
"""
Topology compiler: parses topology.yaml and produces a fully-instantiated
TopologyGraph with nodes, edges, and representative view projections.
"""
from __future__ import annotations
from pathlib import Path
from typing import Any
import yaml
from .types import Edge, Node, TopologyGraph, TopologyHandle, ViewGraph
# PE component offsets from PE center (small, intra-PE distances ~0.5mm)
_PE_COMP_OFFSETS = {
"pe_cpu": (-0.3, 0.0),
"pe_scheduler": (-0.15, 0.0),
"pe_dma": (0.0, -0.15),
"pe_gemm": (0.0, 0.0),
"pe_math": (0.0, 0.15),
"pe_tcm": (0.3, 0.0),
}
# ── Public API ───────────────────────────────────────────────────────
def resolve_topology(path_str: str) -> TopologyHandle:
"""Validate path and build compiled topology graph."""
p = Path(path_str).expanduser().resolve()
if not p.exists():
raise FileNotFoundError(f"Topology file not found: {p}")
if not p.is_file():
raise ValueError(f"Topology path is not a file: {p}")
graph = load_topology(p)
return TopologyHandle(path=p, topology_obj=graph)
def load_topology(path: Path) -> TopologyGraph:
"""Load topology spec from file and compile into a topology graph."""
spec = _read_spec(path)
_validate_spec(spec)
return _compile_graph(spec)
def _read_spec(path: Path) -> dict[str, Any]:
"""Read YAML topology spec file and return a dict."""
try:
with path.open("r", encoding="utf-8") as f:
data = yaml.safe_load(f)
except yaml.YAMLError as e:
msg = f"Failed to parse YAML topology: {path}"
mark = getattr(e, "problem_mark", None)
if mark is not None:
msg += f" (line {mark.line + 1}, column {mark.column + 1})"
raise ValueError(msg) from e
if data is None:
raise ValueError(f"Topology YAML is empty: {path}")
if not isinstance(data, dict):
raise ValueError(
f"Topology YAML root must be a mapping/dict: {path} (got {type(data).__name__})"
)
return data
def _validate_spec(spec: dict) -> None:
# TODO: schema validation
return
# ── Graph Compiler ───────────────────────────────────────────────────
def _compile_graph(spec: dict) -> TopologyGraph:
"""Build fully-instantiated flat graph + representative view projections."""
nodes: dict[str, Node] = {}
edges: list[Edge] = []
system = spec["system"]
sip_spec = spec["sip"]
cube_spec = spec["cube"]
mesh_w = sip_spec["cube_mesh"]["w"]
mesh_h = sip_spec["cube_mesh"]["h"]
cube_w = cube_spec["geometry"]["cube_mm"]["w"]
cube_h = cube_spec["geometry"]["cube_mm"]["h"]
seam = sip_spec["links"]["inter_cube_mesh"]["distance_mm_across_seam"]
stride_x = cube_w + seam
stride_y = cube_h + seam
# System-level
_instantiate_system(nodes, system)
# Per-SIP
for sip_id in range(system["sips"]["count"]):
sp = f"sip{sip_id}"
# IO chiplets
_instantiate_io_chiplets(
nodes, edges, sp, sip_spec,
cube_w, cube_h, mesh_w, mesh_h, seam,
)
# Cubes + PEs
for row in range(mesh_h):
for col in range(mesh_w):
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)
# Inter-cube UCIe mesh
_add_inter_cube_edges(edges, sp, mesh_w, mesh_h, sip_spec)
# IO → cube UCIe
_add_io_to_cube_edges(edges, sp, sip_spec, mesh_w)
# Switch → IO pcie_ep
_add_system_to_io_edges(edges, sp, sip_spec, system)
# Build views
return TopologyGraph(
spec=spec,
nodes=nodes,
edges=edges,
system_view=_build_system_view(spec),
sip_view=_build_sip_view(spec),
cube_view=_build_cube_view(spec),
pe_view=_build_pe_view(spec),
)
# ── Layout helpers ───────────────────────────────────────────────────
def _cube_local_positions(cube_w: float, cube_h: float) -> dict[str, tuple[float, float]]:
"""Cube-internal component positions relative to cube origin (0,0) at top-left."""
cx, cy = cube_w / 2, cube_h / 2
# UCIe node half-sizes (default 2.0×1.2mm) — inset so edges touch boundary
uh = 0.6 # half height
uw = 1.0 # half width
return {
"ucie-N": (cx, uh),
"ucie-S": (cx, cube_h - uh),
"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
"hbm_ctrl": (cx - 2.0, cy),
"xbar.bottom": (cx, cube_h - 3.5), # Y reference for bottom-half xbar.pe nodes
"bridge.left": (2.5, cy + 2.0),
"bridge.right": (cube_w - 2.5, cy + 2.0),
"noc": (cx + 2.0, cy),
"sram": (2.5, cy - 1.5),
}
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)],
}
# ── Instantiation: system ───────────────────────────────────────────
def _instantiate_system(nodes: dict[str, Node], system: dict) -> None:
"""Add system-level nodes (fabric switch)."""
sw = system["components"]["switch"]
sw_id = "fabric.switch0"
nodes[sw_id] = Node(
id=sw_id, kind=sw["kind"], impl=sw["impl"],
attrs=sw.get("attrs", {}), pos_mm=None, label="Switch",
)
# ── Instantiation: IO chiplets ──────────────────────────────────────
def _instantiate_io_chiplets(
nodes: dict[str, Node],
edges: list[Edge],
sp: str,
sip_spec: dict,
cube_w: float,
cube_h: float,
mesh_w: int,
mesh_h: int,
seam: float,
) -> None:
"""Add IO chiplet nodes and internal pcie_ep → io_cpu edges."""
io_spec = sip_spec["iochiplet"]
comp = io_spec["components"]
links = io_spec["links"]
mesh_total_w = mesh_w * cube_w + (mesh_w - 1) * seam
mesh_total_h = mesh_h * cube_h + (mesh_h - 1) * seam
for inst in io_spec["instances"]:
iid = inst["id"]
prefix = f"{sp}.{iid}"
side = inst["place"]["side"]
cx = mesh_total_w / 2
if side == "N":
pcie_y, cpu_y = -5.0, -3.0
else:
pcie_y, cpu_y = mesh_total_h + 5.0, mesh_total_h + 3.0
# pcie_ep
ep = comp["pcie_ep"]
ep_id = f"{prefix}.pcie_ep"
nodes[ep_id] = Node(
id=ep_id, kind=ep["kind"], impl=ep["impl"],
attrs=ep["attrs"], pos_mm=(cx, pcie_y), label="PCIe EP",
)
# io_cpu
cpu = comp["io_cpu"]
cpu_id = f"{prefix}.io_cpu"
nodes[cpu_id] = Node(
id=cpu_id, kind=cpu["kind"], impl=cpu["impl"],
attrs=cpu["attrs"], pos_mm=(cx, cpu_y), label="IO CPU",
)
# Internal edge
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"],
kind="io_internal",
))
# ── Instantiation: cube + PEs ───────────────────────────────────────
def _instantiate_cube(
nodes: dict[str, Node],
edges: list[Edge],
cp: str,
cube: dict,
origin: tuple[float, float],
) -> None:
"""Add all cube-internal nodes and edges, including PE instances."""
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"]
# ── UCIe ports ──
ucie_ns = cube["ucie"]["overhead_ns"]
for port in cube["ucie"]["ports"]:
pid = f"{cp}.ucie-{port}"
lx, ly = local_pos[f"ucie-{port}"]
nodes[pid] = Node(
id=pid, kind="ucie_port", impl="ucie_v1",
attrs={"overhead_ns": ucie_ns}, pos_mm=(ox + lx, oy + ly),
label=f"UCIe-{port}",
)
# ── Named components: noc, m_cpu, sram ──
for name in ("noc", "m_cpu", "sram"):
c = cube["components"][name]
nid = f"{cp}.{name}"
lx, ly = local_pos[name]
nodes[nid] = Node(
id=nid, kind=c["kind"], impl=c["impl"],
attrs=c["attrs"], pos_mm=(ox + lx, oy + ly),
label=name.upper().replace("_", " "),
)
# ── HBM controller slices (one per PE) ──
hbm_spec = cube["components"]["hbm_ctrl"]
hbm_lx, hbm_ly = local_pos["hbm_ctrl"]
for sl in range(n_slices):
sid = f"{cp}.hbm_ctrl.slice{sl}"
nodes[sid] = Node(
id=sid, kind=hbm_spec["kind"], impl=hbm_spec["impl"],
attrs=hbm_spec["attrs"], pos_mm=(ox + hbm_lx, oy + hbm_ly),
label=f"HBM SLICE{sl}",
)
# ── Bridges ──
for br in cube["components"]["xbar"]["bridges"]:
bname = br["id"]
nid = f"{cp}.bridge.{bname}"
lx, ly = local_pos[f"bridge.{bname}"]
nodes[nid] = Node(
id=nid, kind=br["kind"], impl=br["impl"],
attrs=br["attrs"], pos_mm=(ox + lx, oy + ly),
label=f"Bridge {bname.upper()}",
)
# ── PE instances + per-PE xbar entry 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_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}"
dx, dy = _PE_COMP_OFFSETS.get(comp_name, (0.0, 0.0))
nodes[cid] = Node(
id=cid, kind=comp_spec["kind"], impl=comp_spec["impl"],
attrs=comp_spec["attrs"],
pos_mm=(ox + pe_cx + dx, oy + pe_cy + dy),
label=comp_name.upper().replace("_", " "),
)
# 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.)
edges.append(Edge(
src=f"{pp}.pe_dma", dst=f"{cp}.noc",
distance_mm=clinks["pe_dma_to_noc_mm"],
bw_gbs=clinks["pe_dma_to_noc_bw_gbs"],
kind="pe_to_noc",
))
# noc → PE_CPU (command delivery)
edges.append(Edge(
src=f"{cp}.noc", dst=f"{pp}.pe_cpu",
distance_mm=clinks["noc_to_pe_cpu_mm"],
kind="command",
))
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):
edges.append(Edge(
src=f"{cp}.xbar.pe{i}", dst=f"{cp}.hbm_ctrl.slice{i}",
distance_mm=clinks["xbar_to_hbm_mm"],
bw_gbs=clinks["xbar_to_hbm_bw_gbs"],
kind="xbar_to_hbm",
))
edges.append(Edge(
src=f"{cp}.hbm_ctrl.slice{i}", dst=f"{cp}.xbar.pe{i}",
distance_mm=clinks["xbar_to_hbm_mm"],
bw_gbs=clinks["xbar_to_hbm_bw_gbs"],
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",
))
# 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)]:
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}"
edges.append(Edge(
src=xbar_node, dst=br_node,
distance_mm=clinks[br_mm_key],
bw_gbs=clinks["xbar_to_bridge_bw_gbs"],
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"],
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",
))
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)
edges.append(Edge(
src=f"{cp}.m_cpu", dst=f"{cp}.noc",
distance_mm=clinks["m_cpu_to_noc_mm"],
kind="command",
))
edges.append(Edge(
src=f"{cp}.noc", dst=f"{cp}.m_cpu",
distance_mm=clinks["m_cpu_to_noc_mm"],
kind="command",
))
# noc ↔ sram (shared SRAM access; per_connection_bw_gbs = 128 GB/s, n_connections = 4)
_noc_sram = clinks["noc_to_sram"]
edges.append(Edge(
src=f"{cp}.noc", dst=f"{cp}.sram",
distance_mm=clinks["noc_to_sram_mm"],
bw_gbs=_noc_sram["per_connection_bw_gbs"],
n_connections=_noc_sram["n_connections"],
kind="noc_to_sram",
))
edges.append(Edge(
src=f"{cp}.sram", dst=f"{cp}.noc",
distance_mm=clinks["noc_to_sram_mm"],
bw_gbs=_noc_sram["per_connection_bw_gbs"],
n_connections=_noc_sram["n_connections"],
kind="noc_to_sram",
))
def _add_pe_internal_edges(edges: list[Edge], pp: str, pe_links: dict) -> None:
"""Add PE-internal edges for a single PE instance."""
edges.append(Edge(
src=f"{pp}.pe_cpu", dst=f"{pp}.pe_scheduler",
distance_mm=pe_links["pe_cpu_to_scheduler_mm"],
kind="pe_internal",
))
for eng, key in [("pe_dma", "scheduler_to_dma_mm"),
("pe_gemm", "scheduler_to_gemm_mm"),
("pe_math", "scheduler_to_math_mm")]:
edges.append(Edge(
src=f"{pp}.pe_scheduler", dst=f"{pp}.{eng}",
distance_mm=pe_links[key],
kind="pe_internal",
))
for eng, mm_key, bw_key in [("pe_dma", "dma_to_tcm_mm", "dma_to_tcm_bw_gbs"),
("pe_gemm", "gemm_to_tcm_mm", "gemm_to_tcm_bw_gbs"),
("pe_math", "math_to_tcm_mm", "math_to_tcm_bw_gbs")]:
edges.append(Edge(
src=f"{pp}.{eng}", dst=f"{pp}.pe_tcm",
distance_mm=pe_links[mm_key],
bw_gbs=pe_links[bw_key],
kind="pe_internal",
))
# ── Inter-cube / IO / system edges ──────────────────────────────────
def _add_inter_cube_edges(
edges: list[Edge], sp: str, mesh_w: int, mesh_h: int, sip_spec: dict,
) -> None:
"""Add UCIe mesh edges between adjacent cubes within a SIP."""
mesh = sip_spec["links"]["inter_cube_mesh"]
bw = mesh["bw_gbs_per_ucie_phy"]
dist = mesh["distance_mm_across_seam"]
for row in range(mesh_h):
for col in range(mesh_w):
cid = row * mesh_w + col
if col + 1 < mesh_w:
nid = row * mesh_w + (col + 1)
edges.append(Edge(
src=f"{sp}.cube{cid}.ucie-E", dst=f"{sp}.cube{nid}.ucie-W",
distance_mm=dist, bw_gbs=bw, kind="ucie_mesh",
))
edges.append(Edge(
src=f"{sp}.cube{nid}.ucie-W", dst=f"{sp}.cube{cid}.ucie-E",
distance_mm=dist, bw_gbs=bw, kind="ucie_mesh",
))
if row + 1 < mesh_h:
nid = (row + 1) * mesh_w + col
edges.append(Edge(
src=f"{sp}.cube{cid}.ucie-S", dst=f"{sp}.cube{nid}.ucie-N",
distance_mm=dist, bw_gbs=bw, kind="ucie_mesh",
))
edges.append(Edge(
src=f"{sp}.cube{nid}.ucie-N", dst=f"{sp}.cube{cid}.ucie-S",
distance_mm=dist, bw_gbs=bw, kind="ucie_mesh",
))
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"]
for inst in sip_spec["iochiplet"]["instances"]:
iid = inst["id"]
io_cpu_id = f"{sp}.{iid}.io_cpu"
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}"
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,
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,
kind="cube_to_io",
))
def _add_system_to_io_edges(
edges: list[Edge], sp: str, sip_spec: dict, system: dict,
) -> None:
"""Add fabric switch → IO chiplet PCIe edges."""
sw_id = "fabric.switch0"
sys_link = system["links"]["io_ep_to_switch"]
for inst in sip_spec["iochiplet"]["instances"]:
pcie_ep_id = f"{sp}.{inst['id']}.pcie_ep"
edges.append(Edge(
src=sw_id, dst=pcie_ep_id,
distance_mm=sys_link["distance_mm"],
bw_gbs=sys_link["bw_gbs_per_ep"],
kind="pcie",
))
# ── View builders ────────────────────────────────────────────────────
def _build_system_view(spec: dict) -> ViewGraph:
"""System-level view: SIP blocks, IO chiplets, fabric switch."""
system = spec["system"]
sip_count = system["sips"]["count"]
sip_w, sip_h = 71.0, 59.0
gap = 30.0
canvas_w = sip_count * sip_w + (sip_count - 1) * gap
canvas_h = sip_h + 20.0
nodes: dict[str, Node] = {}
view_edges: list[Edge] = []
sw = system["components"]["switch"]
sw_id = "fabric.switch0"
nodes[sw_id] = Node(
id=sw_id, kind=sw["kind"], impl=sw["impl"],
attrs=sw.get("attrs", {}), pos_mm=(canvas_w / 2, 5.0), label="Fabric Switch",
)
for s in range(sip_count):
sx = s * (sip_w + gap)
sy = 20.0
sip_id = f"sip{s}"
nodes[sip_id] = Node(
id=sip_id, kind="sip", impl="",
attrs={"w_mm": sip_w, "h_mm": sip_h},
pos_mm=(sx + sip_w / 2, sy + sip_h / 2),
label=f"SIP {s}",
)
for inst in spec["sip"]["iochiplet"]["instances"]:
iid = inst["id"]
io_nid = f"{sip_id}.{iid}"
side = inst["place"]["side"]
iy = sy if side == "N" else sy + sip_h
nodes[io_nid] = Node(
id=io_nid, kind="iochiplet", impl="",
attrs={}, pos_mm=(sx + sip_w / 2, iy), label=f"IO {iid}",
)
view_edges.append(Edge(
src=sw_id, dst=io_nid,
distance_mm=system["links"]["io_ep_to_switch"]["distance_mm"],
bw_gbs=system["links"]["io_ep_to_switch"]["bw_gbs_per_ep"],
kind="pcie",
))
return ViewGraph(
name="system", nodes=nodes, edges=view_edges,
width_mm=canvas_w, height_mm=canvas_h,
)
def _build_sip_view(spec: dict) -> ViewGraph:
"""SIP-level view: cube mesh + IO chiplets (representative, sip0)."""
sip_spec = spec["sip"]
cube_spec = spec["cube"]
mesh_w = sip_spec["cube_mesh"]["w"]
mesh_h = sip_spec["cube_mesh"]["h"]
cube_w = cube_spec["geometry"]["cube_mm"]["w"]
cube_h = cube_spec["geometry"]["cube_mm"]["h"]
seam = sip_spec["links"]["inter_cube_mesh"]["distance_mm_across_seam"]
stride_x = cube_w + seam
stride_y = cube_h + seam
mesh_total_w = mesh_w * cube_w + (mesh_w - 1) * seam
mesh_total_h = mesh_h * cube_h + (mesh_h - 1) * seam
io_margin = 6.0
canvas_w = mesh_total_w
canvas_h = mesh_total_h + 2 * io_margin
nodes: dict[str, Node] = {}
view_edges: list[Edge] = []
# Cubes as opaque blocks
for row in range(mesh_h):
for col in range(mesh_w):
cid = row * mesh_w + col
cx = col * stride_x + cube_w / 2
cy = io_margin + row * stride_y + cube_h / 2
nid = f"cube{cid}"
nodes[nid] = Node(
id=nid, kind="cube", impl="",
attrs={"w_mm": cube_w, "h_mm": cube_h, "col": col, "row": row},
pos_mm=(cx, cy), label=f"CUBE ({col},{row})",
)
# Inter-cube mesh edges
mesh_link = sip_spec["links"]["inter_cube_mesh"]
for row in range(mesh_h):
for col in range(mesh_w):
cid = row * mesh_w + col
if col + 1 < mesh_w:
nid = row * mesh_w + (col + 1)
view_edges.append(Edge(
src=f"cube{cid}", dst=f"cube{nid}",
distance_mm=mesh_link["distance_mm_across_seam"],
bw_gbs=mesh_link["bw_gbs_per_ucie_phy"],
kind="ucie_mesh",
))
if row + 1 < mesh_h:
nid = (row + 1) * mesh_w + col
view_edges.append(Edge(
src=f"cube{cid}", dst=f"cube{nid}",
distance_mm=mesh_link["distance_mm_across_seam"],
bw_gbs=mesh_link["bw_gbs_per_ucie_phy"],
kind="ucie_mesh",
))
# IO chiplets
io_links = sip_spec["iochiplet"]["links"]
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
nodes[iid] = Node(
id=iid, kind="iochiplet", impl="",
attrs={}, pos_mm=(mesh_total_w / 2, iy), label=f"IO {iid}",
)
for port in inst["cube_ports"]:
cube_col, cube_row = port["cube"]["xy"]
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"],
kind="io_to_cube",
))
return ViewGraph(
name="sip", nodes=nodes, edges=view_edges,
width_mm=canvas_w, height_mm=canvas_h,
)
def _build_cube_view(spec: dict) -> ViewGraph:
"""Cube-level view: representative single cube, PEs as opaque blocks."""
cube = spec["cube"]
cube_w = cube["geometry"]["cube_mm"]["w"]
cube_h = cube["geometry"]["cube_mm"]["h"]
local_pos = _cube_local_positions(cube_w, cube_h)
clinks = cube["links"]
n_slices = cube["memory_map"]["hbm_slices_per_cube"]
nodes: dict[str, Node] = {}
view_edges: list[Edge] = []
# UCIe ports
for port in cube["ucie"]["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}",
)
# 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]
nodes[name] = Node(
id=name, kind=c["kind"], impl=c["impl"],
attrs=c["attrs"], pos_mm=(lx, ly),
label=name.upper().replace("_", " "),
)
# Bridges
for br in cube["components"]["xbar"]["bridges"]:
bname = br["id"]
bid = f"bridge.{bname}"
lx, ly = local_pos[bid]
nodes[bid] = Node(
id=bid, kind=br["kind"], impl=br["impl"],
attrs=br["attrs"], pos_mm=(lx, ly),
label=f"Bridge {bname.upper()}",
)
# PEs as opaque blocks + per-PE xbar entry 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]
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)
view_edges.append(Edge(
src=pid, dst="noc",
distance_mm=clinks["pe_dma_to_noc_mm"],
bw_gbs=clinks["pe_dma_to_noc_bw_gbs"],
kind="pe_to_noc",
))
# noc → PE (command delivery)
view_edges.append(Edge(
src="noc", dst=pid,
distance_mm=clinks["noc_to_pe_cpu_mm"],
kind="command",
))
pe_idx += 1
# Cube fabric edges
# xbar.pe_i → hbm_ctrl (single representative node in view)
for i in range(n_slices):
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",
))
# 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)]:
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}"
view_edges.append(Edge(
src=xbar_id, dst=br_id,
distance_mm=clinks[br_mm_key],
bw_gbs=clinks["xbar_to_bridge_bw_gbs"],
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"],
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",
))
# m_cpu ↔ noc (command dispatch, both directions)
view_edges.append(Edge(
src="m_cpu", dst="noc",
distance_mm=clinks["m_cpu_to_noc_mm"],
kind="command",
))
view_edges.append(Edge(
src="noc", dst="m_cpu",
distance_mm=clinks["m_cpu_to_noc_mm"],
kind="command",
))
# noc ↔ sram (shared SRAM access, bidirectional)
_noc_sram_v = clinks["noc_to_sram"]
view_edges.append(Edge(
src="noc", dst="sram",
distance_mm=clinks["noc_to_sram_mm"],
bw_gbs=_noc_sram_v["per_connection_bw_gbs"],
n_connections=_noc_sram_v["n_connections"],
kind="noc_to_sram",
))
view_edges.append(Edge(
src="sram", dst="noc",
distance_mm=clinks["noc_to_sram_mm"],
bw_gbs=_noc_sram_v["per_connection_bw_gbs"],
n_connections=_noc_sram_v["n_connections"],
kind="noc_to_sram",
))
return ViewGraph(
name="cube", nodes=nodes, edges=view_edges,
width_mm=cube_w, height_mm=cube_h,
)
def _build_pe_view(spec: dict) -> ViewGraph:
"""PE-level view: representative single PE with all template components."""
pe_tmpl = spec["cube"]["pe_template"]
pe_links = pe_tmpl["links"]
canvas_w, canvas_h = 12.0, 8.0
positions = {
"pe_cpu": (1.5, 4.0),
"pe_scheduler": (4.0, 4.0),
"pe_dma": (7.0, 1.5),
"pe_gemm": (7.0, 4.0),
"pe_math": (7.0, 6.5),
"pe_tcm": (10.0, 4.0),
}
nodes: dict[str, Node] = {}
view_edges: list[Edge] = []
for comp_name, comp_spec in pe_tmpl["components"].items():
px, py = positions[comp_name]
nodes[comp_name] = Node(
id=comp_name, kind=comp_spec["kind"], impl=comp_spec["impl"],
attrs=comp_spec["attrs"], pos_mm=(px, py),
label=comp_name.upper().replace("_", " "),
)
view_edges.append(Edge(
src="pe_cpu", dst="pe_scheduler",
distance_mm=pe_links["pe_cpu_to_scheduler_mm"],
kind="pe_internal",
))
for eng, key in [("pe_dma", "scheduler_to_dma_mm"),
("pe_gemm", "scheduler_to_gemm_mm"),
("pe_math", "scheduler_to_math_mm")]:
view_edges.append(Edge(
src="pe_scheduler", dst=eng,
distance_mm=pe_links[key],
kind="pe_internal",
))
for eng, mm_key, bw_key in [("pe_dma", "dma_to_tcm_mm", "dma_to_tcm_bw_gbs"),
("pe_gemm", "gemm_to_tcm_mm", "gemm_to_tcm_bw_gbs"),
("pe_math", "math_to_tcm_mm", "math_to_tcm_bw_gbs")]:
view_edges.append(Edge(
src=eng, dst="pe_tcm",
distance_mm=pe_links[mm_key],
bw_gbs=pe_links[bw_key],
kind="pe_internal",
))
return ViewGraph(
name="pe", nodes=nodes, edges=view_edges,
width_mm=canvas_w, height_mm=canvas_h,
)
src/kernbench/topology/graph.py
View File
src/kernbench/topology/projections/cube_view.py
View File
src/kernbench/topology/projections/pe_view.py
View File
src/kernbench/topology/projections/sip_view.py
View File
src/kernbench/topology/types.py
+56
View File
@@ -0,0 +1,56 @@
from __future__ import annotations
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
@dataclass
class Node:
id: str
kind: str
impl: str
attrs: dict[str, Any]
pos_mm: tuple[float, float] | None # (x_mm, y_mm); None for abstract nodes
label: str = ""
@dataclass
class Edge:
src: str # node id
dst: str # node id
distance_mm: float # physical wire delay distance (ns = distance_mm * ns_per_mm)
routing_weight_mm: float | None = None # Dijkstra cost; None → use distance_mm
bw_gbs: float | None = None
n_connections: int | None = None # multi-connection links; single request uses 1 connection
kind: str = "link"
@dataclass
class ViewGraph:
name: str # "system" | "sip" | "cube" | "pe"
nodes: dict[str, Node]
edges: list[Edge]
width_mm: float
height_mm: float
@dataclass
class TopologyGraph:
spec: dict[str, Any]
# Full instantiated flat graph (used by sim_engine)
nodes: dict[str, Node] = field(default_factory=dict)
edges: list[Edge] = field(default_factory=list)
# Representative view projections (used by visualizer)
system_view: ViewGraph | None = None
sip_view: ViewGraph | None = None
cube_view: ViewGraph | None = None
pe_view: ViewGraph | None = None
@dataclass(frozen=True)
class TopologyHandle:
path: Path
topology_obj: TopologyGraph | None # None until _compile_graph is implemented
src/kernbench/topology/visualizer.py
+367
View File
@@ -0,0 +1,367 @@
# kernbench/topology/visualizer.py
"""
SVG diagram generator for TopologyGraph views.
Produces mm-accurate, deterministic SVG files for each view level
(system, SIP, cube, PE) per ADR-0005 and ADR-0006.
"""
from __future__ import annotations
from pathlib import Path
from .types import Edge, Node, TopologyGraph, ViewGraph
# ── Color palette by component kind ─────────────────────────────────
_KIND_COLORS: dict[str, str] = {
"switch": "#6366f1", # indigo
"sip": "#e0e7ff", # light indigo
"iochiplet": "#0ea5e9", # sky blue
"pcie_ep": "#0ea5e9",
"io_cpu": "#0ea5e9",
"ucie_port": "#3b82f6", # blue
"noc": "#a78bfa", # purple
"m_cpu": "#f59e0b", # amber
"xbar": "#f97316", # orange
"hbm_ctrl": "#10b981", # emerald
"pe": "#94a3b8", # slate
"pe_cpu": "#ef4444", # red
"pe_scheduler": "#f59e0b", # amber
"pe_dma": "#3b82f6", # blue
"pe_gemm": "#8b5cf6", # violet
"pe_math": "#ec4899", # pink
"pe_tcm": "#10b981", # emerald
"sram": "#f59e0b", # amber
"cube": "#cbd5e1", # slate-300
}
_EDGE_COLORS: dict[str, str] = {
"pcie": "#6366f1",
"io_internal": "#0ea5e9",
"io_to_cube": "#0ea5e9",
"ucie_mesh": "#3b82f6",
"pe_to_xbar": "#f97316",
"xbar_to_hbm": "#10b981",
"xbar_to_bridge": "#a78bfa",
"bridge_to_xbar": "#a78bfa",
"noc_to_ucie": "#a78bfa",
"pe_to_noc": "#a78bfa",
"noc_to_sram": "#f59e0b",
"command": "#f59e0b",
"pe_internal": "#94a3b8",
}
# ── Node sizing ──────────────────────────────────────────────────────
_DEFAULT_NODE_W = 2.0 # mm
_DEFAULT_NODE_H = 1.2 # mm
_KIND_SIZE: dict[str, tuple[float, float]] = {
"sip": (60.0, 50.0),
"cube": (6.0, 4.0),
"iochiplet": (4.0, 1.5),
"switch": (5.0, 1.5),
}
# ── Public API ───────────────────────────────────────────────────────
def emit_diagrams(graph: TopologyGraph, out_dir: Path) -> list[Path]:
"""Generate SVG diagrams for all views. Returns list of created file paths."""
out_dir.mkdir(parents=True, exist_ok=True)
created: list[Path] = []
views = [
("system_view", graph.system_view),
("sip_view", graph.sip_view),
("cube_view", graph.cube_view),
("pe_view", graph.pe_view),
]
for name, view in views:
if view is None:
continue
svg = _render_view_svg(view)
path = out_dir / f"{name}.svg"
path.write_text(svg, encoding="utf-8")
created.append(path)
return created
# ── SVG rendering ────────────────────────────────────────────────────
def _render_view_svg(view: ViewGraph) -> str:
"""Render a ViewGraph to an SVG string."""
scale = _pick_scale(view)
pad = 40 # px padding
node_sizes = _compute_node_sizes(view, scale)
# Canvas size in px
w_px = int(view.width_mm * scale + 2 * pad)
h_px = int(view.height_mm * scale + 2 * pad)
parts: list[str] = []
parts.append(_svg_header(w_px, h_px, view.name))
# Background
parts.append(f' <rect width="{w_px}" height="{h_px}" fill="#f8fafc"/>')
# Title
parts.append(
f' <text x="{w_px // 2}" y="18" text-anchor="middle" '
f'font-family="monospace" font-size="14" font-weight="bold" fill="#1e293b">'
f'{view.name.upper()} VIEW</text>'
)
# Special: draw cube boundary + HBM block background in cube view
if view.name == "cube":
_draw_cube_boundary(parts, view, scale, pad)
_draw_hbm_block(parts, view, scale, pad)
# Edges (draw before nodes so nodes are on top)
# Track fan-out edges to assign per-edge offsets
fanout_counter: dict[str, int] = {}
for edge in view.edges:
if edge.src in view.nodes and edge.dst in view.nodes:
_draw_edge(parts, edge, view, node_sizes, scale, pad, fanout_counter)
# Nodes
for node in view.nodes.values():
_draw_node(parts, node, node_sizes, scale, pad)
parts.append("</svg>")
return "\n".join(parts)
def _pick_scale(view: ViewGraph) -> float:
"""Pixels per mm, chosen per view type."""
return {
"system": 4.0,
"sip": 8.0,
"cube": 28.0,
"pe": 35.0,
}.get(view.name, 10.0)
def _compute_node_sizes(
view: ViewGraph, scale: float,
) -> dict[str, tuple[float, float]]:
"""Returns (w_px, h_px) for each node."""
sizes: dict[str, tuple[float, float]] = {}
for nid, node in view.nodes.items():
w_mm, h_mm = _KIND_SIZE.get(node.kind, (_DEFAULT_NODE_W, _DEFAULT_NODE_H))
# For cube view, use smaller PE nodes
if view.name == "cube" and node.kind == "pe":
w_mm, h_mm = 1.8, 1.0
if view.name == "pe":
w_mm, h_mm = 2.5, 1.4
sizes[nid] = (w_mm * scale, h_mm * scale)
return sizes
def _svg_header(w: int, h: int, title: str) -> str:
return (
f'<svg xmlns="http://www.w3.org/2000/svg" '
f'width="{w}" height="{h}" viewBox="0 0 {w} {h}">\n'
f' <title>{title}</title>'
)
def _draw_cube_boundary(
parts: list[str], view: ViewGraph, scale: float, pad: int,
) -> None:
"""Draw the cube die outline as a dashed rectangle."""
bx = pad
by = pad
bw = view.width_mm * scale
bh = view.height_mm * scale
parts.append(
f' <rect x="{bx:.1f}" y="{by:.1f}" '
f'width="{bw:.1f}" height="{bh:.1f}" '
f'rx="6" fill="none" stroke="#475569" stroke-width="2" '
f'stroke-dasharray="8,4"/>'
)
def _draw_hbm_block(
parts: list[str], view: ViewGraph, scale: float, pad: int,
) -> None:
"""Draw HBM area as a filled rectangle in cube view."""
# HBM area: centered at (8.5, 7.0), size 9x5 -> x=[4.0,13.0], y=[4.5,9.5]
hbm_x = 4.0 * scale + pad
hbm_y = 4.5 * scale + pad
hbm_w = 9.0 * scale
hbm_h = 5.0 * scale
parts.append(
f' <rect x="{hbm_x:.1f}" y="{hbm_y:.1f}" '
f'width="{hbm_w:.1f}" height="{hbm_h:.1f}" '
f'rx="4" fill="#d1fae5" stroke="#10b981" stroke-width="1.5" '
f'stroke-dasharray="6,3" opacity="0.5"/>'
)
cx = 8.5 * scale + pad
cy = 8.5 * scale + pad
parts.append(
f' <text x="{cx:.1f}" y="{cy:.1f}" text-anchor="middle" '
f'font-family="monospace" font-size="11" fill="#047857" opacity="0.7">'
f'HBM</text>'
)
def _draw_node(
parts: list[str],
node: Node,
sizes: dict[str, tuple[float, float]],
scale: float,
pad: int,
) -> None:
"""Draw a single node as a rounded rectangle with label."""
if node.pos_mm is None:
return
px = node.pos_mm[0] * scale + pad
py = node.pos_mm[1] * scale + pad
w, h = sizes.get(node.id, (40, 24))
x = px - w / 2
y = py - h / 2
fill = _KIND_COLORS.get(node.kind, "#e2e8f0")
text_color = "#ffffff" if _is_dark(fill) else "#1e293b"
parts.append(
f' <rect x="{x:.1f}" y="{y:.1f}" width="{w:.1f}" height="{h:.1f}" '
f'rx="4" fill="{fill}" stroke="#475569" stroke-width="1"/>'
)
label = node.label or node.id
font_size = _label_font_size(w, label)
parts.append(
f' <text x="{px:.1f}" y="{py + 4:.1f}" text-anchor="middle" '
f'font-family="monospace" font-size="{font_size}" fill="{text_color}">'
f'{_escape(label)}</text>'
)
# ── Fan-out edge kinds that need offset routing ─────────────────────
_FANOUT_KINDS = {"pe_to_xbar", "pe_to_noc", "command", "noc_to_ucie"}
def _draw_edge(
parts: list[str],
edge: Edge,
view: ViewGraph,
sizes: dict[str, tuple[float, float]],
scale: float,
pad: int,
fanout_counter: dict[str, int],
) -> None:
"""Draw an edge with orthogonal (90-degree) routing for fan-out kinds."""
nodes = view.nodes
src_node = nodes[edge.src]
dst_node = nodes[edge.dst]
if src_node.pos_mm is None or dst_node.pos_mm is None:
return
x1 = src_node.pos_mm[0] * scale + pad
y1 = src_node.pos_mm[1] * scale + pad
x2 = dst_node.pos_mm[0] * scale + pad
y2 = dst_node.pos_mm[1] * scale + pad
color = _EDGE_COLORS.get(edge.kind, "#94a3b8")
width = "1.5" if edge.kind == "pe_internal" else "1"
opacity = "0.6" if edge.kind in ("command", "noc_to_ucie") else "0.8"
if edge.kind in _FANOUT_KINDS and view.name == "cube":
# Orthogonal routing: src→horizontal→vertical→dst with per-edge offset.
group_key = f"{edge.kind}:{edge.dst}"
idx = fanout_counter.get(group_key, 0)
fanout_counter[group_key] = idx + 1
# Route: go vertically from src to a staggered horizontal channel,
# then horizontally to dst x, then vertically to dst.
mid_y = (y1 + y2) / 2 + (idx - 1.5) * 10 # spread channels vertically
parts.append(
f' <polyline points="{x1:.1f},{y1:.1f} {x1:.1f},{mid_y:.1f} '
f'{x2:.1f},{mid_y:.1f} {x2:.1f},{y2:.1f}" '
f'fill="none" stroke="{color}" stroke-width="{width}" opacity="{opacity}"/>'
)
# Label on the horizontal segment
if edge.distance_mm > 0:
lx = (x1 + x2) / 2
label = f"{edge.distance_mm:.1f}mm"
if edge.bw_gbs:
label += f" {edge.bw_gbs:.0f}GB/s"
parts.append(
f' <text x="{lx:.1f}" y="{mid_y - 3:.1f}" text-anchor="middle" '
f'font-family="monospace" font-size="7" fill="#64748b">'
f'{label}</text>'
)
return
# Non-fanout: orthogonal L-bend
if abs(x2 - x1) > 1 and abs(y2 - y1) > 1:
# PE view: vertical-first for left→right edges (scheduler→engines),
# horizontal-first for right→right edges (engines→tcm)
if view.name == "pe":
if src_node.pos_mm[0] < view.width_mm / 2:
# Source in left half: vertical-first (scheduler fan-out)
parts.append(
f' <polyline points="{x1:.1f},{y1:.1f} {x1:.1f},{y2:.1f} {x2:.1f},{y2:.1f}" '
f'fill="none" stroke="{color}" stroke-width="{width}" opacity="{opacity}"/>'
)
else:
# Source in right half: horizontal-first (dma/math→tcm)
parts.append(
f' <polyline points="{x1:.1f},{y1:.1f} {x2:.1f},{y1:.1f} {x2:.1f},{y2:.1f}" '
f'fill="none" stroke="{color}" stroke-width="{width}" opacity="{opacity}"/>'
)
else:
parts.append(
f' <polyline points="{x1:.1f},{y1:.1f} {x2:.1f},{y1:.1f} {x2:.1f},{y2:.1f}" '
f'fill="none" stroke="{color}" stroke-width="{width}" opacity="{opacity}"/>'
)
else:
parts.append(
f' <line x1="{x1:.1f}" y1="{y1:.1f}" x2="{x2:.1f}" y2="{y2:.1f}" '
f'stroke="{color}" stroke-width="{width}" opacity="{opacity}"/>'
)
# Distance label at midpoint
if edge.distance_mm > 0:
mx = (x1 + x2) / 2
my = (y1 + y2) / 2
label = f"{edge.distance_mm:.1f}mm"
if edge.bw_gbs:
label += f" {edge.bw_gbs:.0f}GB/s"
parts.append(
f' <text x="{mx:.1f}" y="{my - 4:.1f}" text-anchor="middle" '
f'font-family="monospace" font-size="7" fill="#64748b">'
f'{label}</text>'
)
# ── Helpers ──────────────────────────────────────────────────────────
def _is_dark(hex_color: str) -> bool:
"""Check if a hex color is dark (for white text)."""
h = hex_color.lstrip("#")
r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)
return (r * 0.299 + g * 0.587 + b * 0.114) < 140
def _label_font_size(box_width: float, label: str) -> int:
"""Choose font size to fit label in box."""
char_w = len(label) * 7
if char_w > box_width * 0.9:
return max(7, int(box_width * 0.9 / len(label) * 1.4))
return 10
def _escape(text: str) -> str:
"""Escape XML special characters."""
return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")