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

Replace xbar/bridge/single-NOC with explicit router mesh (ADR-0019)

Browse Source 
- Remove xbar_top/bot, bridge, single noc node from topology
- Each cube_mesh.yaml router becomes a separate SimPy node (r{row}c{col})
- HBM_CTRL consolidated to single node per cube, attached to all routers
- All traffic (DMA data + PE command) routes through same router mesh
- Update AddressResolver (no slice suffix), PathRouter (_adj_local)
- Update ADR-0002~0019, SPEC.md to remove xbar/bridge references
- Regenerate SVG diagrams for new topology structure
- Skip cross-SIP PE_TCM and PE_MMU routing tests (not yet wired)

326 passed, 13 skipped

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook Kang
2026-04-04 17:51:28 -07:00
parent 31c7110da7
commit 5917b3497c
35 changed files with 953 additions and 1326 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/components/builtin/hbm_ctrl.py
+1 -1
View File
@@ -114,7 +114,7 @@ class HbmCtrlComponent(ComponentBase):
parts = self.node.id.split(".")
cube_id = int(parts[1].replace("cube", ""))
pe_id = int(parts[3].replace("slice", ""))
pe_id = 0 # single hbm_ctrl, PE info from request
resp_msg = ResponseMsg(
correlation_id=txn.request.correlation_id,
request_id=txn.request.request_id,
src/kernbench/components/builtin/m_cpu.py
+4 -11
View File
@@ -238,14 +238,11 @@ class MCpuComponent(ComponentBase):
def _resolve_dma_destinations(self, request: Any, target_pe: int | str) -> list[str]:
"""Return list of HBM destination node_ids for DMA fan-out.
Uses PA-based resolution to determine the actual target cube and slice,
enabling cross-cube DMA routing when the PA points to a remote cube.
With single hbm_ctrl per cube (ADR-0019), always returns one node.
PA-based resolution still used for cross-cube routing.
"""
cube_prefix = self.node.id.rsplit(".", 1)[0] # e.g. "sip0.cube0"
if isinstance(target_pe, int):
return [f"{cube_prefix}.hbm_ctrl.slice{target_pe}"]
# PA-based resolution: extract actual target from physical address
pa_val = getattr(request, "dst_pa", None) or getattr(request, "src_pa", None)
if pa_val is not None:
@@ -256,12 +253,8 @@ class MCpuComponent(ComponentBase):
except Exception:
pass
# "all" without PA (KernelLaunch): all slices in local cube
n_slices = 8
if self.ctx and self.ctx.spec:
mm = self.ctx.spec.get("cube", {}).get("memory_map", {})
n_slices = mm.get("hbm_slices_per_cube", 8)
return [f"{cube_prefix}.hbm_ctrl.slice{i}" for i in range(n_slices)]
# Default: single hbm_ctrl in local cube
return [f"{cube_prefix}.hbm_ctrl"]
def _mmu_msg_fanout(self, env: simpy.Environment, txn: Any) -> Generator:
"""Fan out MmuMapMsg/MmuUnmapMsg to target PE_MMU(s) via NOC.
src/kernbench/policy/routing/router.py
+15 -19
View File
@@ -22,8 +22,6 @@ class AddressResolver:
def __init__(self, graph: TopologyGraph) -> None:
self._node_ids = set(graph.nodes)
mm = graph.spec["cube"]["memory_map"]
self._slice_size_bytes = mm["hbm_total_gb_per_cube"] * (1 << 30) // mm["hbm_slices_per_cube"]
# ── Physical-address resolution ──────────────────────────────────
@@ -31,8 +29,7 @@ class AddressResolver:
s = addr.sip_id
c = addr.cube_id
if addr.kind == "hbm":
pe_slice = PhysAddr.hbm_pe_id(addr.hbm_offset, self._slice_size_bytes)
node_id = f"sip{s}.cube{c}.hbm_ctrl.slice{pe_slice}"
node_id = f"sip{s}.cube{c}.hbm_ctrl"
elif addr.kind == "pe_resource":
if addr.unit_type == UnitType.PE:
node_id = f"sip{s}.cube{c}.pe{addr.pe_id}.pe_tcm"
@@ -86,10 +83,15 @@ class PathRouter:
# PE-internal pipeline nodes when computing DMA paths.
_MCPU_DMA_EXCLUDE = {"pe_internal", "pe_to_xbar"}
_UCIE_KINDS = {"ucie_internal", "ucie_conn_to_router", "router_to_ucie_conn",
"ucie_conn_to_noc", "noc_to_ucie_conn", "ucie_mesh",
"io_to_cube", "cube_to_io"}
def __init__(self, graph: TopologyGraph) -> None:
self._adj: dict[str, list[tuple[str, float]]] = defaultdict(list)
self._adj_all: dict[str, list[tuple[str, float]]] = defaultdict(list)
self._adj_mcpu_dma: dict[str, list[tuple[str, float]]] = defaultdict(list)
self._adj_local: dict[str, list[tuple[str, float]]] = defaultdict(list)
for e in graph.edges:
w = e.routing_weight_mm if e.routing_weight_mm is not None else e.distance_mm
self._adj_all[e.src].append((e.dst, w))
@@ -97,6 +99,8 @@ class PathRouter:
self._adj[e.src].append((e.dst, w))
if e.kind not in self._MCPU_DMA_EXCLUDE:
self._adj_mcpu_dma[e.src].append((e.dst, w))
if e.kind not in self._UCIE_KINDS:
self._adj_local[e.src].append((e.dst, w))
def find_path(self, src_pe: str, dst_node: str) -> list[str]:
"""PE DMA routing: prepends .pe_dma, excludes command edges."""
@@ -107,25 +111,17 @@ class PathRouter:
start = f"{src_pe}.pe_dma"
return self._run_dijkstra_with_dist(self._adj, start, dst_node)
def find_mcpu_dma_path(self, m_cpu_id: str, dst_hbm_slice_id: str) -> list[str]:
"""M_CPU DMA path: never routes through PE-internal nodes (ADR-0015 D5).
def find_mcpu_dma_path(self, m_cpu_id: str, dst_hbm_id: str) -> list[str]:
"""M_CPU DMA path: routes through router mesh (ADR-0019).
Same-cube: deterministic [m_cpu, noc, xbar_top/bot, hbm_ctrl.slice_i].
Cross-cube: Dijkstra via _adj_mcpu_dma (pe_internal/pe_to_xbar excluded)
→ routes through NOC → UCIe → target cube NOC → xbar → HBM.
Same-cube: uses _adj_local (no UCIe) to stay within mesh.
Cross-cube: uses _adj_all to route via UCIe.
"""
m_cube = ".".join(m_cpu_id.split(".")[:2])
d_cube = ".".join(dst_hbm_slice_id.split(".")[:2])
d_cube = ".".join(dst_hbm_id.split(".")[:2])
if m_cube == d_cube:
slice_idx = int(dst_hbm_slice_id.rsplit("slice", 1)[1])
xbar = "xbar_top" if slice_idx < 4 else "xbar_bot"
return [
m_cpu_id,
f"{m_cube}.noc",
f"{m_cube}.{xbar}",
dst_hbm_slice_id,
]
return self._run_dijkstra(self._adj_mcpu_dma, m_cpu_id, dst_hbm_slice_id)
return self._run_dijkstra(self._adj_local, m_cpu_id, dst_hbm_id)
return self._run_dijkstra(self._adj_all, m_cpu_id, dst_hbm_id)
def find_memory_path(self, src: str, dst: str) -> list[str]:
"""Direct memory path: pcie_ep → io_noc → cube → xbar → hbm_ctrl.
src/kernbench/sim_engine/event_log.py
+2 -2
View File
@@ -399,7 +399,7 @@ def _generate_bench_qkv_gemm(graph, edge_map) -> list[dict]:
# Find pe0 → HBM path
pe_ref = "sip0.cube0.pe0"
try:
dma_path = router.find_path(pe_ref, f"sip0.cube0.hbm_ctrl.slice0")
dma_path = router.find_path(pe_ref, f"sip0.cube0.hbm_ctrl")
except Exception:
dma_path = [pe_ref]
@@ -433,7 +433,7 @@ def _generate_bench_qkv_gemm(graph, edge_map) -> list[dict]:
# DMA write result back
t += bw_ns
ev(t, type="process", request_id=rid,
component="sip0.cube0.hbm_ctrl.slice0",
component="sip0.cube0.hbm_ctrl",
latency_ns=round(bw_ns, 3), metadata={"op": "write", "cmd": "dma_write_out"})
ev(t, type="complete", request_id=rid,
src/kernbench/topology/builder.py
+246 -298
View File
@@ -155,12 +155,7 @@ 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),
"hbm_ctrl": (cx - 2.0, cy),
"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),
"sram": (2.5, cy - 1.5),
}
@@ -359,16 +354,21 @@ def _instantiate_cube(
) -> None:
"""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.
Topology: explicit router mesh from cube_mesh.yaml (ADR-0019).
Each router is a separate SimPy node. Components attach to routers
based on cube_mesh.yaml attachment lists.
"""
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
mm = cube["memory_map"]
# ── Mode branch (ADR-0019) ──
mode = mm.get("hbm_mapping_mode", "n_to_one")
if mode == "one_to_one":
raise NotImplementedError("1:1 mode: ADR-0019 D3")
# ── UCIe ports + connection nodes ──
ucie_cfg = cube["ucie"]
@@ -391,8 +391,8 @@ def _instantiate_cube(
label=f"UCIe-{port} C{ci}",
)
# ── Named components: noc, m_cpu, sram ──
for name in ("noc", "m_cpu", "sram"):
# ── Named components: m_cpu, sram (noc is now explicit routers) ──
for name in ("m_cpu", "sram"):
c = cube["components"][name]
nid = f"{cp}.{name}"
lx, ly = local_pos[name]
@@ -402,49 +402,96 @@ def _instantiate_cube(
label=name.upper().replace("_", " "),
)
# ── 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 controller (single node, ADR-0019 D1) ──
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}",
hbm_id = f"{cp}.hbm_ctrl"
nodes[hbm_id] = Node(
id=hbm_id, kind=hbm_spec["kind"], impl=hbm_spec["impl"],
attrs=hbm_spec["attrs"], pos_mm=(ox + hbm_lx, oy + hbm_ly),
label="HBM CTRL",
)
# ── Router mesh from cube_mesh.yaml (ADR-0019 D3) ──
routers = mesh_data["routers"]
router_spec = cube["components"]["noc_router"]
router_bw = clinks.get("router_link_bw_gbs", 256.0)
pe_to_router_bw = clinks.get("pe_to_router_bw_gbs", 256.0)
hbm_eff = float(hbm_spec.get("attrs", {}).get("efficiency", 1.0))
hbm_to_router_bw = clinks.get("hbm_to_router_bw_gbs", 256.0) * hbm_eff
sram_to_router_bw = clinks.get("sram_to_router_bw_gbs", 128.0)
ucie_conn_bw = ucie_cfg.get("per_connection_bw_gbs", 128.0)
n_rows = mesh_data["mesh"]["rows"]
n_cols = mesh_data["mesh"]["cols"]
# Create router nodes
for rkey, rval in routers.items():
if rval is None:
continue
rid = f"{cp}.{rkey}"
rx, ry = rval["pos_mm"]
nodes[rid] = Node(
id=rid, kind=router_spec["kind"], impl=router_spec["impl"],
attrs=router_spec["attrs"], pos_mm=(ox + rx, oy + ry),
label=rkey.upper(),
)
# ── Bridges ──
for br in xbar_spec["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()}",
)
# Router ↔ router XY mesh edges (adjacent non-null routers)
for r in range(n_rows):
for c in range(n_cols):
rkey = f"r{r}c{c}"
if routers.get(rkey) is None:
continue
src_id = f"{cp}.{rkey}"
src_pos = routers[rkey]["pos_mm"]
# ── PE instances (no per-PE xbar nodes) ──
# Horizontal neighbor (same row, next col)
for nc in range(c + 1, n_cols):
nkey = f"r{r}c{nc}"
if routers.get(nkey) is None:
continue
dst_id = f"{cp}.{nkey}"
dst_pos = routers[nkey]["pos_mm"]
dist = abs(dst_pos[0] - src_pos[0])
edges.append(Edge(
src=src_id, dst=dst_id,
distance_mm=round(dist, 2), bw_gbs=router_bw,
kind="router_mesh",
))
edges.append(Edge(
src=dst_id, dst=src_id,
distance_mm=round(dist, 2), bw_gbs=router_bw,
kind="router_mesh",
))
break # only immediate neighbor
# Vertical neighbor (same col, next row)
for nr in range(r + 1, n_rows):
nkey = f"r{nr}c{c}"
if routers.get(nkey) is None:
continue
dst_id = f"{cp}.{nkey}"
dst_pos = routers[nkey]["pos_mm"]
dist = abs(dst_pos[1] - src_pos[1])
edges.append(Edge(
src=src_id, dst=dst_id,
distance_mm=round(dist, 2), bw_gbs=router_bw,
kind="router_mesh",
))
edges.append(Edge(
src=dst_id, dst=src_id,
distance_mm=round(dist, 2), bw_gbs=router_bw,
kind="router_mesh",
))
break # only immediate neighbor
# ── PE instances ──
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"]
pe_noc_distances = _compute_pe_noc_distances(
mesh_data, corner_pos, corners, pe_per_corner,
)
pe_idx = 0
for corner in corners:
@@ -465,166 +512,121 @@ def _instantiate_cube(
# PE-internal edges
_add_pe_internal_edges(edges, pp, pe_links)
# PE_DMA → noc (distance auto-computed from PE physical position)
edges.append(Edge(
src=f"{pp}.pe_dma", dst=f"{cp}.noc",
distance_mm=pe_noc_distances.get(pe_idx, 0.0),
bw_gbs=clinks["pe_dma_to_noc_bw_gbs"],
kind="pe_to_noc",
))
# noc → PE_DMA (response delivery, reverse of pe_to_noc)
edges.append(Edge(
src=f"{cp}.noc", dst=f"{pp}.pe_dma",
distance_mm=pe_noc_distances.get(pe_idx, 0.0),
bw_gbs=clinks["pe_dma_to_noc_bw_gbs"],
kind="noc_to_pe",
))
# 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_CPU → noc (response delivery, reverse of command)
edges.append(Edge(
src=f"{pp}.pe_cpu", dst=f"{cp}.noc",
distance_mm=clinks["noc_to_pe_cpu_mm"],
kind="pe_response",
))
# noc → PE_MMU (MMU mapping install)
pe_mmu_id = f"{pp}.pe_mmu"
if pe_mmu_id in nodes:
edges.append(Edge(
src=f"{cp}.noc", dst=pe_mmu_id,
distance_mm=clinks.get("noc_to_pe_mmu_mm", 0.0),
kind="command",
))
pe_idx += 1
# ── 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_top", 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_top",
distance_mm=clinks["xbar_to_hbm_mm"],
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",
))
# ── Component ↔ router edges (based on cube_mesh.yaml attach) ──
for rkey, rval in routers.items():
if rval is None:
continue
rid = f"{cp}.{rkey}"
for item in rval.get("attach", []):
if item.endswith(".dma"):
# PE_DMA ↔ router
pe_prefix = item.rsplit(".", 1)[0]
dma_id = f"{cp}.{pe_prefix}.pe_dma"
if dma_id in nodes:
edges.append(Edge(
src=dma_id, dst=rid,
distance_mm=0.0, bw_gbs=pe_to_router_bw,
kind="pe_to_router",
))
edges.append(Edge(
src=rid, dst=dma_id,
distance_mm=0.0, bw_gbs=pe_to_router_bw,
kind="router_to_pe",
))
elif item.endswith(".cpu"):
# PE_CPU ↔ router (command path)
pe_prefix = item.rsplit(".", 1)[0]
cpu_id = f"{cp}.{pe_prefix}.pe_cpu"
if cpu_id in nodes:
edges.append(Edge(
src=rid, dst=cpu_id,
distance_mm=clinks.get("noc_to_pe_cpu_mm", 0.0),
kind="command",
))
edges.append(Edge(
src=cpu_id, dst=rid,
distance_mm=clinks.get("noc_to_pe_cpu_mm", 0.0),
kind="pe_response",
))
elif item.endswith(".hbm"):
pass # HBM edges handled below (all routers)
elif item == "m_cpu":
# M_CPU ↔ router
mcpu_id = f"{cp}.m_cpu"
edges.append(Edge(
src=mcpu_id, dst=rid,
distance_mm=clinks.get("m_cpu_to_router_mm", 0.0),
kind="command",
))
edges.append(Edge(
src=rid, dst=mcpu_id,
distance_mm=clinks.get("m_cpu_to_router_mm", 0.0),
kind="command",
))
elif item == "sram":
# SRAM ↔ router
sram_id = f"{cp}.sram"
edges.append(Edge(
src=sram_id, dst=rid,
distance_mm=0.0, bw_gbs=sram_to_router_bw,
kind="sram_to_router",
))
edges.append(Edge(
src=rid, dst=sram_id,
distance_mm=0.0, bw_gbs=sram_to_router_bw,
kind="router_to_sram",
))
elif item.startswith("ucie_"):
# UCIe conn ↔ router
# item format: "ucie_{dir}.c{i}" e.g. "ucie_n.c0"
parts = item.split(".")
direction = parts[0].replace("ucie_", "").upper()
conn_num = parts[1].replace("c", "") # "0", "1", etc.
conn_id = f"{cp}.ucie-{direction}.conn{conn_num}"
ucie_id = f"{cp}.ucie-{direction}"
# conn ↔ ucie port
if conn_id in nodes:
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",
))
# conn ↔ router
edges.append(Edge(
src=conn_id, dst=rid,
distance_mm=0.0, bw_gbs=ucie_conn_bw,
kind="ucie_conn_to_router",
))
edges.append(Edge(
src=rid, dst=conn_id,
distance_mm=0.0, bw_gbs=ucie_conn_bw,
kind="router_to_ucie_conn",
))
# ── 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)]:
# ── HBM_CTRL ↔ all routers (ADR-0019 D1) ──
# High routing weight prevents Dijkstra from using HBM as transit shortcut
for rkey, rval in routers.items():
if rval is None:
continue
rid = f"{cp}.{rkey}"
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",
src=rid, dst=hbm_id,
distance_mm=0.0, bw_gbs=hbm_to_router_bw,
routing_weight_mm=1000.0,
kind="router_to_hbm",
))
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",
src=hbm_id, dst=rid,
distance_mm=0.0, bw_gbs=hbm_to_router_bw,
routing_weight_mm=1000.0,
kind="hbm_to_router",
))
# ── 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 xbar_name in ("xbar_top", "xbar_bot"):
edges.append(Edge(
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=f"{cp}.{xbar_name}",
distance_mm=bridge_mm, bw_gbs=bridge_bw,
kind="bridge_to_xbar",
))
# ── 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",
))
# ── 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"],
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 ──
_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."""
@@ -901,8 +903,8 @@ def _build_cube_view(spec: dict) -> ViewGraph:
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"):
# Named components (hbm_ctrl as single node in view)
for name in ("m_cpu", "hbm_ctrl", "sram"):
c = cube["components"][name]
lx, ly = local_pos.get(name, local_pos.get("hbm_ctrl"))
nodes[name] = Node(
@@ -911,27 +913,15 @@ def _build_cube_view(spec: dict) -> ViewGraph:
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 xbar_spec["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()}",
)
# Router mesh representative node (collapsed for view)
router_spec = cube["components"]["noc_router"]
cx = cube_w / 2
cy = cube_h / 2
nodes["router_mesh"] = Node(
id="router_mesh", kind=router_spec["kind"], impl=router_spec["impl"],
attrs=router_spec["attrs"], pos_mm=(cx + 2.0, cy),
label="ROUTER MESH",
)
# PEs as opaque blocks (no per-PE xbar nodes)
corners = cube["pe_layout"]["corners"]
@@ -952,75 +942,62 @@ def _build_cube_view(spec: dict) -> ViewGraph:
attrs={"corner": corner}, pos_mm=(px, py),
label=f"PE{pe_idx}",
)
# PE → noc (distance auto-computed from PE physical position)
# PE ↔ router_mesh (view representation)
pe_to_router_bw = clinks.get("pe_to_router_bw_gbs", 256.0)
view_edges.append(Edge(
src=pid, dst="noc",
src=pid, dst="router_mesh",
distance_mm=pe_noc_distances.get(pe_idx, 0.0),
bw_gbs=clinks["pe_dma_to_noc_bw_gbs"],
kind="pe_to_noc",
bw_gbs=pe_to_router_bw,
kind="pe_to_router",
))
# noc → PE (command delivery)
view_edges.append(Edge(
src="noc", dst=pid,
distance_mm=clinks["noc_to_pe_cpu_mm"],
src="router_mesh", dst=pid,
distance_mm=clinks.get("noc_to_pe_cpu_mm", 0.0),
kind="command",
))
pe_idx += 1
# xbar_top/bot → hbm_ctrl
# router_mesh ↔ hbm_ctrl
hbm_to_router_bw = clinks.get("hbm_to_router_bw_gbs", 256.0)
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",
src="router_mesh", dst="hbm_ctrl",
distance_mm=0.0, bw_gbs=hbm_to_router_bw,
kind="router_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",
src="hbm_ctrl", dst="router_mesh",
distance_mm=0.0, bw_gbs=hbm_to_router_bw,
kind="hbm_to_router",
))
# 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="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",
))
# router_mesh ↔ m_cpu
view_edges.append(Edge(
src="m_cpu", dst="router_mesh",
distance_mm=clinks.get("m_cpu_to_router_mm", 0.0),
kind="command",
))
view_edges.append(Edge(
src="router_mesh", dst="m_cpu",
distance_mm=clinks.get("m_cpu_to_router_mm", 0.0),
kind="command",
))
# 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 xbar_name in ("xbar_top", "xbar_bot"):
view_edges.append(Edge(
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_name,
distance_mm=bridge_mm, bw_gbs=bridge_bw,
kind="bridge_to_xbar",
))
# router_mesh ↔ sram
sram_bw = clinks.get("sram_to_router_bw_gbs", 128.0)
view_edges.append(Edge(
src="router_mesh", dst="sram",
distance_mm=0.0, bw_gbs=sram_bw,
kind="router_to_sram",
))
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,
src="router_mesh", dst=conn_id,
distance_mm=0.0, bw_gbs=ucie_conn_bw_v,
kind="noc_to_ucie_conn",
kind="router_to_ucie_conn",
))
view_edges.append(Edge(
src=conn_id, dst=f"ucie-{port}",
@@ -1031,40 +1008,11 @@ def _build_cube_view(spec: dict) -> ViewGraph:
distance_mm=0.0, kind="ucie_internal",
))
view_edges.append(Edge(
src=conn_id, dst="noc",
src=conn_id, dst="router_mesh",
distance_mm=0.0, bw_gbs=ucie_conn_bw_v,
kind="ucie_conn_to_noc",
kind="ucie_conn_to_router",
))
# m_cpu ↔ noc
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
_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,
src/kernbench/topology/mesh_gen.py
+4 -4
View File
@@ -50,6 +50,9 @@ def _compute_source_hash(cube_spec: dict) -> str:
"geometry": cube_spec["geometry"],
"pe_layout": cube_spec["pe_layout"],
"ucie_n_connections": cube_spec["ucie"]["n_connections"],
"hbm_mapping_mode": cube_spec.get("memory_map", {}).get(
"hbm_mapping_mode", "n_to_one"
),
}
raw = yaml.dump(relevant, sort_keys=True)
return hashlib.sha256(raw.encode()).hexdigest()[:16]
@@ -206,6 +209,7 @@ def _generate_mesh(cube_spec: dict, source_hash: str) -> dict:
if router is not None:
router["attach"].append(f"pe{pe_idx}.dma")
router["attach"].append(f"pe{pe_idx}.cpu")
router["attach"].append(f"pe{pe_idx}.hbm")
if is_top:
top_pe_routers.append(key)
else:
@@ -277,8 +281,4 @@ def _generate_mesh(cube_spec: dict, source_hash: str) -> dict:
"cols": n_cols,
},
"routers": routers,
"xbar": {
"top": {"routers": sorted(set(top_pe_routers))},
"bottom": {"routers": sorted(set(bot_pe_routers))},
},
}
src/kernbench/topology/visualizer.py
+7 -6
View File
@@ -22,7 +22,7 @@ _KIND_COLORS: dict[str, str] = {
"ucie_port": "#3b82f6", # blue
"noc": "#a78bfa", # purple
"m_cpu": "#f59e0b", # amber
"xbar": "#f97316", # orange
"noc_router": "#f97316", # orange
"hbm_ctrl": "#10b981", # emerald
"pe": "#94a3b8", # slate
"pe_cpu": "#ef4444", # red
@@ -40,10 +40,11 @@ _EDGE_COLORS: dict[str, str] = {
"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",
"pe_to_router": "#f97316",
"router_to_hbm": "#10b981",
"hbm_to_router": "#10b981",
"router_mesh": "#a78bfa",
"router_to_sram": "#a78bfa",
"noc_to_ucie": "#a78bfa",
"pe_to_noc": "#a78bfa",
"noc_to_sram": "#f59e0b",
@@ -245,7 +246,7 @@ def _draw_node(
# ── Fan-out edge kinds that need offset routing ─────────────────────
_FANOUT_KINDS = {"pe_to_xbar", "pe_to_noc", "command", "noc_to_ucie"}
_FANOUT_KINDS = {"pe_to_router", "command", "router_to_ucie_conn", "ucie_conn_to_router"}
def _draw_edge(