commit - release 1

2026-03-18 11:47:48 -07:00
commit 6f43807900
109 changed files with 14909 additions and 0 deletions
@@ -0,0 +1,409 @@
+from pathlib import Path
+
+from kernbench.topology.builder import load_topology
+
+TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml"
+
+
+def _graph():
+    return load_topology(TOPOLOGY_PATH)
+
+
+# ── Full graph: node counts ──────────────────────────────────────────
+
+
+def test_full_graph_node_count():
+    g = _graph()
+    # 1 switch
+    # + 2 SIPs × (1 IO × 2 comps + 16 cubes × (cube_comps + 8 PEs × 6 pe_comps))
+    #   cube_comps: 9 (noc, m_cpu, sram, 2 bridge, 4 ucie)
+    #             + 8 xbar.pe{0..7}  [replaced xbar.top/xbar.bottom]
+    #             + 8 hbm_slices = 25
+    # = 1 + 2*(2 + 16*(25+48)) = 1 + 2*(2+1168) = 1 + 2340 = 2341
+    assert len(g.nodes) == 2341
+
+
+def test_full_graph_edge_count():
+    g = _graph()
+    # Per cube: 144 (88 cube-fabric + 56 PE-internal)
+    #   cube-fabric: 8 pe→xbar.pe + 8 pe→noc + 8 noc→pe_cpu
+    #     + 8 xbar.pe→slice + 8 slice→xbar.pe (bidirectional for response)
+    #     + 12 xbar chain (3 pairs × 2 dir × 2 halves)
+    #     + 8 xbar.pe↔bridge (pe0↔bL, pe4↔bL, pe3↔bR, pe7↔bR, ×2 dir each)
+    #     + 4 noc→ucie + 4 ucie→noc (bidirectional)
+    #     + 8 noc→xbar.pe + 8 xbar.pe→noc (bidirectional for response)
+    #     + 1 m_cpu→noc + 1 noc→m_cpu + 1 noc→sram + 1 sram→noc = 88
+    # Per SIP: 16*144 + 48 inter-cube(bidirectional) + 8 io↔cube(bidirectional)
+    #          + 1 io_internal + 1 switch→io = 2362
+    # Total: 2 * 2362 = 4724
+    assert len(g.edges) == 4724
+
+
+# ── Full graph: specific nodes exist ─────────────────────────────────
+
+
+def test_system_switch_exists():
+    g = _graph()
+    assert "fabric.switch0" in g.nodes
+    assert g.nodes["fabric.switch0"].kind == "switch"
+    assert g.nodes["fabric.switch0"].pos_mm is None  # abstract
+
+
+def test_io_chiplet_nodes_exist():
+    g = _graph()
+    for s in range(2):
+        assert f"sip{s}.io0.pcie_ep" in g.nodes
+        assert f"sip{s}.io0.io_cpu" in g.nodes
+
+
+def test_cube_component_nodes_exist():
+    g = _graph()
+    cp = "sip0.cube0"
+    for name in ("noc", "m_cpu",
+                  "bridge.left", "bridge.right",
+                  "ucie-N", "ucie-S", "ucie-E", "ucie-W",
+                  "sram"):
+        assert f"{cp}.{name}" in g.nodes
+    # xbar.top/xbar.bottom replaced by per-PE xbar entry nodes
+    assert "sip0.cube0.xbar.top" not in g.nodes
+    assert "sip0.cube0.xbar.bottom" not in g.nodes
+    for pe in range(8):
+        node_id = f"{cp}.xbar.pe{pe}"
+        assert node_id in g.nodes, f"{node_id} missing"
+        assert g.nodes[node_id].kind == "xbar"
+    # HBM slices (one per PE)
+    for s in range(8):
+        assert f"{cp}.hbm_ctrl.slice{s}" in g.nodes
+        assert g.nodes[f"{cp}.hbm_ctrl.slice{s}"].kind == "hbm_ctrl"
+
+
+def test_pe_component_nodes_exist():
+    g = _graph()
+    for comp in ("pe_cpu", "pe_scheduler", "pe_dma", "pe_gemm", "pe_math", "pe_tcm"):
+        assert f"sip0.cube0.pe0.{comp}" in g.nodes
+        assert f"sip1.cube15.pe7.{comp}" in g.nodes
+
+
+# ── Full graph: positions ────────────────────────────────────────────
+
+
+def test_hbm_ctrl_slices_at_cube_center():
+    g = _graph()
+    # cube0 origin = (0, 0), cx=8.5, cy=7.0, hbm_ctrl at (cx-2, cy)
+    # all slices share the same physical position
+    for s in range(8):
+        node = g.nodes[f"sip0.cube0.hbm_ctrl.slice{s}"]
+        assert node.pos_mm == (6.5, 7.0)
+
+
+def test_hbm_ctrl_slices_cube5_position():
+    g = _graph()
+    # cube5 = col=1, row=1 -> origin = (1*18, 1*15) = (18, 15)
+    # hbm_ctrl = (18 + 6.5, 15 + 7.0) = (24.5, 22.0)
+    node = g.nodes["sip0.cube5.hbm_ctrl.slice0"]
+    assert node.pos_mm == (24.5, 22.0)
+
+
+def test_ucie_ports_at_cube_edges():
+    g = _graph()
+    # cube0 origin = (0, 0), cube_w=17, cube_h=14
+    # UCIe nodes inset by half-size so edges touch boundary
+    assert g.nodes["sip0.cube0.ucie-N"].pos_mm == (8.5, 0.6)
+    assert g.nodes["sip0.cube0.ucie-S"].pos_mm == (8.5, 13.4)
+    assert g.nodes["sip0.cube0.ucie-W"].pos_mm == (1.0, 7.0)
+    assert g.nodes["sip0.cube0.ucie-E"].pos_mm == (16.0, 7.0)
+
+
+# ── Full graph: edges ────────────────────────────────────────────────
+
+
+def _edge_set(g):
+    return {(e.src, e.dst) for e in g.edges}
+
+
+def test_inter_cube_ucie_edges():
+    es = _edge_set(_graph())
+    # cube0 (0,0) E → cube1 (1,0) W
+    assert ("sip0.cube0.ucie-E", "sip0.cube1.ucie-W") in es
+    # cube0 (0,0) S → cube4 (0,1) N
+    assert ("sip0.cube0.ucie-S", "sip0.cube4.ucie-N") in es
+
+
+def test_io_to_cube_edges():
+    es = _edge_set(_graph())
+    # io0 connects to cubes (0,0)..(3,0) on N side
+    assert ("sip0.io0.io_cpu", "sip0.cube0.ucie-N") in es
+    assert ("sip0.io0.io_cpu", "sip0.cube3.ucie-N") in es
+
+
+def test_switch_to_io_edges():
+    es = _edge_set(_graph())
+    assert ("fabric.switch0", "sip0.io0.pcie_ep") in es
+    assert ("fabric.switch0", "sip1.io0.pcie_ep") in es
+
+
+def test_pe_to_xbar_edges():
+    es = _edge_set(_graph())
+    cp = "sip0.cube0"
+    # Each PE connects to its own xbar entry (per-PE chain model)
+    for pe in range(8):
+        assert (f"{cp}.pe{pe}.pe_dma", f"{cp}.xbar.pe{pe}") in es
+    # Old shared xbar.top/bottom edges must NOT exist
+    assert (f"{cp}.pe0.pe_dma", f"{cp}.xbar.top") not in es
+    assert (f"{cp}.pe4.pe_dma", f"{cp}.xbar.bottom") not in es
+
+
+def test_command_path_m_cpu_noc_pe_cpu():
+    es = _edge_set(_graph())
+    cp = "sip0.cube0"
+    # m_cpu ↔ noc (bidirectional)
+    assert (f"{cp}.m_cpu", f"{cp}.noc") in es
+    assert (f"{cp}.noc", f"{cp}.m_cpu") in es
+    # noc → pe_cpu for each PE
+    assert (f"{cp}.noc", f"{cp}.pe0.pe_cpu") in es
+    assert (f"{cp}.noc", f"{cp}.pe7.pe_cpu") in es
+
+
+def test_pe_internal_edges():
+    es = _edge_set(_graph())
+    pp = "sip0.cube0.pe0"
+    assert (f"{pp}.pe_cpu", f"{pp}.pe_scheduler") in es
+    assert (f"{pp}.pe_scheduler", f"{pp}.pe_dma") in es
+    assert (f"{pp}.pe_scheduler", f"{pp}.pe_gemm") in es
+    assert (f"{pp}.pe_scheduler", f"{pp}.pe_math") in es
+    assert (f"{pp}.pe_dma", f"{pp}.pe_tcm") in es
+    assert (f"{pp}.pe_gemm", f"{pp}.pe_tcm") in es
+    assert (f"{pp}.pe_math", f"{pp}.pe_tcm") in es
+
+
+def test_xbar_to_hbm_slice_edges():
+    """Each xbar.pe{i} connects only to its own (local) HBM slice."""
+    es = _edge_set(_graph())
+    cp = "sip0.cube0"
+    # xbar.pe_i -> slice_i only (local Y-direction access)
+    for pe in range(8):
+        assert (f"{cp}.xbar.pe{pe}", f"{cp}.hbm_ctrl.slice{pe}") in es
+    # Negative: xbar.pe_i must NOT directly connect to a different slice
+    assert (f"{cp}.xbar.pe0", f"{cp}.hbm_ctrl.slice1") not in es
+    assert (f"{cp}.xbar.pe0", f"{cp}.hbm_ctrl.slice4") not in es
+    assert (f"{cp}.xbar.pe4", f"{cp}.hbm_ctrl.slice0") not in es
+
+
+# ── Views: system ────────────────────────────────────────────────────
+
+
+def test_system_view_nodes():
+    v = _graph().system_view
+    assert "fabric.switch0" in v.nodes
+    assert "sip0" in v.nodes
+    assert "sip1" in v.nodes
+    assert "sip0.io0" in v.nodes
+    assert "sip1.io0" in v.nodes
+
+
+# ── Views: SIP ───────────────────────────────────────────────────────
+
+
+def test_sip_view_cube_count():
+    v = _graph().sip_view
+    cube_nodes = [n for n in v.nodes if n.startswith("cube")]
+    assert len(cube_nodes) == 16
+
+
+def test_sip_view_io_chiplets():
+    v = _graph().sip_view
+    assert "io0" in v.nodes
+
+
+def test_sip_view_cube_positions():
+    v = _graph().sip_view
+    # cube0 (0,0): center = (8.5, 6+7.0) = (8.5, 13.0)  [io_margin=6]
+    x, y = v.nodes["cube0"].pos_mm
+    assert x == 8.5
+    assert y == 13.0
+    # cube1 (1,0): center = (18+8.5, 13.0) = (26.5, 13.0)
+    x1, y1 = v.nodes["cube1"].pos_mm
+    assert x1 == 26.5
+    assert y1 == 13.0
+
+
+# ── Views: cube ──────────────────────────────────────────────────────
+
+
+def test_cube_view_has_all_components():
+    v = _graph().cube_view
+    expected = {"ucie-N", "ucie-S", "ucie-W", "ucie-E",
+                "m_cpu", "hbm_ctrl",
+                "bridge.left", "bridge.right", "noc", "sram",
+                "xbar.pe0", "xbar.pe1", "xbar.pe2", "xbar.pe3",
+                "xbar.pe4", "xbar.pe5", "xbar.pe6", "xbar.pe7",
+                "pe0", "pe1", "pe2", "pe3", "pe4", "pe5", "pe6", "pe7"}
+    assert set(v.nodes.keys()) == expected
+
+
+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["noc"].pos_mm == (10.5, 7.0)
+    assert v.width_mm == 17.0
+    assert v.height_mm == 14.0
+
+
+def test_cube_view_pe_corner_mapping():
+    v = _graph().cube_view
+    ves = {(e.src, e.dst) for e in v.edges}
+    # Each PE connects to its own xbar entry (chain model)
+    for i in range(8):
+        assert (f"pe{i}", f"xbar.pe{i}") in ves
+    # Old shared xbar.top/bottom mapping must not exist
+    assert ("pe0", "xbar.top") not in ves
+    assert ("pe4", "xbar.bottom") not in ves
+
+
+# ── Views: PE ────────────────────────────────────────────────────────
+
+
+def test_pe_view_has_all_components():
+    v = _graph().pe_view
+    assert set(v.nodes.keys()) == {
+        "pe_cpu", "pe_scheduler", "pe_dma", "pe_gemm", "pe_math", "pe_tcm"
+    }
+
+
+def test_pe_view_edges():
+    v = _graph().pe_view
+    ves = {(e.src, e.dst) for e in v.edges}
+    assert ("pe_cpu", "pe_scheduler") in ves
+    assert ("pe_scheduler", "pe_dma") in ves
+    assert ("pe_scheduler", "pe_gemm") in ves
+    assert ("pe_scheduler", "pe_math") in ves
+    assert ("pe_dma", "pe_tcm") in ves
+    assert ("pe_gemm", "pe_tcm") in ves
+    assert ("pe_math", "pe_tcm") in ves
+
+
+# ── SRAM ────────────────────────────────────────────────────────────
+
+
+def test_sram_node_exists():
+    g = _graph()
+    assert "sip0.cube0.sram" in g.nodes
+    assert g.nodes["sip0.cube0.sram"].kind == "sram"
+
+
+def test_noc_to_sram_edges():
+    es = _edge_set(_graph())
+    cp = "sip0.cube0"
+    assert (f"{cp}.noc", f"{cp}.sram") in es
+    assert (f"{cp}.sram", f"{cp}.noc") in es
+
+
+# ── PE_DMA → NOC (non-HBM data path) ───────────────────────────────
+
+
+def test_pe_dma_to_noc_edges():
+    es = _edge_set(_graph())
+    cp = "sip0.cube0"
+    for i in range(8):
+        assert (f"{cp}.pe{i}.pe_dma", f"{cp}.noc") in es
+
+
+# ── Bridge connects XBAR halves (not NOC) ──────────────────────────
+
+
+def test_bridge_connects_xbar_halves():
+    """bridge.left connects leftmost PE nodes (pe0 top, pe4 bottom).
+    bridge.right connects rightmost PE nodes (pe3 top, pe7 bottom)."""
+    es = _edge_set(_graph())
+    cp = "sip0.cube0"
+    # bridge.left ↔ pe0 (top-left) and pe4 (bottom-left)
+    assert (f"{cp}.xbar.pe0", f"{cp}.bridge.left") in es
+    assert (f"{cp}.bridge.left", f"{cp}.xbar.pe0") in es
+    assert (f"{cp}.xbar.pe4", f"{cp}.bridge.left") in es
+    assert (f"{cp}.bridge.left", f"{cp}.xbar.pe4") in es
+    # bridge.right ↔ pe3 (top-right) and pe7 (bottom-right)
+    assert (f"{cp}.xbar.pe3", f"{cp}.bridge.right") in es
+    assert (f"{cp}.bridge.right", f"{cp}.xbar.pe3") in es
+    assert (f"{cp}.xbar.pe7", f"{cp}.bridge.right") in es
+    assert (f"{cp}.bridge.right", f"{cp}.xbar.pe7") in es
+    # Old xbar.top/bottom ↔ bridge edges must NOT exist
+    assert (f"{cp}.xbar.top", f"{cp}.bridge.left") not in es
+    assert (f"{cp}.xbar.bottom", f"{cp}.bridge.left") not in es
+
+
+def test_no_bridge_to_noc_edges():
+    es = _edge_set(_graph())
+    cp = "sip0.cube0"
+    assert (f"{cp}.bridge.left", f"{cp}.noc") not in es
+    assert (f"{cp}.bridge.right", f"{cp}.noc") not in es
+
+
+# ── Cube view: new edges ────────────────────────────────────────────
+
+
+def test_cube_view_pe_to_noc():
+    v = _graph().cube_view
+    ves = {(e.src, e.dst) for e in v.edges}
+    for i in range(8):
+        assert (f"pe{i}", "noc") in ves
+
+
+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 ("noc", "sram") in ves
+    assert ("sram", "noc") in ves
+
+
+def test_cube_view_bridge_xbar():
+    v = _graph().cube_view
+    ves = {(e.src, e.dst) for e in v.edges}
+    # bridge.left connects pe0 (top-left) ↔ pe4 (bottom-left)
+    assert ("xbar.pe0", "bridge.left") in ves
+    assert ("bridge.left", "xbar.pe0") in ves
+    assert ("xbar.pe4", "bridge.left") in ves
+    assert ("bridge.left", "xbar.pe4") in ves
+    # bridge.right connects pe3 (top-right) ↔ pe7 (bottom-right)
+    assert ("xbar.pe3", "bridge.right") in ves
+    assert ("bridge.right", "xbar.pe3") in ves
+    assert ("xbar.pe7", "bridge.right") in ves
+    assert ("bridge.right", "xbar.pe7") in ves
+
+
+# ── Chain xbar: new topology edges ──────────────────────────────────
+
+
+def test_xbar_chain_edges():
+    """Adjacent xbar.pe nodes within each half are bidirectionally connected."""
+    es = _edge_set(_graph())
+    cp = "sip0.cube0"
+    # Top chain: pe0 ↔ pe1 ↔ pe2 ↔ pe3 (NW→NE direction)
+    for a, b in [(0, 1), (1, 2), (2, 3)]:
+        assert (f"{cp}.xbar.pe{a}", f"{cp}.xbar.pe{b}") in es, f"missing pe{a}→pe{b}"
+        assert (f"{cp}.xbar.pe{b}", f"{cp}.xbar.pe{a}") in es, f"missing pe{b}→pe{a}"
+    # Bottom chain: pe4 ↔ pe5 ↔ pe6 ↔ pe7
+    for a, b in [(4, 5), (5, 6), (6, 7)]:
+        assert (f"{cp}.xbar.pe{a}", f"{cp}.xbar.pe{b}") in es, f"missing pe{a}→pe{b}"
+        assert (f"{cp}.xbar.pe{b}", f"{cp}.xbar.pe{a}") in es, f"missing pe{b}→pe{a}"
+    # Negative: no cross-chain direct edges
+    assert (f"{cp}.xbar.pe0", f"{cp}.xbar.pe2") not in es
+    assert (f"{cp}.xbar.pe0", f"{cp}.xbar.pe4") not in es
+
+
+def test_ucie_noc_reverse_edges():
+    """UCIe ports must have reverse edges back to NOC (bidirectional)."""
+    es = _edge_set(_graph())
+    cp = "sip0.cube1"  # non-edge cube to avoid io-cube edges
+    for port in ("N", "S", "E", "W"):
+        assert (f"{cp}.ucie-{port}", f"{cp}.noc") in es, \
+            f"missing ucie-{port}->noc reverse edge"
+
+
+def test_noc_to_xbar_pe_edges():
+    """NOC connects to all xbar.pe nodes (for remote cube HBM access)."""
+    es = _edge_set(_graph())
+    cp = "sip0.cube0"
+    for pe in range(8):
+        assert (f"{cp}.noc", f"{cp}.xbar.pe{pe}") in es, \
+            f"missing noc->xbar.pe{pe}"