From 91085733ba46e305d0ee01e653577766547d3ef9 Mon Sep 17 00:00:00 2001 From: Yangwook Kang Date: Sat, 4 Apr 2026 18:22:38 -0700 Subject: [PATCH] Show individual routers in cube_view SVG, fix row Y overlap - cube_view now renders all 32 router nodes from cube_mesh.yaml instead of collapsed "router_mesh" placeholder - Fix mesh_gen row Y position overlap (r1/r2 and r3/r4 had same Y) by adding hbm_gap spacing between PE rows and HBM zone - Add noc_router to visualizer KIND_SIZE for proper sizing - Update cube view tests for individual router nodes 339 passed Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/diagrams/cube_view.svg | 334 ++++++++++++++++++++------- src/kernbench/topology/builder.py | 182 +++++++++------ src/kernbench/topology/mesh_gen.py | 12 +- src/kernbench/topology/visualizer.py | 1 + tests/test_topology_compile.py | 49 ++-- 5 files changed, 395 insertions(+), 183 deletions(-) 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