frontier_code/starters/perf_lab.py

"""STARTER for the Frontier F19 lab. Fill in the three TODOs, then run the tests with MANUAL_IMPL=starters.

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:
    """TODO (F19.L1): length in cycles of the longest dependency chain.

    Walk ops in topological order (use _topological). An op with no deps has
    depth 1; otherwise depth = 1 + the largest depth among its deps. Return the
    largest depth (0 for an empty list).
    """
    raise NotImplementedError("F19.L1: write critical_path")


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]]]:
    """TODO (F19.L2): greedy list scheduling into bundles.

    Repeat until every op is issued:
    1. ready = ops not yet issued whose deps were ALL issued in earlier cycles.
    2. Sort ready by (-heights[name], name).
    3. Walk that list, placing each op in this cycle's bundle if its engine
       still has a free slot.
    4. Append the bundle ({engine: [names]}) and advance the cycle.
    Raise RuntimeError if a cycle issues nothing. verify_schedule checks you.
    """
    raise NotImplementedError("F19.L2: write list_schedule")


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]:
    """TODO (F19.L3): search every key together, one level per step.

    Keep one index per key, all starting at 1. For each of the tree's levels,
    replace every index i with 2*i + (1 if key > tree[i] else 0). Finish by
    subtracting 2**depth from each index. Results must equal search_one.
    """
    raise NotImplementedError("F19.L3: write search_batch_lockstep")


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