kernbench2/src/kernbench/ccl/sfr_config.py

"""SFR configuration for the full IPCQ hardware wiring.

Installs PE_IPCQ neighbor tables modeling the physical hardware.
Wiring is independent of DPPolicy / kernel choice — the kernel decides
at runtime which links to use.

Direction label namespaces (disjoint):

  - Intra-cube PE-to-PE:   ``intra_N / intra_S / intra_E / intra_W``
    Logical 2×4 PE grid within a cube (no wrap):

         Row 0:  pe0  pe1  pe2  pe3
         Row 1:  pe4  pe5  pe6  pe7

  - Intercube same-lane:   ``N / S / E / W``
    ``pe_i of cube_A ↔ pe_i of cube_B`` across the 4×4 cube mesh
    (no wrap). Every PE i ∈ [0..7] wired independently.

  - Inter-SIP same-(cube, pe): ``global_N / global_S / global_E / global_W``
    ``pe_i of cube_c on sip_A ↔ pe_i of cube_c on sip_B`` per
    ``topology.yaml → system.sips.topology``.
"""
from __future__ import annotations

import types
from typing import Any

from kernbench.ccl.install import install_ipcq
from kernbench.ccl.topologies import _BUILTIN as _TOPO_BUILTINS


# ── Intra-cube 2×4 PE grid ───────────────────────────────────────────

_PE_GRID_COLS = 4
_PE_GRID_ROWS = 2
_PES_PER_CUBE = _PE_GRID_COLS * _PE_GRID_ROWS  # 8


def _intra_cube_neighbors(pe: int) -> dict[str, int]:
    """Logical 2×4 PE grid neighbors within a cube (no wrap).

    Returns directions in the ``intra_*`` namespace.
    """
    row, col = divmod(pe, _PE_GRID_COLS)
    nbrs: dict[str, int] = {}
    if col < _PE_GRID_COLS - 1:
        nbrs["intra_E"] = row * _PE_GRID_COLS + (col + 1)
    if col > 0:
        nbrs["intra_W"] = row * _PE_GRID_COLS + (col - 1)
    if row < _PE_GRID_ROWS - 1:
        nbrs["intra_S"] = (row + 1) * _PE_GRID_COLS + col
    if row > 0:
        nbrs["intra_N"] = (row - 1) * _PE_GRID_COLS + col
    return nbrs


# ── Public entry point ───────────────────────────────────────────────


def configure_sfr_intercube_multisip(
    engine: Any,
    spec: dict,
    cfg: dict,
) -> dict[str, Any]:
    """Wire the full IPCQ hardware model.

    Every PE on every cube on every SIP gets neighbor table entries for:

      - intra-cube (2×4 grid) in the ``intra_*`` namespace
      - intercube same-lane (4×4 cube mesh, no wrap) in ``N/S/E/W``
      - inter-SIP same-(cube, pe) in ``global_*``

    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``.
    """
    cm = spec["sip"]["cube_mesh"]
    mesh_w = int(cm["w"])
    mesh_h = int(cm["h"])
    n_cubes = mesh_w * mesh_h
    sips_cfg = spec.get("system", {}).get("sips", {})
    n_sips = int(sips_cfg.get("count", 1))
    sip_topology = str(sips_cfg.get("topology", "ring_1d"))
    sip_w = sips_cfg.get("w")
    sip_h = sips_cfg.get("h")
    sip_w = int(sip_w) if sip_w is not None else None
    sip_h = int(sip_h) if sip_h is not None else None

    if sip_topology not in _TOPO_BUILTINS:
        raise ValueError(
            f"Unknown sip topology '{sip_topology}'. "
            f"Available: {list(_TOPO_BUILTINS)}"
        )
    _sip_topo_fn_raw = _TOPO_BUILTINS[sip_topology]

    def sip_topo_fn(rank: int, ws: int) -> dict:
        if sip_w is not None and sip_h is not None:
            try:
                return _sip_topo_fn_raw(rank, ws, w=sip_w, h=sip_h)
            except TypeError:
                pass
        return _sip_topo_fn_raw(rank, ws)

    pes_per_cube = _PES_PER_CUBE
    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
        row = cube // mesh_w
        col = cube % mesh_w

        nbrs: dict[str, int] = {}

        # ── Intra-cube (intra_N/S/E/W) ──
        for d, peer_pe in _intra_cube_neighbors(pe).items():
            nbrs[d] = _pe_idx(sip, cube, peer_pe)

        # ── Intercube same-lane (N/S/E/W, 4×4 no wrap) ──
        if col < mesh_w - 1:
            nbrs["E"] = _pe_idx(sip, row * mesh_w + (col + 1), pe)
        if col > 0:
            nbrs["W"] = _pe_idx(sip, row * mesh_w + (col - 1), pe)
        if row < mesh_h - 1:
            nbrs["S"] = _pe_idx(sip, (row + 1) * mesh_w + col, pe)
        if row > 0:
            nbrs["N"] = _pe_idx(sip, (row - 1) * mesh_w + col, pe)

        # ── Inter-SIP same-(cube, pe) (global_*) ──
        if n_sips > 1:
            sip_nbrs = sip_topo_fn(sip, n_sips)
            for d, peer_sip in sip_nbrs.items():
                nbrs[f"global_{d}"] = _pe_idx(peer_sip, cube, pe)

        return nbrs

    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,
    )


# ── 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,
    )


# ── Inter-cube 1D ring (ADR-0060 §5.5 prefill Ring KV) ─────────────────


def configure_sfr_intercube_ring(
    engine: Any,
    spec: dict,
    cfg: dict,
    *,
    ring_size: int | None = None,
    submesh_shape: tuple[int, int] | None = None,
    submesh_origin: tuple[int, int] = (0, 0),
) -> dict[str, Any]:
    """Install intra-cube PE grid + a CUBE-level ring with wrap.

    Two ring layouts:

      - **1D row** (default; ``submesh_shape=None``): cubes
        ``0..ring_size-1`` form a 1D ring with wrap. ``ring_size`` must
        be ≤ ``mesh_w`` so every hop is a 1-hop CUBE NOC neighbour.
      - **Snake/serpentine** (``submesh_shape=(rows, cols)``,
        ADR-0060 §5.5 prefill Ring KV at C=G=8 on a 2×4 sub-mesh):
        a boustrophedon Hamiltonian cycle through the ``rows × cols``
        sub-mesh rooted at ``submesh_origin``. Every consecutive pair
        on the snake (including the wrap) is a 1-hop CUBE NOC
        neighbour, so the kernel sees a 1D logical E/W ring without
        being aware of the underlying 2D layout.

    Direction namespaces (disjoint, same as
    ``configure_sfr_intercube_multisip``):

      - ``intra_N/S/E/W`` : 2×4 PE grid within each cube (no wrap)
      - ``E/W``           : ring along the resolved path WITH WRAP
                            (symmetric to ``configure_sfr_intracube_pe_ring``
                             at PE level — wrap applied at CUBE level here)
      - ``global_*``      : SIP topology (same as multisip)

    N/S at CUBE level are intentionally NOT installed — use
    ``configure_sfr_intercube_multisip`` for the full 4×4 cube mesh.

    Args:
        ring_size: number of CUBEs in the ring. Defaults to the full
            cube_mesh count for the 1D-row case, or ``rows*cols`` for
            the snake case. If passed alongside ``submesh_shape``, must
            equal ``rows*cols``.
        submesh_shape: ``(rows, cols)`` of the snake sub-mesh. If
            given, ring follows a boustrophedon path through that
            rectangle. If ``None``, falls back to the 1D-row layout.
        submesh_origin: ``(row, col)`` top-left of the sub-mesh inside
            the cube mesh. Defaults to ``(0, 0)``.
    """
    cm = spec["sip"]["cube_mesh"]
    mesh_w = int(cm["w"])
    mesh_h = int(cm["h"])
    n_cubes = mesh_w * mesh_h
    sips_cfg = spec.get("system", {}).get("sips", {})
    n_sips = int(sips_cfg.get("count", 1))
    sip_topology = str(sips_cfg.get("topology", "ring_1d"))
    sip_w = sips_cfg.get("w")
    sip_h = sips_cfg.get("h")
    sip_w = int(sip_w) if sip_w is not None else None
    sip_h = int(sip_h) if sip_h is not None else None

    if submesh_shape is not None:
        sub_h, sub_w = submesh_shape
        origin_row, origin_col = submesh_origin
        if (sub_h <= 0 or sub_w <= 0
                or origin_row < 0 or origin_col < 0
                or origin_row + sub_h > mesh_h
                or origin_col + sub_w > mesh_w):
            raise ValueError(
                f"submesh_shape={submesh_shape} at origin={submesh_origin} "
                f"does not fit cube_mesh ({mesh_h}x{mesh_w})"
            )
        expected_size = sub_h * sub_w
        if ring_size is not None and ring_size != expected_size:
            raise ValueError(
                f"ring_size={ring_size} inconsistent with "
                f"submesh_shape={submesh_shape} (expected {expected_size})"
            )
        ring_size = expected_size
        # Boustrophedon: even rows L→R, odd rows R→L.
        ring_path: list[int] = []
        for r in range(sub_h):
            cols = range(sub_w) if r % 2 == 0 else range(sub_w - 1, -1, -1)
            for c in cols:
                ring_path.append((origin_row + r) * mesh_w + (origin_col + c))
    else:
        if ring_size is None:
            ring_size = n_cubes
        if ring_size > mesh_w:
            raise ValueError(
                f"intercube_ring ring_size={ring_size} > mesh_w={mesh_w}; "
                "multi-row rings cross non-neighbour boundaries"
            )
        ring_path = list(range(ring_size))

    ring_pos: dict[int, int] = {c: i for i, c in enumerate(ring_path)}
    ring_len = len(ring_path)

    if sip_topology not in _TOPO_BUILTINS:
        raise ValueError(
            f"Unknown sip topology '{sip_topology}'. "
            f"Available: {list(_TOPO_BUILTINS)}"
        )
    _sip_topo_fn_raw = _TOPO_BUILTINS[sip_topology]

    def sip_topo_fn(rank: int, ws: int) -> dict:
        if sip_w is not None and sip_h is not None:
            try:
                return _sip_topo_fn_raw(rank, ws, w=sip_w, h=sip_h)
            except TypeError:
                pass
        return _sip_topo_fn_raw(rank, ws)

    pes_per_cube = _PES_PER_CUBE
    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

        nbrs: dict[str, int] = {}

        # ── Intra-cube (intra_N/S/E/W) ──
        for d, peer_pe in _intra_cube_neighbors(pe).items():
            nbrs[d] = _pe_idx(sip, cube, peer_pe)

        # ── Cube ring (E/W along resolved ring_path, with wrap) ──
        pos = ring_pos.get(cube)
        if pos is not None:
            nbrs["E"] = _pe_idx(sip, ring_path[(pos + 1) % ring_len], pe)
            nbrs["W"] = _pe_idx(sip, ring_path[(pos - 1) % ring_len], pe)

        # ── Inter-SIP same-(cube, pe) (global_*) ──
        if n_sips > 1:
            sip_nbrs = sip_topo_fn(sip, n_sips)
            for d, peer_sip in sip_nbrs.items():
                nbrs[f"global_{d}"] = _pe_idx(peer_sip, cube, pe)

        return nbrs

    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,
    )