diff --git a/docs/diagrams/cube_view.svg b/docs/diagrams/cube_view.svg index 26d366b..3743329 100644 --- a/docs/diagrams/cube_view.svg +++ b/docs/diagrams/cube_view.svg @@ -141,8 +141,11 @@ PE1 - + r0c2 + + M_CPU + r0c3 @@ -167,11 +170,8 @@ PE3 - + r2c0 - - M_CPU - r2c1 diff --git a/src/kernbench/topology/mesh_gen.py b/src/kernbench/topology/mesh_gen.py index 48018ed..c90b255 100644 --- a/src/kernbench/topology/mesh_gen.py +++ b/src/kernbench/topology/mesh_gen.py @@ -53,6 +53,7 @@ def _compute_source_hash(cube_spec: dict) -> str: "hbm_mapping_mode": cube_spec.get("memory_map", {}).get( "hbm_mapping_mode", "n_to_one" ), + "placement": cube_spec.get("placement", {}), } raw = yaml.dump(relevant, sort_keys=True) return hashlib.sha256(raw.encode()).hexdigest()[:16] @@ -221,13 +222,29 @@ def _generate_mesh(cube_spec: dict, source_hash: str) -> dict: pe_idx += 1 - # M_CPU and SRAM attachments (HBM row, leftmost available) - mcpu_key = f"r{hbm_row_start}c0" - if routers.get(mcpu_key) is not None: + # M_CPU and SRAM attachments: find nearest router to configured position + placement = cube_spec.get("placement", {}) + + def _nearest_router(target_mm: list[float]) -> str | None: + best_key, best_dist = None, float("inf") + for rk, rv in routers.items(): + if rv is None: + continue + rx, ry = rv["pos_mm"] + dist = math.sqrt((rx - target_mm[0]) ** 2 + (ry - target_mm[1]) ** 2) + if dist < best_dist: + best_dist = dist + best_key = rk + return best_key + + mcpu_pos = placement.get("m_cpu", {}).get("pos_mm", [1.5, 5.5]) + mcpu_key = _nearest_router(mcpu_pos) + if mcpu_key and routers.get(mcpu_key) is not None: routers[mcpu_key]["attach"].append("m_cpu") - sram_key = f"r{hbm_row_end}c0" - if routers.get(sram_key) is not None: + sram_pos = placement.get("sram", {}).get("pos_mm", [1.5, 8.5]) + sram_key = _nearest_router(sram_pos) + if sram_key and routers.get(sram_key) is not None: routers[sram_key]["attach"].append("sram") # UCIe PE rows: top-half rows + bottom-half rows (1 per PE row) diff --git a/tests/test_topology_compile.py b/tests/test_topology_compile.py index 8f0ae15..cb2e8fe 100644 --- a/tests/test_topology_compile.py +++ b/tests/test_topology_compile.py @@ -158,9 +158,9 @@ def test_pe_dma_to_router(): def test_command_path_m_cpu_router_pe_cpu(): es = _edge_set(_graph()) cp = "sip0.cube0" - # m_cpu <-> r2c0 (bidirectional command) - assert (f"{cp}.m_cpu", f"{cp}.r2c0") in es - assert (f"{cp}.r2c0", f"{cp}.m_cpu") in es + # m_cpu <-> r0c2 (bidirectional command) + assert (f"{cp}.m_cpu", f"{cp}.r0c2") in es + assert (f"{cp}.r0c2", f"{cp}.m_cpu") in es # router -> pe_cpu for each PE (command kind) assert (f"{cp}.r0c0", f"{cp}.pe0.pe_cpu") in es assert (f"{cp}.r5c5", f"{cp}.pe7.pe_cpu") in es @@ -416,8 +416,8 @@ def test_cube_view_hbm_router(): def test_cube_view_m_cpu_router(): - """Cube view: m_cpu connects to its router r2c0.""" + """Cube view: m_cpu connects to its router r0c2.""" v = _graph().cube_view ves = {(e.src, e.dst) for e in v.edges} - assert ("m_cpu", "r2c0") in ves - assert ("r2c0", "m_cpu") in ves + assert ("m_cpu", "r0c2") in ves + assert ("r0c2", "m_cpu") in ves diff --git a/topology.yaml b/topology.yaml index 64adf67..81777db 100644 --- a/topology.yaml +++ b/topology.yaml @@ -95,6 +95,11 @@ cube: hbm_ctrl: { kind: hbm_ctrl, impl: hbm_ctrl_v1, attrs: { capacity: 1, efficiency: 1.0 } } sram: { kind: sram, impl: sram_v1, attrs: { size_mb: 32, overhead_ns: 2.0 } } + # Physical placement of non-PE components (mm coordinates) + placement: + m_cpu: { pos_mm: [7.5, 2.0] } # top center area, near UCIe-N + sram: { pos_mm: [1.5, 9.0] } # left side, below HBM zone + ucie: decompose: true ports: [N, S, E, W]