infer .inputs from routing/routing.json (restored): unblocks multiplier8x8, multiplier2x2, float16/32 add/normalize/pack/classify, and ~120 other circuits with explicit inter-gate wiring

build.py

@@ -105,7 +105,7 @@ import argparse
 import json
 import re
 from pathlib import Path
-from typing import Dict, Iterable, List, Set
+from typing import Any, Dict, Iterable, List, Optional, Set
 
 import torch
 from safetensors import safe_open
@@ -2861,9 +2861,123 @@ def _infer_threshold_computer_bit_cascade(gate: str, reg: SignalRegistry) -> Lis
     return None
 
 
+_ROUTING_TABLE_CACHE: Dict[str, Any] | None = None
+
+
+def _load_routing_table() -> Dict[str, Any]:
+    """Load `routing/routing.json` once; on missing file, return {}.
+
+    The JSON maps each top-level circuit name (e.g. ``arithmetic.multiplier8x8``)
+    to a dict with at least ``inputs`` (top-level external ports) and
+    ``internal`` (gate-suffix -> list-of-source-signal-name). The source
+    signal names are either external port references (``$a[3]``), constants
+    (``#0`` / ``#1``), or relative gate names within the same circuit
+    (e.g. ``pp.r0.c3``). Anchored cross-circuit references are not
+    expressed here; circuits in the table are assumed self-contained.
+    """
+    global _ROUTING_TABLE_CACHE
+    if _ROUTING_TABLE_CACHE is not None:
+        return _ROUTING_TABLE_CACHE
+    path = Path(__file__).parent / "routing" / "routing.json"
+    if not path.exists():
+        _ROUTING_TABLE_CACHE = {}
+        return _ROUTING_TABLE_CACHE
+    with open(path, encoding="utf-8") as f:
+        data = json.load(f)
+    _ROUTING_TABLE_CACHE = data.get("circuits", {}) or {}
+    return _ROUTING_TABLE_CACHE
+
+
+def _infer_from_routing_table(
+    gate: str, reg: SignalRegistry,
+    known_gate_suffixes: Optional[Set[str]] = None,
+) -> Optional[List[int]]:
+    """Resolve a gate's inputs by consulting routing.json.
+
+    ``known_gate_suffixes`` is the set of gate-name suffixes that actually
+    exist in the loaded safetensors (within the same circuit prefix).
+    When the routing JSON names a producer like ``stage0.bit0.ha2.sum``
+    but that suffix isn't a real gate (only ``stage0.bit0.ha2.sum.layer2``
+    is), this set is consulted to redirect the reference to the canonical
+    final-layer signal. Returns None when the gate's circuit isn't
+    covered or its suffix isn't in the routing's ``internal`` map.
+    """
+    table = _load_routing_table()
+    if not table:
+        return None
+    # Find the longest circuit prefix that matches.
+    best: Optional[str] = None
+    for circuit in table:
+        if gate == circuit or gate.startswith(circuit + "."):
+            if best is None or len(circuit) > len(best):
+                best = circuit
+    if best is None:
+        return None
+    info = table[best]
+    internal = info.get("internal") or {}
+    suffix = gate[len(best):].lstrip(".")
+    src_names = internal.get(suffix)
+    if src_names is None:
+        return None
+
+    out: List[int] = []
+    for nm in src_names:
+        if nm in ("#0", "#1"):
+            out.append(reg.register(nm))
+            continue
+        if nm.startswith("$"):
+            out.append(reg.register(f"{best}.{nm}"))
+            continue
+        # Internal cross-gate reference. Some routing entries use a
+        # short form (e.g. ``stage0.bit0.ha2.sum``) that resolves to a
+        # multi-layer gate (``stage0.bit0.ha2.sum.layer2``). When the
+        # short name isn't a real gate, redirect to its ``.layer2``
+        # final-output sibling, which is what threshold-network XOR
+        # cells expose.
+        resolved = nm
+        if known_gate_suffixes is not None and nm not in known_gate_suffixes:
+            for cand in (f"{nm}.layer2", f"{nm}.layer1", f"{nm}.out"):
+                if cand in known_gate_suffixes:
+                    resolved = cand
+                    break
+        out.append(reg.register(f"{best}.{resolved}"))
+    return out
+
+
+def _gate_suffixes_under(prefix: str, tensors: Dict[str, torch.Tensor]) -> Set[str]:
+    """Return the set of gate-name suffixes whose full name starts with `prefix`."""
+    out: Set[str] = set()
+    for k in tensors:
+        for suf in (".weight", ".bias", ".inputs"):
+            if k.endswith(suf):
+                gate = k[: -len(suf)]
+                if gate == prefix or gate.startswith(prefix + "."):
+                    out.add(gate[len(prefix):].lstrip("."))
+                break
+    return out
+
+
 def infer_inputs_for_gate(gate: str, reg: SignalRegistry, tensors: Dict[str, torch.Tensor]) -> List[int]:
     if gate.startswith('manifest.'):
         return []
+    # routing.json holds explicit per-gate input lists for circuits where
+    # naming alone is insufficient (multipliers, float pipelines, etc.).
+    # Consult it first; family-specific inferrers below remain the
+    # fallback for anything not covered there.
+    table = _load_routing_table()
+    matched_circuit: Optional[str] = None
+    if table:
+        for circuit in table:
+            if gate == circuit or gate.startswith(circuit + "."):
+                if matched_circuit is None or len(circuit) > len(matched_circuit):
+                    matched_circuit = circuit
+    suffixes: Optional[Set[str]] = (
+        _gate_suffixes_under(matched_circuit, tensors)
+        if matched_circuit else None
+    )
+    rt = _infer_from_routing_table(gate, reg, known_gate_suffixes=suffixes)
+    if rt is not None:
+        return rt
     # Bit-cascade comparators (integer + float + div) and ternary modular detectors.
     # These are the gates emitted by add_bit_cascade_compare and the ternary
     # modular rebuild in quantize.py; cover them up front so they don't fall