kernbench2/tests/test_intercube_snake_ring.py

"""Tests for snake/serpentine ring extension of configure_sfr_intercube_ring.

Verifies the multi-row snake ring SFR wiring used for ADR-0060 §5.5
prefill Ring KV at C=8 on a 2×4 sub-mesh of the 4×4 CUBE mesh.

The snake path for ``submesh_shape=(2, 4), origin=(0, 0)`` is::

    (0,0)→(0,1)→(0,2)→(0,3)→(1,3)→(1,2)→(1,1)→(1,0)→wrap to (0,0)

In cube indices on a mesh_w=4 grid: ``[0, 1, 2, 3, 7, 6, 5, 4]``.

Every consecutive pair (including the wrap) is a 1-hop CUBE NOC neighbour.
The kernel sees a 1D logical E/W ring; the snake is invisible to it.
"""
from __future__ import annotations

from pathlib import Path

import pytest

from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import configure_sfr_intercube_ring
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology

TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml"

N_CUBES = 16          # 4×4 SIP CUBE mesh
PES_PER_CUBE = 8
MESH_W = 4

# Canonical snake path for (2,4) sub-mesh at origin (0,0).
SNAKE_2X4_O00 = [0, 1, 2, 3, 7, 6, 5, 4]

# Canonical snake path for (2,4) sub-mesh at origin (1,0) (rows 1-2).
SNAKE_2X4_O10 = [4, 5, 6, 7, 11, 10, 9, 8]


def _engine_and_spec():
    topo = resolve_topology(str(TOPOLOGY_PATH))
    engine = GraphEngine(topo.topology_obj, enable_data=True)
    return engine, topo.topology_obj.spec


def _merged_cfg():
    cfg = load_ccl_config()
    return resolve_algorithm_config(cfg, name="lrab_hierarchical_allreduce")


def _qp(engine, sip: int, cube: int, pe: int):
    return engine._components[f"sip{sip}.cube{cube}.pe{pe}.pe_ipcq"].queue_pairs


def _row_col(cube: int) -> tuple[int, int]:
    return cube // MESH_W, cube % MESH_W


def _l1(c1: int, c2: int) -> int:
    r1, col1 = _row_col(c1)
    r2, col2 = _row_col(c2)
    return abs(r1 - r2) + abs(col1 - col2)


class TestSnakeRing2x4Origin00:
    """Snake ring through the top 2×4 sub-mesh of the 4×4 SIP."""

    def test_snake_ring_2x4_world_size(self):
        """All PEs on all CUBEs of all SIPs still enumerated (snake only
        restricts which CUBEs have ring links, not the world)."""
        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        plan = configure_sfr_intercube_ring(
            engine, spec, cfg, submesh_shape=(2, 4),
        )
        n_sips = int(spec["system"]["sips"]["count"])
        assert plan["world_size"] == n_sips * N_CUBES * PES_PER_CUBE
        assert len(plan["rank_to_pe"]) == plan["world_size"]

    def test_snake_ring_2x4_path_corners_E(self):
        """The four 'interesting' E hops on the snake path:

          cube0.E → cube1   (start of row 0)
          cube3.E → cube7   (row-bridge: top-right corner down)
          cube7.E → cube6   (row 1, going leftward)
          cube4.E → cube0   (wrap: bottom-left to top-left)
        """
        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        configure_sfr_intercube_ring(
            engine, spec, cfg, submesh_shape=(2, 4),
        )
        assert _qp(engine, 0, 0, 0)["E"]["peer"].cube == 1
        assert _qp(engine, 0, 3, 0)["E"]["peer"].cube == 7
        assert _qp(engine, 0, 7, 0)["E"]["peer"].cube == 6
        assert _qp(engine, 0, 4, 0)["E"]["peer"].cube == 0

    def test_snake_ring_2x4_path_corners_W(self):
        """Symmetric W hops (W is reverse of E along the snake):

          cube0.W → cube4   (wrap reverse)
          cube7.W → cube3   (row-bridge reverse: up)
          cube6.W → cube7   (row 1 reverse: rightward)
          cube4.W → cube5   (row-1 leftmost, W goes right within row 1)

        Note: under the snake path [0,1,2,3,7,6,5,4], W(cube_i) is the
        predecessor on that path. So W(4)=5 (previous on path).
        """
        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        configure_sfr_intercube_ring(
            engine, spec, cfg, submesh_shape=(2, 4),
        )
        assert _qp(engine, 0, 0, 0)["W"]["peer"].cube == 4
        assert _qp(engine, 0, 7, 0)["W"]["peer"].cube == 3
        assert _qp(engine, 0, 6, 0)["W"]["peer"].cube == 7
        assert _qp(engine, 0, 4, 0)["W"]["peer"].cube == 5

    def test_snake_ring_2x4_interior(self):
        """Interior hops along each row (not at row-bridge or wrap)."""
        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        configure_sfr_intercube_ring(
            engine, spec, cfg, submesh_shape=(2, 4),
        )
        # Row 0 interior: cube1 between cube0 and cube2.
        assert _qp(engine, 0, 1, 0)["E"]["peer"].cube == 2
        assert _qp(engine, 0, 1, 0)["W"]["peer"].cube == 0
        # Row 1 interior: cube5 between cube6 and cube4 along the snake.
        assert _qp(engine, 0, 5, 0)["E"]["peer"].cube == 4
        assert _qp(engine, 0, 5, 0)["W"]["peer"].cube == 6

    def test_snake_ring_2x4_all_hops_are_1_hop(self):
        """Every E hop on the snake is L1-distance 1 in the cube mesh.

        Defensive — holds by construction (snake = Hamiltonian cycle
        on the 2×4 grid graph), but explicit guards against future
        path-builder regressions.
        """
        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        configure_sfr_intercube_ring(
            engine, spec, cfg, submesh_shape=(2, 4),
        )
        for cube in SNAKE_2X4_O00:
            qp = _qp(engine, 0, cube, 0)
            e_peer = qp["E"]["peer"].cube
            w_peer = qp["W"]["peer"].cube
            assert _l1(cube, e_peer) == 1, (
                f"snake E hop cube{cube}→cube{e_peer} is not 1-hop"
            )
            assert _l1(cube, w_peer) == 1, (
                f"snake W hop cube{cube}→cube{w_peer} is not 1-hop"
            )

    def test_snake_ring_2x4_off_path_cubes_have_no_ring_links(self):
        """CUBEs not on the snake path (rows 2-3, i.e. cubes 8..15)
        get no E/W ring entries — consistent with the current
        ``if cube < ring_size`` behaviour."""
        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        configure_sfr_intercube_ring(
            engine, spec, cfg, submesh_shape=(2, 4),
        )
        for cube in range(8, 16):
            qp = _qp(engine, 0, cube, 0)
            assert "E" not in qp, (
                f"sip0.cube{cube}.pe0 has E but is off the snake path"
            )
            assert "W" not in qp

    def test_snake_ring_2x4_intra_namespace_unchanged(self):
        """Snake doesn't disturb the intra-cube 2×4 PE grid wiring
        (intra_N/S/E/W) — it only changes cube-level E/W."""
        from kernbench.ccl.sfr_config import _intra_cube_neighbors

        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        configure_sfr_intercube_ring(
            engine, spec, cfg, submesh_shape=(2, 4),
        )
        # Spot-check a few cubes (including off-path) and all 8 PEs.
        for cube in (0, 3, 4, 7, 10, 15):
            for pe in range(PES_PER_CUBE):
                qp = _qp(engine, 0, cube, pe)
                expected = _intra_cube_neighbors(pe)
                for d, expected_pe in expected.items():
                    assert d in qp, f"cube{cube}.pe{pe} missing {d}"
                    assert qp[d]["peer"].pe == expected_pe
                    assert qp[d]["peer"].cube == cube


class TestSnakeRing2x4OriginShift:
    """Snake ring through rows 1-2 (origin=(1, 0)) — proves the
    origin parameter shifts the sub-mesh."""

    def test_snake_ring_origin_shift_path(self):
        """For origin=(1, 0), snake path is [4,5,6,7,11,10,9,8]:

          row 1 left→right:  cube4 → cube5 → cube6 → cube7
                                                       ↓
          row 2 right→left:  cube11 → cube10 → cube9 → cube8
                                                          ↺ wraps to cube4
        """
        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        configure_sfr_intercube_ring(
            engine, spec, cfg, submesh_shape=(2, 4), submesh_origin=(1, 0),
        )
        # Path corners.
        assert _qp(engine, 0, 4, 0)["E"]["peer"].cube == 5
        assert _qp(engine, 0, 7, 0)["E"]["peer"].cube == 11
        assert _qp(engine, 0, 11, 0)["E"]["peer"].cube == 10
        assert _qp(engine, 0, 8, 0)["E"]["peer"].cube == 4  # wrap

    def test_snake_ring_origin_shift_off_path_no_links(self):
        """For origin=(1, 0), cubes 0..3 (row 0) and 12..15 (row 3)
        get no ring links."""
        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        configure_sfr_intercube_ring(
            engine, spec, cfg, submesh_shape=(2, 4), submesh_origin=(1, 0),
        )
        for cube in list(range(0, 4)) + list(range(12, 16)):
            qp = _qp(engine, 0, cube, 0)
            assert "E" not in qp
            assert "W" not in qp


class TestSnakeRingBackwardCompat:
    """The existing 1D-row API (no submesh_shape) must behave identically
    to before the snake extension lands."""

    def test_snake_ring_backward_compat_1d_default(self):
        """ring_size=4 (single-row) without submesh_shape — identical
        behaviour to today's 1D-row ring: cube0..3 with E/W wrap;
        cubes 4..15 have no ring links."""
        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        configure_sfr_intercube_ring(
            engine, spec, cfg, ring_size=4,
        )
        # Row-0 1D ring with wrap.
        assert _qp(engine, 0, 0, 0)["E"]["peer"].cube == 1
        assert _qp(engine, 0, 0, 0)["W"]["peer"].cube == 3  # wrap
        assert _qp(engine, 0, 3, 0)["E"]["peer"].cube == 0  # wrap
        # Off-row cubes get no ring links.
        for cube in range(4, 16):
            qp = _qp(engine, 0, cube, 0)
            assert "E" not in qp
            assert "W" not in qp


class TestSnakeRingValidation:
    """Input validation of the submesh_shape / ring_size combinations."""

    def test_snake_ring_invalid_submesh_overflow(self):
        """submesh_shape that doesn't fit the cube mesh is rejected."""
        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        with pytest.raises(ValueError, match=r"sub.?mesh"):
            configure_sfr_intercube_ring(
                engine, spec, cfg, submesh_shape=(2, 5),  # col overflow
            )

    def test_snake_ring_ring_size_mismatch(self):
        """submesh_shape and ring_size disagree → ValueError."""
        engine, spec = _engine_and_spec()
        cfg = _merged_cfg()
        with pytest.raises(ValueError, match=r"ring_size"):
            configure_sfr_intercube_ring(
                engine, spec, cfg,
                submesh_shape=(2, 4), ring_size=4,
            )