frontier_code/solutions/perf_lab.py

"""Frontier F19 lab: performance engineering on a toy VLIW/SIMD machine.

Python 3.13 standard library only.

The toy machine issues one *bundle* per cycle. A bundle may hold up to
``slots[engine]`` operations for each engine, and every operation takes one
cycle: its result is usable in the next cycle, never the same one. That is
enough to study the two ideas that decide every cycle count:

* a dependency chain cannot finish faster than its length (critical path), and
* an engine cannot do more than its slots allow (resource bound).

The second half of the lab searches many keys through one implicit binary
tree at the same time (a batched tree traversal) and models how its cost
shifts between the load engine and the vector engine.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field

# Toy machine used throughout F19. It is deliberately small and is not a model
# of any real product or of any company's interview problem.
TOY_SLOTS = {"alu": 4, "valu": 2, "load": 1, "store": 1}
TOY_VLEN = 4


# ---------------------------------------------------------------------------
# Part 1: dependency graphs, lower bounds, and list scheduling
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class Op:
    name: str
    engine: str
    deps: tuple[str, ...] = field(default=())


def _index(ops: list[Op]) -> dict[str, Op]:
    table = {}
    for op in ops:
        if op.name in table:
            raise ValueError(f"duplicate op name {op.name!r}")
        table[op.name] = op
    for op in ops:
        for d in op.deps:
            if d not in table:
                raise ValueError(f"{op.name!r} depends on unknown op {d!r}")
    return table


def _topological(ops: list[Op]) -> list[Op]:
    table = _index(ops)
    order, state = [], {}

    def visit(name: str) -> None:
        mark = state.get(name)
        if mark == 1:
            raise ValueError("dependency cycle")
        if mark == 2:
            return
        state[name] = 1
        for d in table[name].deps:
            visit(d)
        state[name] = 2
        order.append(table[name])

    for op in ops:
        visit(op.name)
    return order


def critical_path(ops: list[Op]) -> int:
    """Length in cycles of the longest dependency chain (every op takes 1 cycle)."""
    depth: dict[str, int] = {}
    for op in _topological(ops):
        depth[op.name] = 1 + max((depth[d] for d in op.deps), default=0)
    return max(depth.values(), default=0)


def resource_bound(ops: list[Op], slots: dict[str, int]) -> int:
    """Fewest cycles any schedule needs if only slot limits mattered."""
    counts: dict[str, int] = {}
    for op in ops:
        if op.engine not in slots:
            raise ValueError(f"no slots for engine {op.engine!r}")
        counts[op.engine] = counts.get(op.engine, 0) + 1
    return max((math.ceil(n / slots[e]) for e, n in counts.items()), default=0)


def lower_bound(ops: list[Op], slots: dict[str, int]) -> int:
    return max(critical_path(ops), resource_bound(ops, slots))


def heights(ops: list[Op]) -> dict[str, int]:
    """Longest chain from each op to the end of the graph, counting the op itself."""
    users: dict[str, list[str]] = {op.name: [] for op in ops}
    for op in ops:
        for d in op.deps:
            users[d].append(op.name)
    h: dict[str, int] = {}
    for op in reversed(_topological(ops)):
        h[op.name] = 1 + max((h[u] for u in users[op.name]), default=0)
    return h


def list_schedule(ops: list[Op], slots: dict[str, int]) -> list[dict[str, list[str]]]:
    """Greedy list scheduling: each cycle, fill every engine's slots with ready ops.

    An op is ready when all its dependencies were issued in earlier cycles.
    Among ready ops, the one with the greatest height (longest remaining chain)
    goes first, ties broken by name, so the result is deterministic.
    """
    table = _index(ops)
    h = heights(ops)
    issued_at: dict[str, int] = {}
    remaining = set(table)
    bundles: list[dict[str, list[str]]] = []
    cycle = 0
    while remaining:
        ready = [n for n in remaining if all(issued_at.get(d, cycle) < cycle for d in table[n].deps)]
        ready.sort(key=lambda n: (-h[n], n))
        used: dict[str, int] = {e: 0 for e in slots}
        bundle: dict[str, list[str]] = {}
        for n in ready:
            e = table[n].engine
            if used[e] < slots[e]:
                used[e] += 1
                bundle.setdefault(e, []).append(n)
                issued_at[n] = cycle
        if not bundle:
            raise RuntimeError("no progress: check for an engine with zero slots")
        remaining -= {n for names in bundle.values() for n in names}
        bundles.append(bundle)
        cycle += 1
    return bundles


def verify_schedule(ops: list[Op], slots: dict[str, int], bundles: list[dict[str, list[str]]]) -> None:
    """Raise ValueError unless every op runs once, after its deps, within slot limits."""
    table = _index(ops)
    issued_at: dict[str, int] = {}
    for cycle, bundle in enumerate(bundles):
        for engine, names in bundle.items():
            if len(names) > slots.get(engine, 0):
                raise ValueError(f"cycle {cycle}: {len(names)} ops on {engine!r}, limit {slots.get(engine, 0)}")
            for n in names:
                if n in issued_at:
                    raise ValueError(f"{n!r} issued twice")
                if table[n].engine != engine:
                    raise ValueError(f"{n!r} placed on {engine!r}, belongs on {table[n].engine!r}")
                issued_at[n] = cycle
    for n, op in table.items():
        if n not in issued_at:
            raise ValueError(f"{n!r} never issued")
        for d in op.deps:
            if issued_at[d] >= issued_at[n]:
                raise ValueError(f"{n!r} in cycle {issued_at[n]} uses {d!r} from cycle {issued_at[d]}")


def utilization(bundles: list[dict[str, list[str]]], slots: dict[str, int]) -> dict[str, float]:
    """Fraction of each engine's slots that the schedule filled."""
    cycles = len(bundles)
    if cycles == 0:
        return {e: 0.0 for e in slots}
    return {e: sum(len(b.get(e, [])) for b in bundles) / (slots[e] * cycles) for e in slots}


def independent_chains(chains: int, length: int, engine: str = "valu") -> list[Op]:
    """``chains`` independent chains of ``length`` dependent ops, as in an unrolled loop."""
    ops = []
    for c in range(chains):
        for k in range(length):
            deps = (f"c{c}_{k - 1}",) if k else ()
            ops.append(Op(f"c{c}_{k}", engine, deps))
    return ops


# ---------------------------------------------------------------------------
# Part 2: batched traversal of an implicit binary search tree
# ---------------------------------------------------------------------------

def eytzinger(sorted_keys: list[int]) -> list[int | None]:
    """Lay sorted keys out as a complete binary search tree in an array.

    Index 1 is the root; node i has children 2i and 2i + 1; index 0 is unused.
    len(sorted_keys) must be 2**depth - 1 so every level is full.
    """
    n = len(sorted_keys)
    depth = (n + 1).bit_length() - 1
    if n == 0 or 2 ** depth - 1 != n:
        raise ValueError("need 2**depth - 1 keys")
    if any(a >= b for a, b in zip(sorted_keys, sorted_keys[1:])):
        raise ValueError("keys must be strictly increasing")
    tree: list[int | None] = [None] * (n + 1)
    it = iter(sorted_keys)

    def fill(i: int) -> None:  # in-order walk visits keys in sorted order
        if i <= n:
            fill(2 * i)
            tree[i] = next(it)
            fill(2 * i + 1)

    fill(1)
    return tree


def tree_depth(tree: list[int | None]) -> int:
    return (len(tree)).bit_length() - 1


def search_one(tree: list[int | None], key: int) -> int:
    """Rank of ``key``: how many stored keys are strictly smaller.

    Branch-free descent: go right exactly when key > node, so the path spells
    the answer in binary. After ``depth`` steps, i - 2**depth is the rank.
    """
    depth = tree_depth(tree)
    i = 1
    for _ in range(depth):
        i = 2 * i + (1 if key > tree[i] else 0)
    return i - 2 ** depth


def search_batch_lockstep(tree: list[int | None], keys: list[int]) -> list[int]:
    """Search all keys together, one tree level per step (level-synchronous).

    Every key takes exactly ``depth`` steps, so the whole batch moves down the
    tree in lockstep. That regular shape is what lets SIMD lanes each carry one
    key and a VLIW scheduler overlap independent keys.
    """
    depth = tree_depth(tree)
    idx = [1] * len(keys)
    for _ in range(depth):
        idx = [2 * i + (1 if k > tree[i] else 0) for i, k in zip(idx, keys)]
    return [i - 2 ** depth for i in idx]


def distinct_nodes_per_level(tree: list[int | None], keys: list[int]) -> list[int]:
    """How many different nodes the batch touches at each level."""
    depth = tree_depth(tree)
    idx = [1] * len(keys)
    out = []
    for _ in range(depth):
        out.append(len(set(idx)))
        idx = [2 * i + (1 if k > tree[i] else 0) for i, k in zip(idx, keys)]
    return out


@dataclass(frozen=True)
class TraversalCost:
    loads: int
    valu: int
    load_bound: int
    valu_bound: int

    @property
    def bound(self) -> int:
        return max(self.load_bound, self.valu_bound)


def traversal_cost(batch: int, depth: int, cached_levels: int, *, vlen: int = TOY_VLEN,
                   slots: dict[str, int] | None = None) -> TraversalCost:
    """Operation counts and resource bounds for a vectorized lockstep traversal.

    Model (stated so you can challenge it):
    * Keys ride in vectors of ``vlen`` lanes: ceil(batch / vlen) vectors.
    * Every level costs 2 vector ops per vector: a compare and an index update.
    * An uncached level needs each lane's node value from memory. Addresses
      differ per lane, so that is one scalar load per key per level.
    * The top ``cached_levels`` levels (2**c - 1 nodes) are loaded once, with
      contiguous vector loads, and kept resident. Picking a lane's node value at
      a cached level with 2**l candidates takes 2**l - 1 vector selects per
      vector (level 0 is a single broadcast value and needs none).
    """
    if not 0 <= cached_levels <= depth:
        raise ValueError("cached_levels must be between 0 and depth")
    slots = TOY_SLOTS if slots is None else slots
    vectors = math.ceil(batch / vlen)
    preload = math.ceil((2 ** cached_levels - 1) / vlen)
    loads = batch * (depth - cached_levels) + preload
    selects = vectors * sum(2 ** level - 1 for level in range(cached_levels))
    valu = vectors * 2 * depth + selects
    return TraversalCost(loads, valu, math.ceil(loads / slots["load"]), math.ceil(valu / slots["valu"]))


def best_cached_levels(batch: int, depth: int, **kwargs) -> tuple[int, TraversalCost]:
    """Cached-level count with the smallest resource bound (ties: fewer levels)."""
    options = [(c, traversal_cost(batch, depth, c, **kwargs)) for c in range(depth + 1)]
    return min(options, key=lambda item: (item[1].bound, item[0]))