diff --git a/src/kernbench/ccl/sfr_config.py b/src/kernbench/ccl/sfr_config.py index 7c8516e..e29ac46 100644 --- a/src/kernbench/ccl/sfr_config.py +++ b/src/kernbench/ccl/sfr_config.py @@ -161,3 +161,79 @@ def configure_sfr_intercube_multisip( algo_module=mock_module, rank_to_pe=pe_idx_to_pe, ) + + +# ── Intra-cube PE ring (ADR-0058 Proposed) ───────────────────────────── + + +def configure_sfr_intracube_pe_ring( + engine: Any, + spec: dict, + cfg: dict, +) -> dict[str, Any]: + """Install an 8-PE logical ring inside every cube on every SIP. + + Per cube ``c`` on every SIP, every PE ``i`` ∈ [0, pes_per_cube) gets: + + nbrs[i]["E"] = pe ((i + 1) % pes_per_cube) on the same cube and sip + nbrs[i]["W"] = pe ((i - 1) % pes_per_cube) on the same cube and sip + + No intercube or inter-SIP edges are installed by this function — the + single_user_* attention panels (ADR-0057) operate inside one cube, + so cross-cube traffic is architecturally not part of their kernel. + + The ``E``/``W`` direction namespace is shared with + ``configure_sfr_intercube_multisip`` (which writes cube-mesh + same-lane edges to the same names). The two installs are + **mutually exclusive on the same engine** — the bench is responsible + for picking one per panel-run (ADR-0058 D2). v1 does not add a + runtime guard. + + Args: + engine: GraphEngine with ``_components``. + spec: topology spec dict (from topology.yaml). + cfg: merged algorithm config (from ``resolve_algorithm_config``). + + Returns: + The install plan dict from ``install_ipcq`` (same shape as + ``configure_sfr_intercube_multisip``). + """ + cm = spec["sip"]["cube_mesh"] + n_cubes = int(cm["w"]) * int(cm["h"]) + n_sips = int(spec.get("system", {}).get("sips", {}).get("count", 1)) + pl = spec["cube"]["pe_layout"] + pes_per_cube = int(pl["pe_per_corner"]) * len(pl["corners"]) + + world_size = n_sips * n_cubes * pes_per_cube + pe_idx_to_pe: list[tuple[int, int, int]] = [ + (sip, cube, pe) + for sip in range(n_sips) + for cube in range(n_cubes) + for pe in range(pes_per_cube) + ] + + def _pe_idx(sip: int, cube: int, pe: int) -> int: + return (sip * n_cubes + cube) * pes_per_cube + pe + + def _neighbors(pe_idx: int, ws: int, _base: dict) -> dict[str, int]: + tmp = pe_idx + pe = tmp % pes_per_cube + tmp //= pes_per_cube + cube = tmp % n_cubes + sip = tmp // n_cubes + return { + "E": _pe_idx(sip, cube, (pe + 1) % pes_per_cube), + "W": _pe_idx(sip, cube, (pe - 1) % pes_per_cube), + } + + mock_module = types.SimpleNamespace(neighbors=_neighbors) + + cfg_copy = dict(cfg) + cfg_copy["world_size"] = world_size + cfg_copy["topology"] = "none" + + return install_ipcq( + engine, spec, cfg_copy, + algo_module=mock_module, + rank_to_pe=pe_idx_to_pe, + )