From 6f61657ba46a956a28a6ef99f0f686231e7de09d Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Tue, 28 Jul 2026 22:07:50 -0700 Subject: [PATCH] =?UTF-8?q?analytical-viz:=20topology=20map=20=E2=80=94=20?= =?UTF-8?q?color=20PEs=20by=20CP=20rank=20when=20cp=5Fplacement=3Dpe?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The topology map previously colored a whole cube (border + all PE fills) by the single cp_rank assigned to that cube in the cube→(pp,cp) mapping. Under cp_placement=cube this is right (each cube = one CP rank). Under cp_placement=pe, however, multiple CP ranks are PACKED into the same cube's PEs, so the whole-cube coloring makes every PE look identical (defaulting to cp_rank=0's palette entry — bright red). Fix: in the per-PE loop, if cp_packed = (cp_placement=="pe" and cp>1), compute each PE's own cp_rank = pe_id // tp and look up its color via _cp_color(pp_stage, pe_cp_rank, cp_size). Border still uses the cube-level color (cp_rank=0), so the outer bounding box is unchanged, but the interior PEs now show all four (or however many) CP-rank colors visibly. For cp_placement=cube: unchanged (single pe_fill per cube). Verified with a headless render at Qwen 3 8B, CP=4, TP=2, cp_placement=pe: 145 patches → 8 distinct facecolors (4 for the CP ranks in the packed cube + inactive/border tints), where before it was fewer distinct colors. Co-Authored-By: Claude Opus 4.7 --- tests/analytical_visualization/topology_map.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/analytical_visualization/topology_map.py b/tests/analytical_visualization/topology_map.py index 53a974c..e37e629 100644 --- a/tests/analytical_visualization/topology_map.py +++ b/tests/analytical_visualization/topology_map.py @@ -258,12 +258,27 @@ def _draw_one_sip(ax, cfg: FullConfig, sip_idx: int, pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS pes_used = cfg.topo.pes_per_cube_used if is_used else 0 + # When cp_placement=pe, multiple CP ranks live inside this cube's PEs + # (cp_rank = pe_id // tp). Color each PE by its cp_rank so all four + # groups are visible; otherwise fall back to the whole-cube pe_fill. + _cp_packed = is_used and cfg.topo.cp_placement == "pe" and cfg.topo.cp > 1 for pr in range(PE_ROWS): for pc in range(PE_COLS): pe_id = pr * PE_COLS + pc px = x + pe_pad + pc * (pe_w + pe_gap) py = y + pe_pad + (PE_ROWS - 1 - pr) * (pe_h + pe_gap) - fill = pe_fill if pe_id < pes_used else _INACTIVE_PE + if pe_id >= pes_used: + fill = _INACTIVE_PE + elif _cp_packed: + _pe_cp_rank = pe_id // cfg.topo.tp + _pe_color, _pe_pe_fill = _cp_color( + info[0] if is_used else 0, + _pe_cp_rank, + cfg.topo.cp, + ) + fill = _pe_pe_fill + else: + fill = pe_fill pe_rect = patches.Rectangle( (px, py), pe_w, pe_h, facecolor=fill, edgecolor="#666", linewidth=0.3,