src/paper_stats.py

"""Structural statistics for the host, its levelization, and the interpreter.

Reproduces the level assignment of matrix8.compile_net without materialising the
dense matrices, so the layer widths, the number of identity units, the total
size, and the maximal effective fan-in can be reported for maps that are too
large to compile in memory.
"""

from __future__ import annotations

import json
import os
import sys
from collections import Counter, defaultdict
from typing import Dict, List, Tuple

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def runs_path(name: str) -> str:
    d = os.path.join(REPO, "paper", "runs")
    os.makedirs(d, exist_ok=True)
    return os.path.join(d, name)


def levelize(net, inputs: List[str], outputs: List[str]) -> Dict:
    """The level assignment of matrix8.compile_net, counted and not built."""
    allgates = net.gates
    level: Dict[str, int] = {s: 0 for s in inputs}
    consumers: Dict[str, List[str]] = defaultdict(list)
    indeg: Dict[str, int] = {}
    for gname, (ins, _) in allgates.items():
        indeg[gname] = len([s for s, _ in ins if s in allgates])
        for s, _ in ins:
            consumers[s].append(gname)
    order = [g for g, d in indeg.items() if d == 0]
    i = 0
    while i < len(order):
        for c in consumers.get(order[i], []):
            indeg[c] -= 1
            if indeg[c] == 0:
                order.append(c)
        i += 1
    if len(order) != len(allgates):
        raise ValueError("cycle")
    for gname in order:
        ins, _ = allgates[gname]
        level[gname] = 1 + max((level[s] for s, _ in ins), default=0)

    live: set = set()
    stack = [s for s in outputs if s in allgates]
    while stack:
        g = stack.pop()
        if g in live:
            continue
        live.add(g)
        for s, _ in allgates[g][0]:
            if s in allgates and s not in live:
                stack.append(s)
    gates = {g: allgates[g] for g in live}

    kmax = max(level[g] for g in gates)
    K = kmax

    last_use: Dict[str, int] = {}
    for gname, (ins, _) in gates.items():
        for s, _ in ins:
            last_use[s] = max(last_use.get(s, 0), level[gname])
    for s in outputs:
        if s not in ("#0", "#1"):
            last_use[s] = K

    by_level: Dict[int, List[str]] = defaultdict(list)
    for gname in gates:
        by_level[level[gname]].append(gname)

    vec = list(inputs)
    widths = [len(vec)]
    n_new, n_pass = [], []
    for lv in range(1, K + 1):
        if lv < K:
            new_gates = sorted(by_level.get(lv, []))
            passes = [s for s in vec if last_use.get(s, 0) > lv]
            rows = new_gates + passes
            n_new.append(len(new_gates))
            n_pass.append(len(passes))
        else:
            rows = list(outputs)
            new_here = sum(1 for s in rows if s in gates and level[s] == lv)
            n_new.append(new_here)
            n_pass.append(len(rows) - new_here)
        widths.append(len(rows))
        vec = rows

    fanins = [len(ins) for ins, _ in gates.values()]
    biases = [b for _, b in gates.values()]
    nonzero = sum(fanins)
    return {
        "depth": K,
        "gates": len(gates),
        "netlist_gates": len(allgates),
        "state_bits": len(inputs),
        "widths": widths,
        "M": sum(widths[1:]),
        "dense_entries": sum(widths[i] * widths[i + 1] for i in range(len(widths) - 1)),
        "identity_units": sum(n_pass),
        "new_units": sum(n_new),
        "max_fanin_netlist": max(fanins),
        "nonzero_weights_netlist": nonzero,
        "bias_min": min(biases),
        "bias_max": max(biases),
        "fanin_hist": dict(sorted(Counter(fanins).items())),
    }


DATAPATH = ["p + 3", "subtractor", "zero test", "branch", "counter and halt override",
            "halt latch"]
MEMORY = ["decode p", "operand fetch", "decode A and B", "read mux", "write cells"]

_COMPONENT = {
    "P": "decode p", "FA": "operand fetch", "F": "operand fetch",
    "PA": "decode A and B", "PB": "decode A and B",
    "RX": "read mux", "RY": "read mux", "X": "read mux", "Y": "read mux",
    "W": "write cells", "KP": "write cells", "WR": "write cells", "NM": "write cells",
    "Q": "p + 3",
    "PR": "subtractor", "T": "subtractor", "C": "subtractor", "O": "subtractor",
    "N": "subtractor", "R": "subtractor", "H": "subtractor",
    "AX": "zero test", "NX": "zero test", "XN": "zero test", "EQ": "zero test",
    "LEQ": "branch",
    "HP": "counter and halt override", "TA": "counter and halt override",
    "TB": "counter and halt override", "NP": "counter and halt override",
    "AC": "halt latch", "HA": "halt latch", "HB": "halt latch", "NH": "halt latch",
}


def component(name: str) -> str:
    """The block of the host that a unit of host_family.build belongs to."""
    import re
    return _COMPONENT[re.match(r"[A-Z]+", name).group(0)]


def host_component_table() -> Dict:
    """Per-component counts of the step netlist of the host."""
    from selfrep import host_netlist
    net, inputs, outputs = host_netlist()
    d: Dict[str, int] = {}
    rows: Dict[str, Dict] = {}
    for name, (ins, _) in net.gates.items():
        d[name] = 1 + max((d.get(x, 0) for x, _ in ins), default=0)
        c = component(name)
        r = rows.setdefault(c, {"units": 0, "weights": 0, "max_fanin": 0,
                                "level_lo": 10 ** 9, "level_hi": 0})
        r["units"] += 1
        r["weights"] += len(ins)
        r["max_fanin"] = max(r["max_fanin"], len(ins))
        r["level_lo"] = min(r["level_lo"], d[name])
        r["level_hi"] = max(r["level_hi"], d[name])
    return {
        "rows": rows,
        "built_units": len(net.gates),
        "units": len(net.gates),
        "weights": sum(r["weights"] for r in rows.values()),
        "datapath_units": sum(rows[k]["units"] for k in DATAPATH if k in rows),
        "datapath_weights": sum(rows[k]["weights"] for k in DATAPATH if k in rows),
        "memory_units": sum(rows[k]["units"] for k in MEMORY if k in rows),
        "memory_weights": sum(rows[k]["weights"] for k in MEMORY if k in rows),
    }


def main() -> int:
    out = {}
    print("=" * 74)
    print(" The step netlist of the host, by component")
    print("=" * 74)
    h = host_component_table()
    for name in DATAPATH + MEMORY:
        if name not in h["rows"]:
            continue
        d = h["rows"][name]
        lv = (f"{d['level_lo']}" if d["level_lo"] == d["level_hi"]
              else f"{d['level_lo']}-{d['level_hi']}")
        print(f"  {name:<18} units={d['units']:>6}  weights={d['weights']:>6}  "
              f"levels={lv:<7} max eff fan-in={d['max_fanin']}")
    print(f"  {'datapath':<18} units={h['datapath_units']:>6}  "
          f"weights={h['datapath_weights']:>6}")
    print(f"  {'memory':<18} units={h['memory_units']:>6}  "
          f"weights={h['memory_weights']:>6}")
    print(f"  {'step netlist':<18} units={h['units']:>6}  "
          f"weights={h['weights']:>6}")
    out["host_components"] = h

    print()
    print("=" * 74)
    print(" Lev(N_host): levelization of the SUBLEQ step circuit")
    print("=" * 74)
    from netlist_io import net_of_sigma
    from selfrep import read_host
    net, inputs, outputs, _ = net_of_sigma(read_host())
    lev = levelize(net, inputs, outputs)
    for k in ("depth", "gates", "netlist_gates", "state_bits", "M",
              "dense_entries", "identity_units",
              "new_units", "max_fanin_netlist", "nonzero_weights_netlist",
              "bias_min", "bias_max"):
        print(f"  {k:<26} {lev[k]:,}" if isinstance(lev[k], int) else f"  {k:<26} {lev[k]}")
    print(f"  fan-in histogram          {lev['fanin_hist']}")
    print(f"  widths                    {lev['widths']}")
    out["lev_host"] = lev

    print()
    print("=" * 74)
    print(" Interpreter U (reflect.Cfg default)")
    print("=" * 74)
    from reflect import Cfg, build_net
    cfg = Cfg()
    unet, uin, uout = build_net(cfg)
    ulev = levelize(unet, uin, uout)
    for k in ("depth", "gates", "netlist_gates", "state_bits", "M",
              "dense_entries", "identity_units",
              "new_units", "max_fanin_netlist", "nonzero_weights_netlist",
              "bias_min", "bias_max"):
        print(f"  {k:<26} {ulev[k]:,}" if isinstance(ulev[k], int) else f"  {k:<26} {ulev[k]}")
    print(f"  fan-in histogram          {ulev['fanin_hist']}")
    print(f"  cfg: S={cfg.S} G={cfg.G} F={cfg.F} BB={cfg.BB} banks={cfg.banks} "
          f"record_bits={cfg.R} acc_bits={cfg.ACC} slot={cfg.SLOT}")
    print(f"  regions: PTR@{cfg.PTR_BASE} width {cfg.A}; control {cfg.OUT_DATA}..{cfg.WORK_BASE-1}; "
          f"data {cfg.WORK_BASE}..{cfg.NET0-1} ({cfg.NET0-cfg.WORK_BASE} bits); "
          f"Q0 {cfg.NET0}..{cfg.NET1-1}; Q1 {cfg.NET1}..{cfg.S-1}; "
          f"|Q|={cfg.G*cfg.R} b={cfg.R}")
    out["lev_U"] = ulev
    out["U_cfg"] = {"S": cfg.S, "G": cfg.G, "F": cfg.F, "BB": cfg.BB,
                    "banks": cfg.banks, "R": cfg.R, "ACC": cfg.ACC,
                    "SLOT": cfg.SLOT, "A": cfg.A, "NET0": cfg.NET0,
                    "NET1": cfg.NET1, "WORK_BASE": cfg.WORK_BASE,
                    "span": cfg.span, "STATE_BITS": cfg.STATE_BITS}

    path = runs_path("paper_stats.json")
    json.dump(out, open(path, "w"), indent=1)
    print(f"\nwrote {path}")
    return 0


if __name__ == "__main__":
    sys.exit(main())