Emit 2-D weight-matrix occupancy in the debug sidecar

The floor-plan viz currently colors sublayers by a 1-D signal (which
residual columns an op writes). The real object is 2-D: each op occupies
a region of the layer's actual weight matrices. Record that during weight
writing and emit it in <stem>.debug.json for the viz to join by canonical
node id.

- PlacementRecorder (forward/weight_writer.py): accumulates each op's
  2-D region per weight matrix as weights are written. Holds only ints
  (O(ops), never pins trimmed tensors). One hook in _scatter_attn_head
  covers all attention ops; three cover the MLP write sites. Dense
  weight blocks vs. paired identity diagonals are recorded from the
  actual write pattern, so identity passthroughs don't overstate area.

- compile.py / transformer.py: recorder created once, set_layer per
  sublayer (incl. the overlay delta layer), attached as net.placements
  beside residual_assignment. Always on, so a schedule-cache hit still
  yields it (weight writing always runs).

- export._write_debug_sidecar: emits a `matrices` table (six per layer,
  heads flattened into one `head` axis = d), `placements` keyed by
  canonical id (dense run-rectangles; diagonals coalesced with a `diag`
  flag), and globals n_heads + per-layer d_hidden. Describes the full
  logical (untrimmed) matrices, consistent with the sidecar's existing
  residual view. Additive keys; no format bump. embed/unembed deferred.

Tests (test_matrix_occupancy.py): pure-function dense/diagonal encoding,
completeness (every nonzero compiled weight is covered by a placement),
within-shape bounds, canonical-id-or-bucket keying, ONNX-sidecar
integration, and cache-hit stability.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: physicsrob

tests/compile/forward/test_matrix_occupancy.py

@@ -0,0 +1,268 @@
+"""2-D weight-matrix occupancy instrumentation (floor-plan viz).
+
+Covers the ``matrices`` table + ``placements`` map emitted in the debug
+sidecar (see ``_build_matrix_occupancy`` in compiler/export.py) and the
+``PlacementRecorder`` that feeds it from weight writing.
+
+The load-bearing invariant is **completeness**: the union of an op's
+recorded placements must cover every nonzero entry the compile actually
+wrote into that weight matrix.  If instrumentation misses a write site,
+this fails.
+"""
+
+import json
+import os
+import tempfile
+
+import pytest
+import torch
+
+from torchwright.compiler.export import (
+    DEBUG_META_FORMAT,
+    _build_matrix_occupancy,
+    _dense_rects,
+    _diag_rects,
+    compile_headless_to_onnx,
+    debug_meta_path_for,
+)
+from torchwright.compiler.forward.compile import forward_compile
+from torchwright.compiler.graph_identity import canonical_ids, unwrap_debug
+from torchwright.ops.arithmetic_ops import add, signed_multiply
+from torchwright.ops.inout_nodes import create_input, create_pos_encoding
+
+D = 256
+D_HEAD = 16
+N_HEADS = D // D_HEAD
+
+
+def _build_graph():
+    """A small graph that needs multiple layers (attention + MLP chains)."""
+    a = create_input("a", 1)
+    b = create_input("b", 1)
+    c = create_input("c", 1)
+    prod = signed_multiply(a, b, max_abs1=10, max_abs2=10)
+    out = add(prod, c)
+    return out, create_pos_encoding()
+
+
+def _compile():
+    out, pos = _build_graph()
+    net = forward_compile(D, D_HEAD, out, pos, verbose=False, device="cpu")
+    canon = canonical_ids(unwrap_debug(out))
+    return net, canon, out, pos
+
+
+# ---------------------------------------------------------------------------
+# Pure rectangle encoding — pins dense (cross product) vs diagonal (paired)
+# ---------------------------------------------------------------------------
+
+
+def _expand(rects):
+    """Expand placement rectangles to the set of (row, col) cells they fill."""
+    cells = set()
+    for r in rects:
+        a0s, a0l = r["a0"]
+        a1s, a1l = r["a1"]
+        if r.get("diag"):
+            assert a0l == a1l, "diagonal rect must have equal-length axes"
+            cells.update((a0s + k, a1s + k) for k in range(a0l))
+        else:
+            for i in range(a0l):
+                for j in range(a1l):
+                    cells.add((a0s + i, a1s + j))
+    return cells
+
+
+def test_dense_rects_are_full_cross_product():
+    rows = [0, 1, 2, 5, 6]  # runs (0,3) and (5,2)
+    cols = [10, 11]  # run (10,2)
+    rects = _dense_rects("M", rows, cols)
+    # One rectangle per row-run x col-run pair; together the full product.
+    assert len(rects) == 2
+    assert all("diag" not in r for r in rects)
+    assert _expand(rects) == {(r, c) for r in rows for c in cols}
+
+
+def test_dense_rects_order_independent():
+    # Unsorted input still yields the same cell set (dense = set semantics).
+    assert _expand(_dense_rects("M", [6, 0, 1], [3, 2])) == {
+        (r, c) for r in (0, 1, 6) for c in (2, 3)
+    }
+
+
+def test_diag_rects_pair_in_lockstep():
+    rows = [0, 1, 2, 5]
+    cols = [3, 4, 5, 9]
+    rects = _diag_rects("M", rows, cols)
+    # (0,3)(1,4)(2,5) coalesce into one len-3 segment; (5,9) is its own.
+    assert len(rects) == 2
+    assert all(r["diag"] for r in rects)
+    assert _expand(rects) == {(0, 3), (1, 4), (2, 5), (5, 9)}
+
+
+def test_diag_rects_break_when_only_one_axis_advances():
+    # rows advance by 1 but cols jump -> no coalescing.
+    rects = _diag_rects("M", [0, 1], [0, 5])
+    assert len(rects) == 2
+    assert _expand(rects) == {(0, 0), (1, 5)}
+
+
+# ---------------------------------------------------------------------------
+# Matrices table — shapes and axis labels
+# ---------------------------------------------------------------------------
+
+
+def test_matrices_table_shapes_and_axes():
+    net, canon, _, _ = _compile()
+    matrices, placements, n_heads, d_hidden = _build_matrix_occupancy(
+        net, canon, D, D_HEAD
+    )
+
+    assert n_heads == N_HEADS
+    assert isinstance(d_hidden, list)
+    assert len(d_hidden) == len(net.layers)
+    assert all(isinstance(h, int) and h > 0 for h in d_hidden)
+
+    for i, dh in enumerate(d_hidden):
+        expect = {
+            f"L{i}.attn.W_Q": ([D, D], "residual_in", "head"),
+            f"L{i}.attn.W_K": ([D, D], "residual_in", "head"),
+            f"L{i}.attn.W_V": ([D, D], "residual_in", "head"),
+            f"L{i}.attn.W_O": ([D, D], "head", "residual_out"),
+            f"L{i}.mlp.W_in": ([D, dh], "residual_in", "hidden"),
+            f"L{i}.mlp.W_out": ([dh, D], "hidden", "residual_out"),
+        }
+        for mid, (shape, a0, a1) in expect.items():
+            rec = matrices[mid]
+            assert rec["layer"] == i
+            assert rec["shape"] == shape, mid
+            assert rec["axis0"] == a0 and rec["axis1"] == a1, mid
+
+
+# ---------------------------------------------------------------------------
+# Completeness — placements cover every nonzero weight the compile wrote
+# ---------------------------------------------------------------------------
+
+
+def _flatten_weight(net, layer_i, kind):
+    """The compiled weight tensor for one matrix kind, in flat 2-D coords."""
+    layer = net.layers[layer_i]
+    attn = layer.attn.attn
+    if kind == "attn.W_Q":
+        return attn.query_matrix.permute(1, 0, 2).reshape(D, N_HEADS * D_HEAD)
+    if kind == "attn.W_K":
+        return attn.key_matrix.permute(1, 0, 2).reshape(D, N_HEADS * D_HEAD)
+    if kind == "attn.W_V":
+        return attn.value_matrix.permute(1, 0, 2).reshape(D, N_HEADS * D_HEAD)
+    if kind == "attn.W_O":
+        return attn.output_matrix.reshape(N_HEADS * D_HEAD, D)
+    if kind == "mlp.W_in":
+        return layer.mlp.linear1.output_matrix
+    if kind == "mlp.W_out":
+        return layer.mlp.linear2.output_matrix
+    raise ValueError(kind)
+
+
+def test_placements_cover_all_nonzero_weights():
+    # trim_heads=False keeps the full (n_heads, d, d_head) tensors so the flat
+    # layout matches the sidecar's full-matrix description.
+    out, pos = _build_graph()
+    net = forward_compile(
+        D, D_HEAD, out, pos, verbose=False, device="cpu", trim_heads=False
+    )
+    canon = canonical_ids(unwrap_debug(out))
+    matrices, placements, _, _ = _build_matrix_occupancy(net, canon, D, D_HEAD)
+
+    # Gather every rectangle per matrix id (across node keys AND reserved
+    # "_<op_type>" buckets — completeness needs the node-less ops too).
+    rects_by_matrix: dict[str, list] = {}
+    for rects in placements.values():
+        for r in rects:
+            rects_by_matrix.setdefault(r["matrix"], []).append(r)
+
+    n_checked = 0
+    for mid, rec in matrices.items():
+        layer_i, kind = rec["layer"], rec["kind"]
+        weight = _flatten_weight(net, layer_i, kind)
+        rows_n, cols_n = rec["shape"]
+        assert tuple(weight.shape) == (rows_n, cols_n), mid
+
+        covered = _expand(rects_by_matrix.get(mid, []))
+        # All claimed cells are inside the matrix.
+        for r, c in covered:
+            assert 0 <= r < rows_n and 0 <= c < cols_n, f"{mid} cell ({r},{c}) OOB"
+
+        nz = torch.nonzero(weight, as_tuple=False)
+        for r, c in nz.tolist():
+            assert (r, c) in covered, (
+                f"{mid}: nonzero weight at ({r},{c}) not covered by any "
+                f"placement — instrumentation missed a write site."
+            )
+            n_checked += 1
+
+    assert n_checked > 0, "expected at least some nonzero weights to verify"
+
+
+# ---------------------------------------------------------------------------
+# Keying — every placement key is a canonical id or a reserved bucket
+# ---------------------------------------------------------------------------
+
+
+def test_placement_keys_are_canonical_or_bucket():
+    net, canon, _, _ = _compile()
+    _, placements, _, _ = _build_matrix_occupancy(net, canon, D, D_HEAD)
+
+    valid_ids = set(canon.values())
+    real_node_keys = 0
+    for key in placements:
+        if key.startswith("_"):
+            continue  # reserved bucket (e.g. "_cancel", "_unreachable")
+        assert int(key) in valid_ids, f"placement key {key!r} is not a canonical id"
+        real_node_keys += 1
+    assert real_node_keys > 0, "expected some placements keyed by real graph nodes"
+
+
+# ---------------------------------------------------------------------------
+# Sidecar integration — the ONNX export carries the occupancy data
+# ---------------------------------------------------------------------------
+
+
+def _read_sidecar(tmpdir, name="model.onnx"):
+    out, pos = _build_graph()
+    onnx_path = os.path.join(tmpdir, name)
+    compile_headless_to_onnx(
+        out, pos, onnx_path, d=D, d_head=D_HEAD, max_seq_len=32, verbose=False
+    )
+    with open(debug_meta_path_for(onnx_path)) as f:
+        return json.load(f)
+
+
+def test_sidecar_carries_occupancy():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        data = _read_sidecar(tmpdir)
+
+    assert data["format"] == DEBUG_META_FORMAT
+    assert data["n_heads"] == N_HEADS
+    assert isinstance(data["d_hidden"], list)
+    assert len(data["d_hidden"]) == data["n_layers"]
+    assert data["matrices"], "matrices table should be non-empty"
+    assert data["placements"], "placements should be non-empty"
+    # Every placement references a declared matrix id.
+    matrix_ids = set(data["matrices"])
+    for rects in data["placements"].values():
+        for r in rects:
+            assert r["matrix"] in matrix_ids
+
+
+def test_occupancy_stable_across_cache_hit(monkeypatch):
+    # The schedule cache replays the CP-SAT solve, but weight writing (and thus
+    # placement recording) always runs — a cache hit must still yield identical
+    # occupancy data.
+    with tempfile.TemporaryDirectory() as cache_dir:
+        monkeypatch.setenv("TW_SCHEDULE_CACHE_DIR", cache_dir)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            first = _read_sidecar(tmpdir, "a.onnx")  # cold: solves + caches
+            second = _read_sidecar(tmpdir, "b.onnx")  # warm: cache hit
+
+    for key in ("n_heads", "d_hidden", "matrices", "placements"):
+        assert first[key] == second[key], f"{key} differs across cache hit"