diff --git a/docs/diagrams/cube_view.svg b/docs/diagrams/cube_view.svg
index a3d55a2..051a1ac 100644
--- a/docs/diagrams/cube_view.svg
+++ b/docs/diagrams/cube_view.svg
@@ -5,95 +5,201 @@
HBM
-
-
-
-
-
- 4.0mm 256GB/s
-
-
- 4.0mm 256GB/s
-
-
- 4.0mm 256GB/s
-
-
- 4.0mm 256GB/s
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
+
+ 3.0mm
+
+ 2.5mm
+
+ 3.0mm
+
+ 2.5mm
+
+ 3.0mm
+
+ 2.5mm
+
+ 2.0mm
+
+ 2.5mm
+
+ 3.0mm
+
+ 2.5mm
+
+ 2.5mm
+
+ 3.0mm
+
+ 1.5mm
+
+ 3.0mm
+
+ 1.5mm
+
+ 3.0mm
+
+ 6.0mm
+
+ 2.0mm
+
+ 6.0mm
+
+ 3.0mm
+
+ 1.5mm
+
+ 1.5mm
+
+ 3.0mm
+
+ 3.0mm
+
+ 8.0mm
+
+ 3.0mm
+
+ 3.0mm
+
+ 3.0mm
+
+ 3.0mm
+
+ 3.0mm
+
+ 1.5mm
+
+ 8.0mm
+
+ 1.5mm
+
+ 3.0mm
+
+ 1.5mm
+
+ 1.5mm
+
+ 3.0mm
+
+ 2.5mm
+
+ 3.0mm
+
+ 2.5mm
+
+ 3.0mm
+
+ 2.5mm
+
+ 2.0mm
+
+ 2.5mm
+
+ 3.0mm
+
+ 2.5mm
+
+ 2.5mm
+
+ 3.0mm
+
+ 3.0mm
+
+ 3.0mm
+
+ 2.0mm
+
+ 3.0mm
+
+
+
+
+
-
-
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
-
-
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
-
-
+
+
+
+
+
+
+
+
+
-
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
UCIe-N
@@ -140,8 +246,70 @@
HBM CTRL
SRAM
-
- ROUTER MESH
+
+ R0C0
+
+ R0C1
+
+ R0C2
+
+ R0C3
+
+ R0C4
+
+ R0C5
+
+ R1C0
+
+ R1C1
+
+ R1C2
+
+ R1C3
+
+ R1C4
+
+ R1C5
+
+ R2C0
+
+ R2C1
+
+ R2C4
+
+ R2C5
+
+ R3C0
+
+ R3C1
+
+ R3C4
+
+ R3C5
+
+ R4C0
+
+ R4C1
+
+ R4C2
+
+ R4C3
+
+ R4C4
+
+ R4C5
+
+ R5C0
+
+ R5C1
+
+ R5C2
+
+ R5C3
+
+ R5C4
+
+ R5C5
PE0
diff --git a/src/kernbench/topology/builder.py b/src/kernbench/topology/builder.py
index e6dc220..3b2297e 100644
--- a/src/kernbench/topology/builder.py
+++ b/src/kernbench/topology/builder.py
@@ -921,21 +921,26 @@ def _build_cube_view(spec: dict) -> ViewGraph:
label=name.upper().replace("_", " "),
)
- # 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",
- )
+ # Load mesh data early (needed for router nodes + PE distances)
+ mesh_data = spec.get("_mesh", {})
- # PEs as opaque blocks (no per-PE xbar nodes)
+ # Router nodes from cube_mesh.yaml (explicit in view)
+ router_spec = cube["components"]["noc_router"]
+ routers = mesh_data.get("routers", {})
+ for rkey, rval in routers.items():
+ if rval is None:
+ continue
+ rx, ry = rval["pos_mm"]
+ nodes[rkey] = Node(
+ id=rkey, kind=router_spec["kind"], impl=router_spec["impl"],
+ attrs=router_spec["attrs"], pos_mm=(rx, ry),
+ label=rkey.upper(),
+ )
+
+ # PEs as opaque blocks
corners = cube["pe_layout"]["corners"]
pe_per_corner = cube["pe_layout"]["pe_per_corner"]
corner_pos = _corner_pe_positions(cube_w, cube_h)
- mesh_data = spec.get("_mesh", {})
pe_noc_distances = _compute_pe_noc_distances(
mesh_data, corner_pos, corners, pe_per_corner,
) if mesh_data else {}
@@ -950,76 +955,101 @@ def _build_cube_view(spec: dict) -> ViewGraph:
attrs={"corner": corner}, pos_mm=(px, py),
label=f"PE{pe_idx}",
)
- # 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="router_mesh",
- distance_mm=pe_noc_distances.get(pe_idx, 0.0),
- bw_gbs=pe_to_router_bw,
- kind="pe_to_router",
- ))
- view_edges.append(Edge(
- src="router_mesh", dst=pid,
- distance_mm=clinks.get("noc_to_pe_cpu_mm", 0.0),
- kind="command",
- ))
pe_idx += 1
- # router_mesh ↔ hbm_ctrl
+ # View edges based on cube_mesh.yaml attach (mirrors _instantiate_cube logic)
+ pe_to_router_bw = clinks.get("pe_to_router_bw_gbs", 256.0)
hbm_to_router_bw = clinks.get("hbm_to_router_bw_gbs", 256.0)
- view_edges.append(Edge(
- 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="hbm_ctrl", dst="router_mesh",
- distance_mm=0.0, bw_gbs=hbm_to_router_bw,
- kind="hbm_to_router",
- ))
-
- # 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",
- ))
-
- # 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="router_mesh", dst=conn_id,
- distance_mm=0.0, bw_gbs=ucie_conn_bw_v,
- kind="router_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="router_mesh",
- distance_mm=0.0, bw_gbs=ucie_conn_bw_v,
- kind="ucie_conn_to_router",
- ))
+ n_rows = mesh_data.get("mesh", {}).get("rows", 6)
+ n_cols = mesh_data.get("mesh", {}).get("cols", 6)
+
+ # Router ↔ router mesh edges
+ 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_pos = routers[rkey]["pos_mm"]
+ # Horizontal neighbor
+ for nc in range(c + 1, n_cols):
+ nkey = f"r{r}c{nc}"
+ if routers.get(nkey) is None:
+ continue
+ dist = abs(routers[nkey]["pos_mm"][0] - src_pos[0])
+ view_edges.append(Edge(
+ src=rkey, dst=nkey, distance_mm=round(dist, 2),
+ kind="router_mesh",
+ ))
+ break
+ # Vertical neighbor
+ for nr in range(r + 1, n_rows):
+ nkey = f"r{nr}c{c}"
+ if routers.get(nkey) is None:
+ continue
+ dist = abs(routers[nkey]["pos_mm"][1] - src_pos[1])
+ view_edges.append(Edge(
+ src=rkey, dst=nkey, distance_mm=round(dist, 2),
+ kind="router_mesh",
+ ))
+ break
+
+ # Component ↔ router edges from attach lists
+ for rkey, rval in routers.items():
+ if rval is None:
+ continue
+ for item in rval.get("attach", []):
+ if item.endswith(".dma"):
+ pe_prefix = item.rsplit(".", 1)[0]
+ pid = pe_prefix.replace("pe", "pe") # "pe0" → "pe0"
+ if pid in nodes:
+ view_edges.append(Edge(
+ src=pid, dst=rkey, distance_mm=0.0,
+ bw_gbs=pe_to_router_bw, kind="pe_to_router",
+ ))
+ view_edges.append(Edge(
+ src=rkey, dst=pid, distance_mm=0.0,
+ kind="command",
+ ))
+ elif item.endswith(".hbm"):
+ view_edges.append(Edge(
+ src=rkey, dst="hbm_ctrl", distance_mm=0.0,
+ bw_gbs=hbm_to_router_bw, kind="router_to_hbm",
+ ))
+ elif item == "m_cpu":
+ view_edges.append(Edge(
+ src="m_cpu", dst=rkey, distance_mm=0.0, kind="command",
+ ))
+ view_edges.append(Edge(
+ src=rkey, dst="m_cpu", distance_mm=0.0, kind="command",
+ ))
+ elif item == "sram":
+ view_edges.append(Edge(
+ src="sram", dst=rkey, distance_mm=0.0,
+ bw_gbs=sram_bw, kind="router_to_sram",
+ ))
+ elif item.startswith("ucie_"):
+ parts = item.split(".")
+ direction = parts[0].replace("ucie_", "").upper()
+ conn_num = parts[1].replace("c", "")
+ conn_id = f"ucie-{direction}.conn{conn_num}"
+ view_edges.append(Edge(
+ src=rkey, dst=conn_id, distance_mm=0.0,
+ bw_gbs=ucie_conn_bw_v, kind="router_to_ucie_conn",
+ ))
+ view_edges.append(Edge(
+ src=conn_id, dst=rkey, distance_mm=0.0,
+ bw_gbs=ucie_conn_bw_v, kind="ucie_conn_to_router",
+ ))
+ view_edges.append(Edge(
+ src=conn_id, dst=f"ucie-{direction}",
+ distance_mm=0.0, kind="ucie_internal",
+ ))
+ view_edges.append(Edge(
+ src=f"ucie-{direction}", dst=conn_id,
+ distance_mm=0.0, kind="ucie_internal",
+ ))
return ViewGraph(
name="cube", nodes=nodes, edges=view_edges,
diff --git a/src/kernbench/topology/mesh_gen.py b/src/kernbench/topology/mesh_gen.py
index 6b5cc72..48018ed 100644
--- a/src/kernbench/topology/mesh_gen.py
+++ b/src/kernbench/topology/mesh_gen.py
@@ -111,6 +111,7 @@ def _compute_row_positions(
# Top half: evenly spaced from top PE y to just above HBM zone
top_pe_y = 1.5
+ hbm_gap = 1.5 # minimum gap between PE rows and HBM rows
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
@@ -119,21 +120,24 @@ def _compute_row_positions(
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
+ # End before HBM zone with gap
+ top_end = hbm_top_y - hbm_gap
+ step = (top_end - 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
+ # Bottom half: mirror of top, start after HBM zone with gap
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
+ bot_start = hbm_bot_y + hbm_gap
+ step = (bot_pe_y - bot_start) / (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))
+ bot_rows.append(round(bot_start + i * step, 1))
return top_rows + hbm_rows + bot_rows, rows_per_half
diff --git a/src/kernbench/topology/visualizer.py b/src/kernbench/topology/visualizer.py
index 3df0138..c02b18f 100644
--- a/src/kernbench/topology/visualizer.py
+++ b/src/kernbench/topology/visualizer.py
@@ -62,6 +62,7 @@ _KIND_SIZE: dict[str, tuple[float, float]] = {
"cube": (6.0, 4.0),
"iochiplet": (4.0, 1.5),
"switch": (5.0, 1.5),
+ "noc_router": (1.2, 0.8),
}
diff --git a/tests/test_topology_compile.py b/tests/test_topology_compile.py
index 9845c0b..8f0ae15 100644
--- a/tests/test_topology_compile.py
+++ b/tests/test_topology_compile.py
@@ -247,8 +247,14 @@ def test_sip_view_cube_positions():
def test_cube_view_has_all_components():
v = _graph().cube_view
expected = {"ucie-N", "ucie-S", "ucie-W", "ucie-E",
- "m_cpu", "hbm_ctrl", "router_mesh", "sram",
- "pe0", "pe1", "pe2", "pe3", "pe4", "pe5", "pe6", "pe7"}
+ "m_cpu", "hbm_ctrl", "sram",
+ "pe0", "pe1", "pe2", "pe3", "pe4", "pe5", "pe6", "pe7",
+ "r0c0", "r0c1", "r0c2", "r0c3", "r0c4", "r0c5",
+ "r1c0", "r1c1", "r1c2", "r1c3", "r1c4", "r1c5",
+ "r2c0", "r2c1", "r2c4", "r2c5",
+ "r3c0", "r3c1", "r3c4", "r3c5",
+ "r4c0", "r4c1", "r4c2", "r4c3", "r4c4", "r4c5",
+ "r5c0", "r5c1", "r5c2", "r5c3", "r5c4", "r5c5"}
# Add UCIe connection nodes (4 ports x 4 connections)
for port in ("N", "S", "E", "W"):
for ci in range(4):
@@ -259,17 +265,19 @@ def test_cube_view_has_all_components():
def test_cube_view_hbm_at_center():
v = _graph().cube_view
assert v.nodes["hbm_ctrl"].pos_mm == (6.5, 7.0)
- assert v.nodes["router_mesh"].pos_mm == (10.5, 7.0)
+ assert "r0c0" in v.nodes # routers exist in cube view
assert v.width_mm == 17.0
assert v.height_mm == 14.0
-def test_cube_view_pe_to_router_mesh():
- """PEs connect to router_mesh in cube view."""
+def test_cube_view_pe_to_router():
+ """PEs connect to their assigned routers in cube view."""
v = _graph().cube_view
ves = {(e.src, e.dst) for e in v.edges}
- for i in range(8):
- assert (f"pe{i}", "router_mesh") in ves
+ pe_router_map = {"pe0": "r0c0", "pe1": "r0c1", "pe2": "r1c4", "pe3": "r1c5",
+ "pe4": "r4c0", "pe5": "r4c1", "pe6": "r5c4", "pe7": "r5c5"}
+ for pe, router in pe_router_map.items():
+ assert (pe, router) in ves, f"{pe} should connect to {router}"
# -- Views: PE ----------------------------------------------------------------
@@ -383,32 +391,33 @@ def test_cross_cube_path_includes_conn():
# -- Cube view: edges ---------------------------------------------------------
-def test_cube_view_pe_to_router_mesh_edges():
- """All PEs connect to router_mesh in cube view."""
+def test_cube_view_pe_to_router_edges():
+ """All PEs connect to their routers in cube view."""
v = _graph().cube_view
ves = {(e.src, e.dst) for e in v.edges}
- for i in range(8):
- assert (f"pe{i}", "router_mesh") in ves
+ pe_router_map = {"pe0": "r0c0", "pe1": "r0c1", "pe2": "r1c4", "pe3": "r1c5",
+ "pe4": "r4c0", "pe5": "r4c1", "pe6": "r5c4", "pe7": "r5c5"}
+ for pe, router in pe_router_map.items():
+ assert (pe, router) in ves, f"{pe} should connect to {router}"
def test_cube_view_sram():
v = _graph().cube_view
assert "sram" in v.nodes
ves = {(e.src, e.dst) for e in v.edges}
- assert ("router_mesh", "sram") in ves
+ assert ("sram", "r3c0") in ves
-def test_cube_view_hbm_router_mesh():
- """Cube view: hbm_ctrl connects to router_mesh."""
+def test_cube_view_hbm_router():
+ """Cube view: PE routers connect to hbm_ctrl."""
v = _graph().cube_view
ves = {(e.src, e.dst) for e in v.edges}
- assert ("router_mesh", "hbm_ctrl") in ves
- assert ("hbm_ctrl", "router_mesh") in ves
+ assert ("r0c0", "hbm_ctrl") in ves # PE0's router → HBM
-def test_cube_view_m_cpu_router_mesh():
- """Cube view: m_cpu connects to router_mesh."""
+def test_cube_view_m_cpu_router():
+ """Cube view: m_cpu connects to its router r2c0."""
v = _graph().cube_view
ves = {(e.src, e.dst) for e in v.edges}
- assert ("router_mesh", "m_cpu") in ves
- assert ("m_cpu", "router_mesh") in ves
+ assert ("m_cpu", "r2c0") in ves
+ assert ("r2c0", "m_cpu") in ves